#!/usr/bin/env python3
"""Regression test for always-on ray failure diagnostics.

Every INCOMPLETE ray endpoint must be reported on stderr with its reason and
the affected frame id / camera coordinate time even without ``--verbose``, so a
long movie never needs a rerun to be diagnosed.  ``--verbose`` (or any Debug
build) adds a bounded set of representative samples with film/cost
localization.  Budget-incomplete ``UNRESOLVED/BUDGET_EXHAUSTED`` frames are
reported separately from numerical ``INCOMPLETE`` failures, and the publication
refusal must not recommend a larger retry budget for a pure integration error.

A deterministic DP54 integration error is produced without a large workload:
strict tolerances, ``min == initial == max`` step, and one allowed rejection
make the first trial step fail.  The rejection quota is reached before any
minimum-step check, so the exact detailed reason is ``REJECTION_LIMIT``.  Small
16x8 scenes keep every case short.  Two-frame movie coverage uses the existing
``test_observer_<backend>`` fixture track duplicated to two rows (the
Schwarzschild metric is stationary, so the same tetrad is valid at both
times); the test then checks that each failing sample is counted once across
all time slabs.
"""
import os
import re
import struct
import subprocess
import sys
import tempfile
from pathlib import Path

# Keep scratch data inside the pre-approved OpenCode scratch directory.
TMP_ROOT = Path('/tmp/opencode')
TMP_ROOT.mkdir(parents=True, exist_ok=True)

BUILD = Path(sys.argv[1] if len(sys.argv) > 1 else 'build/Release').resolve()
TESTDIR = Path(sys.argv[2]).resolve() if len(sys.argv) > 2 else BUILD
ENV = dict(os.environ, OMP_NUM_THREADS='4')

# v3 lens-map layout (see src/lens_map.c); only enough to count stored
# INCOMPLETE vertices.  The first frame header follows the 40..176 v3
# provenance block.
VERSION_OFFSET = 8
FRAME_COUNT_OFFSET = 32
FRAME_HEADER_START = 176
VERTEX_SIZE = 108
OUTCOME_INDEX = 10
INCOMPLETE_OUTCOME = 3

COMMON = ['--catalog', 'assets/sky_grid_5deg.csv', '--width', 16, '--height', 8,
          '--fov-deg', 80, '--exposure', 1e-3, '--coarse-cell-pixels', 8,
          '--refine-max-level', 0, '--psf-relative-tail', 1e-4]

# Deterministic DP54 rejection-quota failure: tight tolerances, a single fixed
# step bound, and one allowed rejection.  Every ray rejects its first trial step
# and reports the exact REJECTION_LIMIT reason instead of a fabricated terminal
# category.
INJECT = ['--integrator', 'dp54',
          '--ode-rtol', '1e-15', '--ode-atol-x', '1e-15',
          '--ode-atol-pi', '1e-15', '--ode-atol-l', '1e-15',
          '--ode-initial-step', '0.5', '--ode-min-step', '0.5',
          '--ode-max-step', '0.5', '--ode-max-rejections', '1']


def run(binary, *args, ok=True, env=ENV):
    result = subprocess.run([str(binary), *map(str, args)], env=env,
                            capture_output=True, text=True)
    if (result.returncode == 0) != ok:
        raise AssertionError(
            f'{binary.name} {args}: rc={result.returncode}\n'
            f'stdout:\n{result.stdout}\nstderr:\n{result.stderr}')
    return result


def is_debug(result):
    return 'Debug build:' in result.stdout


def incomplete_total(stderr):
    return sum(int(m) for m in re.findall(r'INCOMPLETE=(\d+)', stderr))


def frame_reported(stderr, frame_id):
    return re.search(rf'Ray failures: frame {frame_id} camera_t=', stderr) is not None


def map_incomplete_vertices(path):
    data = path.read_bytes()
    assert data[:8] == b'GRLENS\x01\x00', f'not a lens map: {path}'
    version = struct.unpack_from('<I', data, VERSION_OFFSET)[0]
    assert version == 3, f'expected v3 map, got v{version}'
    frame_count = struct.unpack_from('<Q', data, FRAME_COUNT_OFFSET)[0]
    offset = FRAME_HEADER_START
    incomplete = 0
    for _ in range(frame_count):
        vertices, triangles = struct.unpack_from('<QQ', data, offset + 24)
        offset += 48
        for _ in range(vertices):
            fields = struct.unpack_from('<9dIIIQQQ', data, offset)
            offset += VERTEX_SIZE
            if fields[OUTCOME_INDEX] == INCOMPLETE_OUTCOME:
                incomplete += 1
        offset += triangles * 32 + 4  # triangles plus the frame payload CRC
    return incomplete


# A required renderer or fixture must not be skipped silently under make test.
binary = BUILD / 'schwarzschild_sky'
assert binary.exists(), f'missing {binary}'
observer_test = TESTDIR / 'test_observer_schwarzschild'
assert observer_test.exists(), f'missing {observer_test}'

