"""Finite fixed inventory; no overwrites, no automatic scientific retries."""
import argparse
import json
import os
import time
import traceback
import urllib.request
from context import ROOT, MODEL, EXP, AK, CFG, read, save, digest, check, budget


def preflight():
    os.environ['CUDA_VISIBLE_DEVICES'] = CFG['gpu_uuid']
    os.environ['CUDA_HOME'] = '/usr/local/cuda'
    os.environ['PATH'] = str(ROOT / 'models/male-cns-lif/.venv/bin') + ':/usr/local/cuda/bin:' + os.environ['PATH']
    os.environ['MAX_JOBS'] = '2'
    os.environ['TORCH_CUDA_ARCH_LIST'] = '8.9'
    for port in (8188, 8189, 8190):
        with urllib.request.urlopen(f'http://192.168.10.51:{port}/queue', timeout=5) as response:
            q = json.load(response)
        assert not q['queue_running'] and not q['queue_pending'], 'ComfyUI busy: defer, never evict'


def output_inventory(folder):
    return [dict(path=p.relative_to(folder).as_posix(), bytes=p.stat().st_size, sha256=digest(p))
        for p in sorted(folder.rglob('*')) if p.is_file() and p.name != 'output-hashes.json']


def compare_expected(stage, row, folder):
    from comparison import compare_runs, validate_run
    case = AK / 'cases' / row['cohort'] / row['name']
    if stage == 'zero':
        reference = AK / 'gpu' / row['cohort'] / (row['name'] + '-r' + str(row['repeat']))
        return compare_runs(case, reference, folder, mode='zero')
    if row['repeat'] == 1:
        reference = EXP / 'gpu' / stage / (row['name'] + '-r0')
        return compare_runs(case, reference, folder, mode='repeat')
    return validate_run(case, folder)


def verify_identity(folder, stage, row):
    assert stage in CFG['stages'] and row['repeat'] in (0, 1)
    assert folder == EXP / 'gpu' / stage / (row['name'] + '-r' + str(row['repeat']))
    case = AK / 'cases' / row['cohort'] / row['name']
    summary = read(folder / 'summary.json')
    assert summary['case'] == row['name'] and summary['complete'] and summary['pass_all']
    assert summary['input_sha256']['spec.json'] == digest(case / 'spec.json')
    assert summary['input_sha256']['requests.npz'] == digest(case / 'requests.npz')
    overlay = read(folder / 'overlay-audit.json')
    assert overlay['complete'] and overlay['overlay_name'] == stage
    assert overlay['input']['overlay_sha256'] == digest(EXP / 'overlays' / (stage + '.npz'))
    assert overlay['input']['manifest_sha256'] == digest(EXP / 'overlays' / (stage + '.json'))
    report = read(folder / 'comparison.json')
    assert report['pass_all'] and report['overlay_name'] == stage
    assert report['case_dir'] == str(case) and report['right'] == str(folder)
    if stage == 'zero':
        assert report['mode'] == 'zero'
        assert report['left'] == str(AK / 'gpu' / row['cohort'] / (row['name'] + '-r' + str(row['repeat'])))
    elif row['repeat'] == 1:
        assert report['mode'] == 'repeat' and report['left'] == str(EXP / 'gpu' / stage / (row['name'] + '-r0'))
    else:
        assert report['mode'] == 'self' and report['left'] == str(folder)


def verify_completed(folder, stage=None, row=None, recompute=False):
    rows = read(folder / 'output-hashes.json')
    assert rows == output_inventory(folder), 'Completed output bytes changed: ' + str(folder)
    assert read(folder / 'comparison.json')['pass_all']
    assert read(folder / 'summary.json')['complete']
    if stage is not None:
        verify_identity(folder, stage, row)
        if recompute:
            report = compare_expected(stage, row, folder)
            assert report['pass_all'] and report == read(folder / 'comparison.json'), 'Recomputed comparison differs'


def run(stage):
    assert stage in CFG['stages']
    reg = check(full=True); budget()
    assert not list((EXP / 'audit').glob('stopped-*.json')), 'Previous failure: preserve and inspect, never silently retry'
    receipt = EXP / 'audit' / (stage + '-complete.json')
    assert not receipt.exists(), 'Stage already completed'
    if stage == 'static25':
        prior = read(EXP / 'audit/zero-complete.json')
        assert prior['pass_all'] and prior['trials'] == 144
        for row in reg['planned']:
            verify_completed(EXP / 'gpu/zero' / (row['name'] + '-r' + str(row['repeat'])), 'zero', row)
    preflight()
    from adapter import run_gpu_overlay
    began = time.perf_counter()
    completed = []
    try:
        for ordinal, row in enumerate(reg['planned'], 1):
            key = row['name'] + '-r' + str(row['repeat'])
            folder = EXP / 'gpu' / stage / key
            case = AK / 'cases' / row['cohort'] / row['name']
            if folder.exists():
                verify_completed(folder, stage, row, recompute=True)
                print('VERIFIED-COMPLETED', stage, ordinal, key, flush=True)
            else:
                preflight(); check(); budget()
                print('START', stage, ordinal, key, flush=True)
                summary = run_gpu_overlay(case, folder, EXP / 'extension-build',
                    EXP / 'overlays' / (stage + '.npz'), EXP / 'overlays' / (stage + '.json'))
                assert summary['pass_all'] and summary['complete']
                report = compare_expected(stage, row, folder)
                save(folder / 'comparison.json', report)
                assert report['pass_all'], 'Independent recording/equivalence gate failed: ' + key
                save(folder / 'output-hashes.json', output_inventory(folder))
                verify_completed(folder, stage, row)
                print('COMPLETE', stage, ordinal, key, 'simulation_seconds', summary['diagnostic_simulation_seconds'], flush=True)
            completed.append(dict(name=row['name'], repeat=row['repeat'],
                output_hashes_sha256=digest(folder / 'output-hashes.json'),
                comparison_sha256=digest(folder / 'comparison.json')))
        assert len(completed) == 144
        check(full=True)
        result = dict(stage=stage, pass_all=True, trials=144, primary_trials=120, technical_repeats=24,
            completed=completed, wall_seconds=time.perf_counter() - began, training=False,
            scope='All zero-delta runs exactly match sealed AK' if stage == 'zero' else
                  'Full closed-loop fixed overlay, internal event/queue accounting and 24 exact technical repeats; not learning')
        save(receipt, result); budget()
        return result
    except Exception as error:
        save(EXP / 'audit' / ('stopped-' + stage + '.json'), dict(stage=stage,
            completed_before_error=completed, type=type(error).__name__, message=str(error),
            traceback=traceback.format_exc(), partial_outputs_preserved=True, no_parameter_rescue=True))
        raise


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('stage', choices=CFG['stages'])
    run(parser.parse_args().stage)