"""Audit and package only science outputs; retain all weights/profile caches."""
import argparse
import csv
import hashlib
import json
from pathlib import Path
import shutil
import time
import numpy as np
from functional.run import write_json
from .protocol import LAYERS,PAIRS
from .core import assignment


def digest(path):
    h=hashlib.sha256()
    with path.open('rb') as f:
        while b:=f.read(4*1024*1024):h.update(b)
    return h.hexdigest()


def main():
    p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);p.add_argument('--source',type=Path,required=True);p.add_argument('--output',type=Path,required=True);a=p.parse_args();root=a.root.resolve();source=a.source.resolve();out=a.output.resolve();assert root!=out
    status=json.loads((root/'status.json').read_text());assert status['state']=='completed' and status['last_source_layer']==23 and status['last_pair']==[19,23]
    expected={'full_attention':15,'qkv_head':960,'qk_map':30,'qk_head':960,'v_group':60};assert status['candidate_counts']==expected
    summary=json.loads((root/'calibration-summary.json').read_text());assert len(summary)==2025
    assert not status['selected'] and all(not x['strict_pass'] and not x['approximate_pass'] for x in summary)
    selection=json.loads((root/'selection.json').read_text());assert selection['qk_head_prequalified']==selection['v_group_prequalified']==0
    from collections import defaultdict
    rawout=defaultdict(list)
    for r in csv.DictReader((root/'calibration-output-metrics.csv').open()):rawout[r['candidate']].append(r)
    for c in summary:
        if c['output_verified']:
            rr=rawout[c['id']];assert len(rr)==c['output_observations']
            worst=max(max(float(r['error_a']),float(r['error_b'])) for r in rr)
            assert abs(worst-c['worst_output_error'])<1e-12 and worst>.20
    assert min(c['worst_mean_js'] for c in summary if c['kind']=='qk_head')>.05
    assert min(c['worst_value_error'] for c in summary if c['kind']=='v_group')>.20

    rows=list(csv.DictReader((root/'head-assignment-costs.csv').open()));assert len(rows)==960
    for i,j in PAIRS:
        r=[x for x in rows if int(x['i'])==i and int(x['j'])==j];cost=np.empty((8,8));chosen=[]
        for x in r:
            ha,hb=int(x['head_i']),int(x['head_j']);cost[ha,hb]=float(x['calibration_mean_js'])
            if x['chosen']=='True':chosen.append((ha,hb))
        assert len(chosen)==8;assert [b for _,b in sorted(chosen)]==assignment(cost)
    tests=(root/'final-tests.txt').read_text();assert 'Ran 20 tests' in tests and '\nOK\n' in tests
    cond=json.loads((root/'conditional-status.json').read_text());assert cond['state']=='completed' and cond['fresh_function_pass'] is False
    case=json.loads((root/'conditional-case-audit.json').read_text());assert case['passed_recomputation'] and case['no_success_relabeling']
    assert len(case['strict_failures'])==3 and case['network_intervention_not_run_due_to_fresh_gate_failure']
    assert cond['shared_budget_seconds_including_60s_early_reserve']+600<7200
    prob=np.load(root/'conditional-fresh-probabilities.npz');target=prob['target_layer23_head6'];donor=prob['donor_layer3_head7_on_H23']
    from .core import distribution_metrics
    from .conditional import tv_metrics
    metric={**distribution_metrics(target,donor),**tv_metrics(target,donor)}
    fresh=[r for r in csv.DictReader((root/'conditional-function-metrics.csv').open()) if r['split']=='fresh_test']
    assert len(fresh)==4
    for k,values in metric.items():
        if k in fresh[0] and values[0] is not None:np.testing.assert_allclose(values,[float(r[k]) for r in fresh],atol=1e-12,rtol=1e-12)
    for file in root.glob('*.csv'):
        with file.open() as f:
            for r in csv.DictReader(f):
                for value in r.values():
                    if value in ('',None):continue
                    try:x=float(value)
                    except ValueError:continue
                    assert np.isfinite(x),(file,value)
    weights=json.loads((source/'manifest.json').read_text());assert weights['revision']=='15852e8c16360a2fea060d615a32b45270f8a8fc'
    for entry in weights['files']:
        if entry['name'].endswith('.safetensors'):assert (source/'weights'/entry['name']).stat().st_size==entry['bytes'] and entry['sha256']==entry['upstream_sha256']
    assert json.loads((source/'retention-policy.json').read_text())['automatic_cleanup'] is False
    reference=json.loads((source/'reference-check.json').read_text());assert reference['implementation_sha256']==digest(Path('functional/model.py'))
    write_json(root/'final-source-hashes.json',{str(file):digest(file) for file in [*sorted(Path('crossqkv').glob('*.py')),Path('crossqkv/fresh_samples.json'),Path('tests/test_crossqkv.py'),Path('functional/model.py')]})
    caches=[]
    for file in sorted((root/'cache').glob('*')):
        if file.is_file():caches.append({'server_path':str(file),'bytes':file.stat().st_size,'sha256':digest(file),'retained':True})
    old_cache=root.parent/'cross-layer-qkv'/'cache'/'calibration-source-3.npz'
    if old_cache.exists():caches.append({'server_path':str(old_cache),'bytes':old_cache.stat().st_size,'sha256':digest(old_cache),'retained':True,'failed_attempt_cache':True})
    write_json(root/'external-caches.json',{'host':'8.211.154.134','source_weights_directory':str(source/'weights'),'source_baseline_cache':str(source/'baseline-activations.npz'),
               'files':caches,'automatic_deletion':False,'all_figure_source_tables_and_selected_probabilities_in_git':True})
    (root/'DELIVERY.md').write_text('# Cross-layer QKV scientific delivery\n\nAll plot CSV/JSON, the selected fresh probability matrices, protocols, reference tests and failure logs are included.\nLarge intermediate profiles remain on 8.211.154.134 with locations/sizes/SHA256 in external-caches.json.\nThe fixed checkpoint and all caches are retained until explicit user deletion. No weights are committed.\nThe primary protocol has no passing candidate. The later fixed conditional audit is exploratory and its strict fresh checks failed; observed readout overlap is reported without relabeling it a complete pass.\n')
    out.mkdir(exist_ok=True,parents=True)
    files=[f for f in root.iterdir() if f.is_file() and f.name not in ('checksums.json','audit-validation.json','archive-files.json')]
    files += list((root/'figures').glob('*'))
    for file in files:
        dest=out/file.relative_to(root);dest.parent.mkdir(exist_ok=True,parents=True);shutil.copy2(file,dest)
    checks={str(f.relative_to(out)):digest(f) for f in sorted(out.rglob('*')) if f.is_file() and f.name not in ('checksums.json','audit-validation.json','archive-files.json')}
    write_json(out/'checksums.json',checks)
    write_json(out/'archive-files.json',[*checks,'checksums.json','audit-validation.json','archive-files.json'])
    write_json(out/'audit-validation.json',{'passed':True,'validated_unix':time.time(),'checksummed_files':len(checks),'primary_counts_rechecked':expected,
               'fresh_probabilities_independently_recomputed':True,'no_unrun_network_quality_claims':True,'all_weights_and_caches_retained':True})
    print(json.dumps({'passed':True,'checksummed_files':len(checks),'external_cache_files':len(caches),'weights_retained':True}))


if __name__=='__main__':main()