"""Package and verify the completed dependency experiment without prior archives.

The archive includes the current source tree, this entire experiment, the two
frozen starting predictors, and the prior public data/second VM needed here.
It reproduces this round from the included parents, not every older experiment.
No package operation is allowed before the independent final audit succeeds.
"""
from __future__ import annotations

import argparse
from datetime import datetime, timezone
import hashlib
import io
import json
import lzma
import os
from pathlib import Path
import subprocess
import tarfile

ROOT = Path(__file__).resolve().parents[1]
WORK = ROOT / 'work/morphology-20261017'


def sha(path):
    result = hashlib.sha256()
    with Path(path).open('rb') as source:
        for block in iter(lambda: source.read(1024 * 1024), b''):
            result.update(block)
    return result.hexdigest()


def write_json(path, value):
    with Path(path).open('x') as output:
        json.dump(value, output, ensure_ascii=False, sort_keys=True, indent=2)
        output.write('\n')


def collect(bundle):
    git_files = subprocess.check_output(
        ['git', 'ls-files', '-z', '--cached', '--others', '--exclude-standard'], cwd=ROOT)
    paths = {ROOT / item.decode() for item in git_files.split(b'\0') if item}
    paths.update(path for path in WORK.rglob('*') if path.is_file()
                 and '__pycache__' not in path.parts and not path.name.endswith('.pyc'))
    dependencies = [
        'work/architecture-20261017/audit_evaluation.py',
        'work/architecture-20261017/data/train.jsonl',
        'work/architecture-20261017/data/development.jsonl',
        'work/architecture-20261017/data/development-episodes.json',
        'work/architecture-20261017/data-protocol.json',
        'work/architecture-20261017/relation/formal-protocol.json',
    ]
    for seed in (104729, 130363):
        for name in ('final.pt', 'protocol.json', 'report.json', 'training-order.json'):
            dependencies.append(f'work/architecture-20261017/relation/formal-dataflow-seed{seed}/{name}')
    paths.update(ROOT / relative for relative in dependencies)
    result = []
    for path in sorted(paths):
        if path.is_symlink() or not path.is_file() or not path.resolve().is_relative_to(ROOT):
            raise ValueError('unexpected missing or linked input: ' + str(path))
        result.append(dict(path=path, name='lightnest/' + str(path.relative_to(ROOT)),
                           bytes=path.stat().st_size, sha256=sha(path), mode=path.stat().st_mode & 0o777))
    if not bundle.is_file():
        raise ValueError('missing source-history bundle')
    result.append(dict(path=bundle, name='lightnest/source-history.bundle',
                       bytes=bundle.stat().st_size, sha256=sha(bundle), mode=0o644))
    if len({entry['name'] for entry in result}) != len(result):
        raise ValueError('duplicate archive names')
    return result


def verify(archive):
    seen = {}
    manifest = None
    with lzma.open(archive, 'rb') as compressed:
        with tarfile.open(fileobj=compressed, mode='r|') as tar:
            for member in tar:
                if not member.isfile() or member.name in seen or member.name.startswith('/') or '..' in Path(member.name).parts:
                    raise ValueError('unexpected or duplicate archive member')
                stream = tar.extractfile(member)
                digest, size = hashlib.sha256(), 0
                capture = bytearray() if member.name == 'MANIFEST.json' else None
                for block in iter(lambda: stream.read(1024 * 1024), b''):
                    digest.update(block)
                    size += len(block)
                    if capture is not None:
                        capture.extend(block)
                if size != member.size:
                    raise ValueError('truncated archive member')
                seen[member.name] = dict(bytes=size, sha256=digest.hexdigest(), mode=member.mode)
                if capture is not None:
                    manifest = json.loads(capture)
    if manifest is None:
        raise ValueError('archive manifest missing')
    expected = {entry['name']: {key: entry[key] for key in ('bytes', 'sha256', 'mode')}
                for entry in manifest['files']}
    if set(seen) != set(expected) | {'MANIFEST.json'}:
        raise ValueError('archive manifest membership differs')
    for name, entry in expected.items():
        if seen[name] != entry:
            raise ValueError('archive member differs from committed input: ' + name)
    return dict(all_checks_passed=True, archive=str(archive), archive_bytes=archive.stat().st_size,
                archive_sha256=sha(archive), file_count=len(expected),
                payload_bytes=sum(entry['bytes'] for entry in expected.values()),
                complete_member_hash_size_and_mode_check=True,
                source_commit=manifest['source_commit'],
                scope=manifest['scope'])


