"""Audit historical artifacts without calling a model or using gold answers."""
import argparse
import ast
import hashlib
import importlib.util
import json
from pathlib import Path

HERE = Path(__file__).resolve().parent


def load_module(name, path):
    spec = importlib.util.spec_from_file_location(name, path)
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return m


def check(condition, message):
    if not condition:
        raise ValueError(message)


def audit(data=None):
    merger = load_module('v9_merge', HERE / 'training/ensemble_v9.py')
    runner = load_module('v9_runner', HERE / 'inference.py')
    manifest = json.loads((HERE / 'submission_manifest.json').read_text())
    submitted = (HERE / 'predictions.jsonl').read_bytes()
    check(hashlib.sha256(submitted).hexdigest() == manifest['predictions_sha256'], 'Submitted prediction hash changed')
    maps = {}
    for name, entry in manifest['passes'].items():
        path = HERE / entry['file']
        check(hashlib.sha256(path.read_bytes()).hexdigest() == entry['sha256'], f'Pass hash changed: {name}')
        maps[name] = merger.load_map(path)
        check(len(maps[name]) == entry['rows'], f'Pass row count changed: {name}')
    records = [json.loads(line) for line in submitted.splitlines() if line.strip()]
    keys = [merger.key_of(r) for r in records]
    check(len(keys) == len(set(keys)), 'Duplicate submitted key')
    check(all(r['answer'] in {'1','2','3','4'} for r in records), 'Invalid submitted answer')
    check(set(keys) == set(maps['p1']), 'Pass 1 and submission keys differ')
    check(all(set(m) <= set(keys) for m in maps.values()), 'Unknown key in a pass')
    pos = {tuple(k) for k in json.loads((HERE / 'training/passes/positional_option_keys.json').read_text())}
    check(pos <= set(keys), 'Unknown positional key')
    rebuilt = ''.join(json.dumps({**r, 'answer': merger.merge_one(merger.key_of(r), pos, **maps)}, ensure_ascii=False)+'\n' for r in records).encode()
    check(rebuilt == submitted, 'Archival vote does not reproduce submitted bytes')
    check(all(runner.merge_one(k, pos, **maps) == merger.merge_one(k, pos, **maps) for k in keys), 'Runner and archive merger disagree')
    result = {'submission_id': manifest['codabench_submission_id'], 'rows':len(records), 'sha256':manifest['predictions_sha256'],
              'saved_pass_merge_byte_identical':True, 'pass_counts':{k:len(v) for k,v in maps.items()},
              'fresh_model_inference_verified':False, 'rules_compliance_certified':False}
    if data:
        qs = runner.prepare(runner.load_questions('', local_data=data))
        check({runner.kof(q) for q in qs} == set(keys), 'Dataset key set differs')
        actual_pos = {runner.kof(q) for q in qs if runner.POS_RE.search(' '.join(map(str,q['choices'])))}
        check(actual_pos == pos, 'Recorded positional filter differs from input text')
        priority = runner.priority_order(qs, maps['cot'], maps['dr'])
        t2 = {runner.kof(q) for q in priority if q['_level']=='hard' or maps['cot'][runner.kof(q)] != maps['dr'][runner.kof(q)]}
        t3 = {runner.kof(q) for q in qs if (maps['p1'][runner.kof(q)] != maps['p2'][runner.kof(q)] if runner.kof(q) in maps['p2'] else maps['p1'][runner.kof(q)] != maps['cot'][runner.kof(q)] or maps['p1'][runner.kof(q)] != maps['dr'][runner.kof(q)])}
        check(t2 == set(maps['p2']), 'Pass 2 target rule mismatch')
        check(t3 == set(maps['p3']), 'Pass 3 target rule mismatch')
        mac_rows = [json.loads(l) for l in (HERE / manifest['passes']['mac']['file']).read_text().splitlines() if l.strip()]
        check([merger.key_of(r) for r in mac_rows] == [runner.kof(q) for q in priority[:len(mac_rows)]], 'Mac prefix does not match original scheduler')
        source = ast.parse((HERE / 'training/evidence/original_cell5_pass4_pass5.py').read_text())
        literal = None
        for node in source.body:
            if isinstance(node, ast.Assign) and any(isinstance(t, ast.Name) and t.id=='TARGET_IDX' for t in node.targets):
                literal = next(n.value for n in ast.walk(node.value) if isinstance(n,ast.Constant) and isinstance(n.value,str) and ',' in n.value and len(n.value) > 10)
        check(literal is not None, 'Original target indices not found')
        indices = {int(v.strip()) for v in literal.split(',') if v.strip()}
        canonical = sorted(qs, key=lambda q:(runner.LEVEL_RANK[q['_level']],q['_file']))
        check({runner.kof(q) for j,q in enumerate(canonical) if j in indices} == t3, 'Original literal target set differs')
        t45 = t3 - pos
        for n in ['p4','p5']:
            check(set(maps[n]) <= t45, f'Unexpected {n} key')
        examples = json.loads((HERE / 'data/few_shot_examples.json').read_text())
        qtexts = {' '.join(q['question'].split()) for q in qs}
        overlap = [{'category':cat,'index':j} for cat,exs in examples.items() for j,ex in enumerate(exs) if ' '.join(ex['question'].split()) in qtexts]
        result['input_checks'] = {'pass2_targets':len(t2),'pass3_targets':len(t3),'literal_notebook_targets_match':True,
                                 'pass45_targets':len(t45),'mac_priority_prefix_match':True,
                                 'omitted_keys':{n:sorted([list(k) for k in t45-set(maps[n])],key=str) for n in ['p4','p5']},
                                 'few_shot_exact_text_overlaps':overlap,
                                 'overlap_check_limit':'Exact normalized question text only; not semantic overlap or contamination certification.'}
    return result


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--data', help='optional local folder containing easy/medium/hard question JSON')
    ap.add_argument('--report', help='optional JSON report path')
    args = ap.parse_args()
    result = audit(args.data)
    text = json.dumps(result, ensure_ascii=False, indent=2)+'\n'
    if args.report:
        Path(args.report).write_text(text)
    print(text)


if __name__ == '__main__':
    main()