"""Deterministic newc and GPT writers for host image construction (ordinary files only)."""
import hashlib
from pathlib import Path
import shutil
import struct
import uuid
import zlib

NAMESPACE = uuid.UUID('34c5ce42-2282-4b58-b633-28d410000001')
LINUX_FILESYSTEM = uuid.UUID('0fc63daf-8483-4772-8e79-3d69d8477de4')
EFI_SYSTEM = uuid.UUID('c12a7328-f81f-11d2-ba4b-00a0c93ec93b')

def digest(path):
    with Path(path).open('rb') as stream:
        return hashlib.file_digest(stream, 'sha256').hexdigest()

def newc(entries, epoch):
    """entries: (relative name, mode, bytes, device major, device minor)."""
    result = bytearray()
    for inode, (name, mode, data, major, minor) in enumerate([*entries, ('TRAILER!!!', 0, b'', 0, 0)], 1):
        if name.startswith('/') or '..' in Path(name).parts or '\0' in name:
            raise ValueError(f'Unsafe cpio path: {name}')
        encoded = name.encode() + b'\0'
        fields = [inode, mode, 0, 0, 1, epoch, len(data), 0, 0, major, minor, len(encoded), 0]
        result += b'070701' + ''.join(f'{field:08x}' for field in fields).encode() + encoded
        result += b'\0' * (-len(result) % 4)
        result += data
        result += b'\0' * (-len(result) % 4)
    result += b'\0' * (-len(result) % 512)
    return bytes(result)

def gpt(output, partitions):
    """Create GPT from [(label, type UUID, payload path)] with 1 MiB alignment."""
    if not partitions or len(partitions) > 128:
        raise ValueError('GPT needs 1..128 partitions')
    names = [name for name, _, _ in partitions]
    if len(set(names)) != len(names):
        raise ValueError('Duplicate GPT partition names')
    identities = []
    next_lba = 2048
    for name, kind, payload in partitions:
        path = Path(payload)
        if not path.is_file() or path.is_symlink() or path.stat().st_size == 0:
            raise ValueError('GPT payload must be a nonempty regular file')
        encoded = name.encode('utf-16le')
        if len(encoded) > 72 or '\0' in name:
            raise ValueError('Invalid GPT name')
        sectors = (path.stat().st_size + 511) // 512
        allocated = ((sectors + 2047) // 2048) * 2048
        identities.append(dict(name=name, type=str(kind), path=str(path), sha256=digest(path),
                               start=next_lba, size=allocated, payload_bytes=path.stat().st_size))
        next_lba += allocated
    total = ((next_lba + 33 + 2047) // 2048) * 2048
    disk_id = uuid.uuid5(NAMESPACE, '|'.join(item['name'] + ':' + item['sha256'] for item in identities))
    array = bytearray(128 * 128)
    for index, item in enumerate(identities):
        unique = uuid.uuid5(disk_id, item['name'])
        entry = struct.pack('<16s16sQQQ72s', uuid.UUID(item['type']).bytes_le, unique.bytes_le,
                            item['start'], item['start'] + item['size'] - 1, 0,
                            item['name'].encode('utf-16le'))
        array[index*128:(index+1)*128] = entry
    array_crc = zlib.crc32(array)
    def header(current, backup, entries):
        data = bytearray(struct.pack('<8sIIIIQQQQ16sQIII', b'EFI PART', 0x10000, 92, 0, 0,
                                    current, backup, 34, total-34, disk_id.bytes_le,
                                    entries, 128, 128, array_crc))
        struct.pack_into('<I', data, 16, zlib.crc32(data))
        return data + bytes(512-len(data))
    mbr = bytearray(512)
    mbr[446:462] = struct.pack('<B3sB3sII', 0, b'\0\x02\0', 0xee, b'\xff'*3, 1, min(total-1, 0xffffffff))
    mbr[510:512] = b'\x55\xaa'
    with Path(output).open('xb') as stream:
        stream.truncate(total*512)
        stream.write(mbr)
        stream.write(header(1, total-1, 2))
        stream.write(array)
        for item in identities:
            stream.seek(item['start']*512)
            with Path(item['path']).open('rb') as payload:
                shutil.copyfileobj(payload, stream)
        stream.seek((total-33)*512)
        stream.write(array)
        stream.write(header(total-1, 1, total-33))
    # Verify the written payloads independently of the GPT header calculations.
    with Path(output).open('rb') as stream:
        for item in identities:
            stream.seek(item['start']*512)
            remaining = item['payload_bytes']
            checksum = hashlib.sha256()
            while remaining:
                chunk = stream.read(min(1024*1024, remaining))
                if not chunk: raise ValueError('Truncated GPT payload')
                remaining -= len(chunk)
                checksum.update(chunk)
            if checksum.hexdigest() != item['sha256']:
                raise ValueError('GPT payload verification failed')
    return {'disk_uuid': str(disk_id), 'bytes': total*512, 'partitions': identities}