#!/usr/bin/env python3
"""Audit stage 2 in the existing pinned Lean environment, without editing sibling manifests."""
import argparse
import hashlib
import json
from pathlib import Path
import re
import subprocess
import tempfile
from check_proofs import PROJECT, checked, parse_axioms


def declarations(path, namespace):
    text = path.read_text()
    if re.search(r'\b(sorry|admit|axiom|native_decide)\b|debug\.skipKernelTC', text):
        raise ValueError(f'disallowed proof construct in {path.name}')
    if re.findall(r'^namespace (\w+)$', text, re.M) != [namespace]:
        raise ValueError('unexpected namespace')
    return [namespace+'.'+n for n in re.findall(r'^theorem (\w+)', text, re.M)]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--certificates', type=Path, help='generated Lean certificate file')
    args = parser.parse_args()
    source = PROJECT/'VibeQCStage2.lean'
    manifest = PROJECT/'stage2-claims.json'
    audit = PROJECT/'Stage2Audit.lean'
    files = [source, PROJECT/'VibeQCCore/Contractions.lean', manifest, audit, PROJECT/'lean-toolchain', PROJECT/'lake-manifest.json',
             Path(__file__).resolve(), Path(__file__).with_name('check_proofs.py')]
    names = declarations(source, 'VibeQCStage2')
    claims = json.loads(manifest.read_text())
    if sorted(names) != sorted(claims['theorems']):
        raise ValueError('stage-2 manifest does not match declarations')
    if sorted(re.findall(r'^#print axioms (\S+)', audit.read_text(), re.M)) != sorted(names):
        raise ValueError('stage-2 audit does not cover declarations')
    extra = []
    if args.certificates:
        args.certificates = args.certificates.resolve()
        extra = declarations(args.certificates, 'VibeQCStage2Cases')
        if not extra:
            raise ValueError('empty certificate file')
        files.append(args.certificates)
    snapshots = {p: p.read_bytes() for p in files}
    target = PROJECT/'.lake/build/lib/lean/VibeQCStage2.olean'
    target.parent.mkdir(parents=True, exist_ok=True)
    checked(['lake', 'build', 'VibeQCCore.Contractions'])
    checked(['lake', 'env', 'lean', '-o', str(target), str(source)])
    output = checked(['lake', 'env', 'lean', str(audit)])
    axioms = parse_axioms(output, names, set(claims['allowed_axioms']))
    if args.certificates:
        # Only this generated file is copied; no repository checkout is created.
        with tempfile.TemporaryDirectory(prefix='vibeqc-stage2-audit-') as tmp:
            path = Path(tmp)/'CertificateAudit.lean'
            path.write_bytes(snapshots[args.certificates] + ('\n'+''.join(
                '#print axioms '+name+'\n' for name in extra)).encode())
            output = checked(['lake', 'env', 'lean', str(path)])
        axioms.update(parse_axioms(output, extra, set(claims['allowed_axioms'])))
    for p, data in snapshots.items():
        if p.read_bytes() != data:
            raise ValueError('audit input changed during checking: '+str(p))
    print(json.dumps({'verdict': 'PASS', 'universal_claims': len(names),
        'certificate_claims': len(extra), 'axioms': {k: sorted(v) for k, v in axioms.items()},
        'lean_version': checked(['lake', 'env', 'lean', '--version']).strip(),
        'source_sha256': {str(p): hashlib.sha256(v).hexdigest() for p, v in snapshots.items()}}, indent=2))


if __name__ == '__main__':
    main()