"""Audit saved functional evidence and archive only outputs, never checkpoint weights."""
import argparse
import csv
import hashlib
import json
from pathlib import Path
import platform
import subprocess
import time
import numpy as np
from .model import StreamingLM
from .norm_repetition import uniform_bounds
from .run import write_json

EXCLUDE={'tokenizer.json','reference.py','checksums-functional.json','audit-validation.json','archive-files.json'}


def archive_files(root):
    files=[p for p in root.iterdir() if p.is_file() and p.name not in EXCLUDE and not p.name.endswith(('.partial','.safetensors'))]
    files += [p for p in (root/'figures').glob('*') if p.is_file()]
    return sorted(files)


def main():
    p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);a=p.parse_args();root=a.root
    manifest=json.loads((root/'manifest.json').read_text())
    assert manifest['revision']=='15852e8c16360a2fea060d615a32b45270f8a8fc'
    for f in manifest['files']:
        if f['name'].endswith('.safetensors'):
            path=root/'weights'/f['name']
            assert path.is_file() and path.stat().st_size==f['bytes']
            assert f['sha256']==f['upstream_sha256']
    retention=json.loads((root/'retention-policy.json').read_text());assert retention['automatic_cleanup'] is False
    reference=json.loads((root/'reference-check.json').read_text())
    assert reference['passed'] and len(reference['tests'])==22
    assert reference['implementation_sha256']==hashlib.sha256(Path('functional/model.py').read_bytes()).hexdigest()
    tests=(root/'final-tests.txt').read_text();assert 'Ran 12 tests' in tests and '\nOK\n' in tests
    result=json.loads((root/'norm-repetition-summary.json').read_text());assert result['state']=='completed'
    assert result['learned_function_repetition_positive'] and result['bidirectional_network_quality_pass']
    cfg=json.loads((root/'model_config.json').read_text());lm=StreamingLM(root/'weights',cfg,seconds=60)
    c=result['selected'];ga=1+lm.small(f'layers.{c["a"]}.{c["role"]}.weight');gb=1+lm.small(f'layers.{c["b"]}.{c["role"]}.weight')
    bounds=uniform_bounds(ga,gb)
    np.testing.assert_allclose(bounds,[c['a_to_b_uniform_bound'],c['b_to_a_uniform_bound']],atol=1e-14,rtol=1e-14)
    with (root/'norm-repetition-uniform-bounds.csv').open() as f:assert len(list(csv.DictReader(f)))==552
    baseline=np.load(root/'baseline-heldout-token-nll.npy');assert baseline.shape==(4,31)
    for rec in result['interventions']:
        nll=np.load(root/f'norm-repetition-{rec["target"]}-uses-{rec["donor"]}-token-nll.npy')
        delta=nll.mean(1)-baseline.mean(1)
        np.testing.assert_allclose(delta,rec['sequence_nll_deltas'],atol=1e-13,rtol=1e-13)
        assert abs(delta.mean()-rec['mean_nll_delta'])<1e-13
    with (root/'norm-repetition-function-metrics.csv').open() as f:rows=list(csv.DictReader(f))
    held=[r for r in rows if r['split']=='heldout'];assert len(held)==8
    worst=max(max(float(r['relative_l2']),float(r['reverse_relative_l2'])) for r in held)
    assert abs(worst-result['worst_heldout_relative_l2'])<1e-14
    neuron=json.loads((root/'neuron-reuse-search-summary.json').read_text());assert sum(x['eligible_unique_pairs_checked'] for x in neuron)==7450139
    ledger=json.loads((root/'compute-budget-ledger.json').read_text());assert ledger['spent_seconds']+600<7200
    environment={'python':platform.python_version(),'platform':platform.platform(),'numpy':np.__version__,
                 'packages':subprocess.run([str(Path(__import__('sys').executable)),'-m','pip','freeze'],capture_output=True,text=True,check=True).stdout.splitlines(),
                 'weights_retained':True,'no_automatic_cleanup':True,'ledger_seconds_including_initial_reserve':ledger['spent_seconds'],
                 'final_tests_report_audit_additional_budget_upper_reserve_seconds':600,
                 'historical_error_logs':'Early reference adapter errors and preflight pressure stops; latest summaries/tests are authoritative.'}
    write_json(root/'functional-environment.json',environment)
    source_paths=[*Path('functional').glob('*.py'),Path('functional/reference.json'),Path('pilot/tensors.py'),Path('tests/test_functional.py')]
    write_json(root/'final-source-hashes.json',{str(p):hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted(source_paths)})
    files=archive_files(root)
    # Every numeric CSV value must be finite. Empty diagnostic fields stay unknown.
    for path in files:
        if path.suffix=='.csv':
            with path.open() as f:
                for row in csv.DictReader(f):
                    for value in row.values():
                        if value in (None,''):continue
                        try:number=float(value)
                        except ValueError:continue
                        assert np.isfinite(number),(path,value)
    checksums={str(p.relative_to(root)):hashlib.sha256(p.read_bytes()).hexdigest() for p in files}
    write_json(root/'checksums-functional.json',checksums)
    selected=[*checksums,'checksums-functional.json','audit-validation.json','archive-files.json']
    write_json(root/'archive-files.json',selected)
    write_json(root/'audit-validation.json',{'passed':True,'validated_unix':time.time(),'checksummed_files':len(checksums),
               'proof_recomputed_from_retained_weights':True,'heldout_nll_recomputed_from_saved_arrays':True,
               'weights_present_and_retained':True,'archive_excludes_weights_and_tokenizer':True})
    print(json.dumps({'passed':True,'checksummed_files':len(checksums),'weights_retained':True}))


if __name__=='__main__':main()