help_text = run(binary, '--help').stdout
ext = 'png' if '.png' in help_text else 'ppm'
for option in ('--verbose', '--allow-incomplete', '--observer-track',
               '--movie-track-samples', '--lens-map-output', '--lens-map-input',
               '--refine-max-level'):
    assert option in help_text, (option, help_text)

with tempfile.TemporaryDirectory(prefix='gr-ray-diagnostics-',
                                 dir=str(TMP_ROOT)) as directory:
    tmp = Path(directory)

    # 1) A normal scene emits no ray-failure or budget diagnostics at all, so a
    #    clean run is not polluted by the always-on summary.  The refinement
    #    variant exercises the always-installed diagnostics callback across
    #    finished generations without producing any failure.
    normal = tmp / f'normal.{ext}'
    normal_run = run(binary, *COMMON, '--output', normal)
    assert normal_run.returncode == 0
    assert 'Ray failures:' not in normal_run.stderr, normal_run.stderr
    assert 'UNRESOLVED/BUDGET_EXHAUSTED' not in normal_run.stderr, normal_run.stderr
    refined = tmp / f'refined.{ext}'
    refined_run = run(binary, *COMMON, '--refine-max-level', '2',
                      '--refine-angle-abs-deg', '0.1', '--output', refined)
    assert refined_run.returncode == 0
    assert 'Ray failures:' not in refined_run.stderr, refined_run.stderr

    # 2) Non-verbose INCOMPLETE: the reason histogram, frame id, and camera time
    #    are on stderr even though --verbose was not passed.
    single = tmp / f'single.{ext}'
    single_run = run(binary, *COMMON, *INJECT, '--allow-incomplete',
                     '--output', single)
    err = single_run.stderr
    assert 'Ray failures:' in err, err
    assert 'INCOMPLETE=' in err and 'REJECTION_LIMIT' in err, err
    assert frame_reported(err, 0), err
    if not is_debug(single_run):
        assert 'ray failure:' not in err, err
        assert 'film=' not in err, err
    else:
        # A Debug build prints representative detail without --verbose.
        assert 'ray failure:' in err and 'film=' in err, err

    # 3) Verbose adds bounded representative samples with localization.
    verbose = tmp / f'verbose.{ext}'
    verbose_run = run(binary, *COMMON, *INJECT, '--allow-incomplete', '--verbose',
                      '--output', verbose)
    err = verbose_run.stderr
    assert 'ray failure:' in err, err
    assert 'reason=REJECTION_LIMIT' in err, err
    assert 'vertex=' in err and 'film=(' in err, err
    assert 'accepted=' in err and 'rejected=' in err and 'rhs=' in err, err
    # Bounded detail: at most RAY_DIAG_REPS_PER_REASON per reason, remainder
    # reported as suppressed.
    reps = sum(1 for line in err.splitlines()
               if 'ray failure: frame=0 reason=' in line
               and 'suppressed=' not in line)
    assert reps <= 3, err
    assert 'suppressed=' in err, err

    # 4) Refinement enabled with the failing trace: the always-installed
    #    diagnostics callback is part of the build, but the uniform strict
    #    injection fails every initial vertex, so no probe generation runs and
    #    only the initial scan reports (documented limit: probe/refinement
    #    witness failures are not reachable from this deterministic scene).
    refine_map = tmp / 'refine_inc.grlens'
    refine_run = run(binary, *COMMON, *INJECT, '--refine-max-level', '1',
                     '--refine-angle-abs-deg', '0.1', '--allow-incomplete',
                     '--lens-map-output', refine_map,
                     '--output', tmp / f'refine_inc.{ext}')
    err = refine_run.stderr
    assert frame_reported(err, 0), err
    stored = map_incomplete_vertices(refine_map)
    assert stored > 0
    assert incomplete_total(err) == stored, \
        f'refine histogram {incomplete_total(err)} != stored {stored}'

    # 5) Two-frame movie.  (a) A default, non-verbose run must identify both
    #    frames and count each failing sample exactly once across the time
    #    slabs, matching the saved map.
    track_single = tmp / 'track_single.csv'
    run(observer_test, track_single)
    rows = [line for line in track_single.read_text().splitlines()
            if line.strip() and not line.startswith('#')
            and not line[0].isalpha()]
    assert len(rows) == 1, rows
    fields = rows[0].split(',')
    fields[0], fields[1] = '1', '1'
    header = next(line for line in track_single.read_text().splitlines()
                  if line.startswith('t,'))
    track = tmp / 'track2.csv'
    track.write_text(header + '\n' + rows[0] + '\n' + ','.join(fields) + '\n')

    frames_dir = tmp / 'frames'
    frames_dir.mkdir()
    movie_map = tmp / 'movie.grlens'
    movie_run = run(binary, *COMMON, *INJECT, '--allow-incomplete',
                    '--slab-duration', '0.4', '--observer-track', track,
                    '--movie-track-samples', '--frames-dir', frames_dir,
                    '--lens-map-output', movie_map)
    err = movie_run.stderr
    assert frame_reported(err, 0) and frame_reported(err, 1), err
    stored = map_incomplete_vertices(movie_map)
    assert stored > 0, 'movie map has no INCOMPLETE vertices'
    assert incomplete_total(err) == stored, \
        f'histogram {incomplete_total(err)} != stored {stored}; ' \
        'a sample was re-counted across slabs'
    if not is_debug(movie_run):
        assert 'ray failure:' not in err, err

    # (b) Verbose movie adds request kind, persistent vertex id, film position
    #     and cost, bounded per reason.
    movie_verbose_dir = tmp / 'frames_verbose'
    movie_verbose_dir.mkdir()
    movie_verbose = run(binary, *COMMON, *INJECT, '--allow-incomplete',
                        '--verbose', '--slab-duration', '0.4',
                        '--observer-track', track, '--movie-track-samples',
                        '--frames-dir', movie_verbose_dir)
    err = movie_verbose.stderr
    assert 'sample=' in err and 'kind=vertex' in err and 'vertex=' in err, err
    assert 'film=(' in err, err
    assert 'accepted=' in err and 'rejected=' in err and 'rhs=' in err, err
    reps = sum(1 for line in err.splitlines()
               if 'ray failure:' in line and 'suppressed=' not in line)
    assert reps <= 6, err  # two frames, one reason each, <=3 reps per reason

    # (c) Render-only replay of the two-frame map without --allow-incomplete
    #     reports both frames' reasons before the publication gate refuses, and
    #     never invents an unpersisted trusted stop state.
    map_frames_dir = tmp / 'map_frames'
    map_frames_dir.mkdir()
    refused = run(binary, '--catalog', 'assets/sky_grid_5deg.csv',
                  '--lens-map-input', movie_map, '--frames-dir', map_frames_dir,
                  '--output', tmp / f'map_refused.{ext}', ok=False)
    err = refused.stderr
    assert frame_reported(err, 0) and frame_reported(err, 1), err
    assert 'REJECTION_LIMIT' in err, err
    assert 'phase=import' in err, err
    assert 'Incomplete render refused' in err, err
    assert err.index('frame 0') < err.index('Incomplete render refused'), err

    single_map = tmp / 'single.grlens'
    run(binary, *COMMON, *INJECT, '--allow-incomplete', '--lens-map-output',
        single_map, '--output', tmp / f'maplive.{ext}')
    replay = tmp / f'replay.{ext}'
    replay_run = run(binary, '--catalog', 'assets/sky_grid_5deg.csv',
                     '--lens-map-input', single_map, '--allow-incomplete',
                     '--verbose', '--output', replay)
    err = replay_run.stderr
    assert 'Ray failures:' in err and frame_reported(err, 0), err
    assert 'REJECTION_LIMIT' in err and 'phase=import' in err, err
    assert 'stop_t=' not in err and 'trusted=1' not in err, err

    # 6) Budget exhaustion is a distinct, always-on message, and its refusal
    #    still points at the retry budget.  A replay of an allowed budget map
    #    must not present the unpersisted continuation time as observed.
    budget_map = tmp / 'budget.grlens'
    run(binary, *COMMON, '--integrator', 'dp54',
        '--trace-lookback-time', '1e-6', '--retry-lookback-increment', '0',
        '--max-total-lookback-time', '1e-6', '--allow-incomplete',
        '--lens-map-output', budget_map, '--output', tmp / f'budget_allow.{ext}')
    replay_budget = run(binary, '--catalog', 'assets/sky_grid_5deg.csv',
                        '--lens-map-input', budget_map, '--allow-incomplete',
                        '--verbose', '--output', tmp / f'budget_replay.{ext}')
    err = replay_budget.stderr
    assert 'UNRESOLVED/BUDGET_EXHAUSTED' in err, err
    assert 'continuation_t=-' in err, err
    assert 'continuation_t=-1' not in err, \
        'replay invented an unpersisted continuation time'

    budget = tmp / f'budget.{ext}'
    budget_run = run(binary, *COMMON, '--integrator', 'dp54',
                     '--trace-lookback-time', '1e-6',
                     '--retry-lookback-increment', '0',
                     '--max-total-lookback-time', '1e-6',
                     '--output', budget, ok=False)
    err = budget_run.stderr
    assert 'UNRESOLVED/BUDGET_EXHAUSTED' in err, err
    assert re.search(r'UNRESOLVED/BUDGET_EXHAUSTED: frame 0 camera_t=', err), err
    assert 'blocking_triangles=' in err and 'unresolved_samples=' in err, err
    assert 'Incomplete render refused' in err, err
    assert 'budget' in err.lower(), err
    assert not budget.exists()

    print('ray diagnostics checks passed: always-on reasons + frame/time, '
          'bounded verbose samples, movie slab de-duplication and replay',
          flush=True)