"""Stage only verified weights, reproducible tools and paired evaluation evidence."""
import argparse
import json
import shutil
from pathlib import Path
from scripts.common import REVISION,RUNTIME_REVISION,sha256,header,write_json
from scripts.report import REQUIRED_CASES,pair_records,records

def manifest(root):
    root=Path(root)
    return [dict(path=p.relative_to(root).as_posix(),size=p.stat().st_size,sha256=sha256(p))
            for p in sorted(root.rglob('*')) if p.is_file() and p.name!='MANIFEST.json'
            and not any(x.startswith('.') or x=='__pycache__' for x in p.relative_to(root).parts)]

def validate_manifest(root):
    root=Path(root).resolve()
    rows=json.loads((root/'MANIFEST.json').read_text())
    names=set()
    for row in rows:
        name=row['path']; original=root/name; path=original.resolve()
        if original.is_symlink() or not path.is_relative_to(root) or name in names:
            raise ValueError(f'Invalid manifest path: {name}')
        names.add(name)
        if path.stat().st_size!=row['size'] or sha256(path)!=row['sha256']:
            raise ValueError(f'Manifest mismatch: {name}')
    actual={p.relative_to(root).as_posix() for p in root.rglob('*') if p.is_file()
            and not any(x.startswith('.') or x=='__pycache__' for x in p.relative_to(root).parts)}
    if actual!=names|{'MANIFEST.json'}: raise ValueError('Unexpected/missing release files')
    return rows

def validate_evaluation(root):
    root=Path(root)
    pairs=list(pair_records(records(root/'bf16'),records(root/'8bit')))
    if {a['case_id'] for a,b in pairs}!=REQUIRED_CASES: raise ValueError('Incomplete seven-case suite')
    for a,b in pairs:
        if a['steps']!=40 or a['width']!=1024 or a['height']!=1024 or not a['warmup']:
            raise ValueError('Release requires the full 1024px/40-step warmed protocol')
        if b['mode']!='RGBA' or b['rgb_std']<1: raise ValueError('Broken output')
    review=json.loads((root/'visual-review.json').read_text())
    if review.get('approved_for_research_release') is not True:
        raise ValueError('Visual review incomplete or failed')
    if set(review.get('cases',{}))!=REQUIRED_CASES: raise ValueError('Missing visual case reviews')
    for name in ('bf16','8bit'):
        if review['results_sha256'][name]!=sha256(root/name/'results.jsonl'):
            raise ValueError('Visual review does not match current results')

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument('--model',type=Path,default=Path('models/Image21-MLX-8bit'))
    ap.add_argument('--evaluation',type=Path,default=Path('artifacts/eval'))
    ap.add_argument('--platform',choices=('huggingface','modelscope'),required=True)
    ap.add_argument('--output',type=Path)
    args=ap.parse_args(); out=args.output or Path('release')/args.platform
    if out.exists(): raise FileExistsError(out)
    validate_evaluation(args.evaluation)
    conversion=json.loads((args.model/'conversion.json').read_text())
    if conversion['status']!='converted_and_roundtrip_verified' or conversion['source_revision']!=REVISION or conversion['runtime_revision']!=RUNTIME_REVISION:
        raise ValueError('Unverified conversion')
    for comp in conversion['components'].values():
        if not comp['exact_roundtrip']: raise ValueError('Roundtrip not verified')
    for p in args.model.glob('*/*.safetensors'):
        if 'modification_notice' not in header(p).get('__metadata__',{}): raise ValueError('Missing modification notice')
    shutil.copytree(args.model,out,copy_function=shutil.copy2)
    shutil.copytree(args.evaluation,out/'evaluation',ignore=shutil.ignore_patterns('._*','*.tmp'))
    # Publish useful machine measurements without the workstation's absolute paths.
    for precision in ('bf16','8bit'):
        path=out/'evaluation'/precision/'environment.json'
        env=json.loads(path.read_text())
        env['model']=f'Image21-MLX-{precision}'
        write_json(path,env)
    audit=json.loads(Path('artifacts/source-audit.json').read_text())
    audit.pop('source')
    audit['source_model']='Qwen/Qwen-Image-2.1'
    audit['note']='Public source-file hashes and precision inventory; local source directory omitted.'
    write_json(out/'evaluation'/'source-files.json',audit)
    shutil.copytree('scripts',out/'scripts',ignore=shutil.ignore_patterns('__pycache__'))
    shutil.copytree('benchmarks',out/'benchmarks')
    for name in ('requirements.txt','requirements.lock.txt'):
        shutil.copy2(name,out/name)
    shutil.copy2(f'cards/{args.platform}.md',out/'README.md')
    shutil.copy2('artifacts/precision-probe.json',out/'evaluation'/'precision-probe.json')
    for name in ('cache-parity-bf16.json','cache-parity-8bit.json','offload-parity.json'):
        shutil.copy2(Path('artifacts')/name,out/'evaluation'/name)
    for row in json.loads((args.evaluation/'runtime-manifest.json').read_text()):
        if sha256(out/row['path'])!=row['sha256']:
            raise ValueError(f'Runtime changed after evaluation: {row["path"]}')
    write_json(out/'MANIFEST.json',manifest(out))
    rows=validate_manifest(out)
    print(f'Staged and verified {len(rows)} files in {out}')

if __name__=='__main__': main()