def package(output, bundle, training_audit):
    if output.exists() or output.with_suffix('.verification.json').exists():
        raise ValueError('refuse to replace a delivered package')
    final_audit = json.loads((WORK / 'independent-audit.json').read_text())
    train_audit = json.loads(training_audit.read_text())
    if (final_audit.get('all_checks_passed') is not True or
            train_audit.get('schema') != 'lightnest-independent-dependency-training-audit-v1' or
            train_audit.get('passed') is not True or
            train_audit.get('complete_run_count') != 12):
        raise ValueError('completed independent training and final audits are required')
    inventory_path = WORK / 'models-frozen.json'
    inventory = json.loads(inventory_path.read_text())
    if (final_audit['inventory_sha256'] != sha(inventory_path) or
            train_audit['protocol'] != inventory['training_protocol']):
        raise ValueError('audit commitments do not match the frozen experiment')
    committed = [inventory['training_protocol'], inventory['data_protocol'],
                 inventory['specification'], *inventory['role_reports'],
                 *inventory['sources'], *inventory['models'],
                 *final_audit['artifacts'], *train_audit['verified_inputs'],
                 *inventory['suites'].values()]
    for record in committed:
        path = (ROOT / record['path']).resolve()
        if (not path.is_relative_to(ROOT) or not path.is_file() or
                sha(path) != record['sha256'] or
                ('bytes' in record and path.stat().st_size != record['bytes'])):
            raise ValueError('audited evidence changed before packaging: ' + record['path'])
    subprocess.run(['git', 'bundle', 'verify', str(bundle)], cwd=ROOT, check=True,
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    entries = collect(bundle)
    manifest = dict(schema='lightnest-dependency-delivery-v1',
        created_utc=datetime.now(timezone.utc).isoformat(),
        source_commit=subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=ROOT, text=True).strip(),
        scope='Current source and complete dependency-round evidence, including starting parent '
              'predictors, required prior public data, and the second VM. Earlier full experimental '
              'histories remain in their separate preserved archives.',
        model_inventory_sha256=sha(WORK / 'models-frozen.json'),
        final_audit_sha256=sha(WORK / 'independent-audit.json'),
        training_audit_sha256=sha(training_audit),
        files=[{key: value for key, value in entry.items() if key != 'path'} for entry in entries])
    payload = (json.dumps(manifest, ensure_ascii=False, sort_keys=True, separators=(',', ':')) + '\n').encode()
    output.parent.mkdir(parents=True, exist_ok=True)
    # Keep the bounded (~120 MB input) delivery in memory until its exact
    # compressed size is known. Avoid partially filling the nearly full scratch
    # disk. This does not change any experiment member or its committed bytes.
    with io.BytesIO() as packed:
        with lzma.LZMAFile(packed, 'wb', preset=9) as compressed:
            with tarfile.open(fileobj=compressed, mode='w|') as tar:
                item = tarfile.TarInfo('MANIFEST.json')
                item.size, item.mode, item.mtime = len(payload), 0o644, 0
                tar.addfile(item, io.BytesIO(payload))
                for entry in entries:
                    info = tar.gettarinfo(str(entry['path']), arcname=entry['name'])
                    info.uid = info.gid = 0
                    info.uname = info.gname = ''
                    info.mode = entry['mode']
                    with entry['path'].open('rb') as source:
                        tar.addfile(info, source)
        required = packed.tell()
        filesystem = os.statvfs(output.parent)
        available = filesystem.f_bavail * filesystem.f_frsize
        if required + 1024 * 1024 > available:
            raise ValueError(f'archive needs {required} bytes plus 1 MiB reserve; available {available}')
        with output.open('xb') as raw:
            with packed.getbuffer() as payload_view:
                raw.write(payload_view)
            raw.flush()
            os.fsync(raw.fileno())
    checked = verify(output)
    write_json(output.with_suffix('.verification.json'), checked)
    return checked


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('action', choices=('package', 'verify'))
    parser.add_argument('--archive', type=Path, required=True)
    parser.add_argument('--bundle', type=Path)
    parser.add_argument('--training-audit', type=Path)
    args = parser.parse_args()
    if args.action == 'package':
        if args.bundle is None or args.training_audit is None:
            parser.error('package requires --bundle and --training-audit')
        result = package(args.archive.resolve(), args.bundle.resolve(), args.training_audit.resolve())
    else:
        result = verify(args.archive.resolve())
    print(json.dumps(result, sort_keys=True))


if __name__ == '__main__':
    main()