#!/usr/bin/env python3
"""Stage full shipped ASB manifests for early startup; retain originals.

Observed SRMI selector, SRAM v5 manifest and SRBM v1 block-table formats.
This changes startup availability only; it does not select a login scene.
All installation paths are explicit arguments. Restore refuses unrelated edits.
"""
import argparse
import hashlib
import struct
from pathlib import Path

from install_baked_lua_stage import atomic_copy, selector_digest


KINDS = ('AsbV', 'BlockV')


def selected(root, name):
    selector = root / ('M_' + name + '.bytes')
    raw = selector.read_bytes()
    digest = selector_digest(raw)
    data = root / (name + '_' + digest + '.bytes')
    content = data.read_bytes()
    if hashlib.md5(content).hexdigest() != digest or struct.unpack_from('<I', raw, 44)[0] != len(content):
        raise ValueError('selector size or digest mismatch: ' + name)
    return selector, raw, data, content


def stage(source, output):
    source, output = source.resolve(strict=True), output.resolve()
    if source == output or source in output.parents:
        raise ValueError('stage must be outside installation')
    if output.exists() and any(output.iterdir()):
        raise ValueError('stage must be empty')
    prepared = []
    for kind in KINDS:
        _, old, _, before = selected(source, 'Start_' + kind)
        _, full, path, content = selected(source, kind)
        if old[:28] != full[:28] or before[:16] != content[:16]:
            raise ValueError('startup and full manifest schema/build differ')
        if kind == 'BlockV':
            if content[:8] != b'SRBM\x00\x01\x00\x01':
                raise ValueError('unsupported block table')
            offset, count, stride = struct.unpack_from('<III', content, 16)
            if offset != 28 or stride != 24 or len(content) != offset + count * stride:
                raise ValueError('unsupported block table bounds')
            for i in range(count):
                row = content[offset+i*stride:offset+(i+1)*stride]
                digest = b''.join(row[j:j+4][::-1] for j in range(0, 16, 4)).hex()
                if not (source / (digest + '.block')).is_file():
                    raise ValueError('missing referenced block: ' + digest)
        prepared.append((kind, full, path, content))
    output.mkdir(parents=True, exist_ok=True)
    for kind, selector, path, content in prepared:
        (output / ('Start_' + path.name)).write_bytes(content)
        (output / ('M_Start_' + kind + '.bytes')).write_bytes(selector)
    print('Staged full manifests; all referenced block files exist')


def apply(source, staged, backup, restore=False):
    source, staged, backup = source.resolve(strict=True), staged.resolve(strict=True), backup.resolve()
    if any(p == source or source in p.parents for p in (staged, backup)):
        raise ValueError('stage and backup must be outside installation')
    entries = []
    for kind in KINDS:
        name = 'M_Start_' + kind + '.bytes'
        current = source / name
        _, expected, blob, _ = selected(staged, 'Start_' + kind)
        saved = backup / name
        if restore:
            if current.read_bytes() != expected:
                raise ValueError('current selector is not this overlay: ' + name)
            selected_backup = saved.read_bytes()
            original = source / ('Start_' + kind + '_' + selector_digest(selected_backup) + '.bytes')
            if not original.is_file():
                raise ValueError('original manifest missing')
        else:
            selected(source, 'Start_' + kind)
            if saved.exists() and saved.read_bytes() != current.read_bytes():
                raise ValueError('existing backup differs')
        entries.append((current, saved, staged / name, blob))
    backup.mkdir(parents=True, exist_ok=True)
    for current, saved, selector, blob in entries:
        if restore:
            atomic_copy(saved, current)
        else:
            if not saved.exists():
                atomic_copy(current, saved)
            atomic_copy(blob, source / blob.name)
            atomic_copy(selector, current)
        if current.read_bytes() != (saved if restore else selector).read_bytes():
            raise ValueError('selector write verification failed')
    print('Restored original startup manifests' if restore else 'Installed startup manifest overlay; originals retained')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('mode', choices=('stage', 'install', 'restore'))
    parser.add_argument('source', type=Path)
    parser.add_argument('staged', type=Path)
    parser.add_argument('--backup', type=Path)
    args = parser.parse_args()
    if args.mode == 'stage':
        stage(args.source, args.staged)
    else:
        if args.backup is None:
            parser.error('--backup is required for install/restore')
        apply(args.source, args.staged, args.backup, args.mode == 'restore')