"""Export complete review artifacts to lossless non-executable tensor trees."""
from pathlib import Path
import argparse,hashlib,json,shutil,time
import torch
from safe_tree import save_tree,load_tree


def sha(path):return hashlib.sha256(path.read_bytes()).hexdigest()


def compare(a,b):
    if torch.is_tensor(a):assert torch.is_tensor(b) and torch.equal(a,b);return 1
    if isinstance(a,dict):assert list(a)==list(b);return sum(compare(a[k],b[k]) for k in a)
    if isinstance(a,(tuple,list)):assert type(a)==type(b) and len(a)==len(b);return sum(compare(x,y) for x,y in zip(a,b))
    assert a==b;return 0


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--source',default=str(Path(__file__).resolve().parent));parser.add_argument('--output',required=True);parser.add_argument('--device',default='cuda:3');args=parser.parse_args()
    source=Path(args.source).resolve();out=Path(args.output).resolve();assert not out.exists();out.mkdir(parents=True);torch.cuda.set_device(args.device);start=time.perf_counter();trees=[];tensor_count=0
    for name in ['run_01','objective_extension_01']:
        directory=out/name;directory.mkdir()
        for file in sorted((source/name).iterdir()):
            if file.suffix=='.pt':
                original=torch.load(file,map_location=args.device,weights_only=True);base=directory/file.stem;save_tree(original,base)
                restored=load_tree(base,args.device);count=compare(original,restored);tensor_count+=count
                trees.append(dict(file=str(base.relative_to(out)),source_sha256=sha(file),tensors=count,json_sha256=sha(base.with_suffix('.json')),safetensors_sha256=sha(base.with_suffix('.safetensors'))))
            elif file.suffix in ('.json','.csv','.md','.py'):shutil.copyfile(file,directory/file.name)
    for name in ['analysis_spec.json','objective_extension_spec.json','findings.json','findings.md','primary_table.csv','full_replay.json','verify_all.py','safe_tree.py','objective_extension.py','export_safe.py']:
        shutil.copyfile(source/name,out/name)
    trainer=source.parents[2]/'reproduce/review_r1_sensor.py';shutil.copyfile(trainer,out/'review_r1_sensor.py')
    assert not list(out.rglob('*.pt'))
    report=dict(format='Lossless Safetensors + JSON tensor trees; no executable pickle checkpoints',trees=trees,roundtrip_tensor_count=tensor_count,roundtrip_maximum_difference=0,seconds=time.perf_counter()-start,
                source_artifacts_modified=False,device=args.device,export_uses_cpu_only_for_serialization=True,
                replay_command='python verify_all.py --release /path/to/original-release --base run_01 --extension objective_extension_01 --output replay.json --device cuda:0')
    (out/'EXPORT_MANIFEST.json').write_text(json.dumps(report,indent=2)+'\n');print(json.dumps({k:v for k,v in report.items() if k!='trees'},indent=2))


if __name__=='__main__':main()