"""Discover the first failure of the unchanged matched-root GPU test, then replay.

This script never runs the body/Jet test, changes the solver, or changes its
assertions. GPU admission remains the caller's responsibility.
"""
import argparse
import hashlib
import importlib.util
import json
from pathlib import Path
import subprocess
import sys
import traceback


def capture_case(exc):
    """Extract the original test frame; no inferred profile/block/configuration."""
    frames = []
    tb = exc.__traceback__
    while tb is not None:
        frames.append(tb.tb_frame)
        tb = tb.tb_next
    caller = next((f for f in frames if f.f_code.co_name == 'test_uniform_matched_root'), None)
    if caller is None or not any(f.f_code.co_name == 'assert_array_equal' for f in frames):
        return None
    local = caller.f_locals
    required = ('cfg', 'shifted', 'profile', 'clamp', 'damping', 'block', 'expected', 'actual')
    if not all(key in local for key in required):
        return None
    cfg = local['cfg']
    return dict(shifted=int(local['shifted']), profile=local['profile'],
                clamp=int(local['clamp']), damping=float(local['damping']),
                block=[int(v) for v in local['block']], steps=8,
                config_repr=repr(cfg)), local['expected'].copy(), local['actual'].copy()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--checkout', type=Path, default=Path.cwd())
    parser.add_argument('--output', type=Path, required=True, help='New directory; existing evidence is never overwritten')
    args = parser.parse_args()
    checkout, output = args.checkout.resolve(), args.output.resolve()
    output.mkdir(parents=True, exist_ok=False)
    import os
    os.chdir(checkout)
    sys.path.insert(0, str(checkout / 'src'))
    import numpy as np
    gate_path = checkout / 'tests/validation/multigrid/test_channel_root_gpu.py'
    spec = importlib.util.spec_from_file_location('original_channel_gpu_gate', gate_path)
    gate = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(gate)
    revision = subprocess.run(['git', 'rev-parse', 'HEAD'], capture_output=True, text=True, check=True).stdout.strip()
    manifest = dict(checkout=str(checkout), revision=revision,
                    gate_sha256=hashlib.sha256(gate_path.read_bytes()).hexdigest(),
                    test='ChannelRootGPU.test_uniform_matched_root')
    (output / 'manifest.json').write_text(json.dumps(manifest, indent=2))
    test = gate.ChannelRootGPU('test_uniform_matched_root')
    try:
        # The original method and np.testing.assert_array_equal are unchanged.
        # Its existing loops naturally stop at the first failing assertion.
        test.test_uniform_matched_root()
    except AssertionError as exc:
        original_trace = traceback.format_exc()
        (output / 'original_failure.txt').write_text(original_trace)
        print(original_trace, file=sys.stderr, flush=True)
        case = capture_case(exc)
        if case is None:
            (output / 'status.json').write_text(json.dumps(dict(status='different_assertion_failure', replay_started=False), indent=2))
            return 3
        config, expected, actual = case
        (output / 'failed_case.json').write_text(json.dumps(config, indent=2))
        np.savez(output / 'original_failure_arrays.npz', expected=expected, actual=actual)
        print('CAPTURED_CASE ' + json.dumps(config), flush=True)
    except Exception:
        (output / 'discovery_error.txt').write_text(traceback.format_exc())
        (output / 'status.json').write_text(json.dumps(dict(status='discovery_error', replay_started=False), indent=2))
        raise
    else:
        (output / 'status.json').write_text(json.dumps(dict(status='original_failure_not_reproduced', replay_started=False), indent=2))
        print('The original matched-root test completed without reproducing the failure. No targeted replay was started.', flush=True)
        return 2

    # Discovery's context has closed. Run targeted diagnostics in a fresh process,
    # serially, never concurrently with another compilation or GPU run here.
    replay = Path(__file__).resolve().with_name('replay.py')
    command = [sys.executable, str(replay), '--shifted', str(config['shifted']),
               '--profile', config['profile'], '--clamp', str(config['clamp']),
               '--damping', str(config['damping']), '--block', *map(str, config['block']),
               '--steps', str(config['steps']), '--output', str(output / 'targeted')]
    (output / 'replay_command.json').write_text(json.dumps(command, indent=2))
    with (output / 'targeted_stdout.txt').open('w') as stdout, (output / 'targeted_stderr.txt').open('w') as stderr:
        result = subprocess.run(command, cwd=checkout, stdout=stdout, stderr=stderr)
    targeted = output / 'targeted/result.json'
    details = json.loads(targeted.read_text()) if targeted.exists() else None
    status = dict(status='original_assertion_failure_retained', replay_started=True,
                  replay_returncode=result.returncode, replay_result=details)
    (output / 'status.json').write_text(json.dumps(status, indent=2))
    print(json.dumps(status), flush=True)
    # Preserve the original failed gate even if replay fails or does not reproduce.
    return 1


if __name__ == '__main__':
    raise SystemExit(main())