#!/usr/bin/env python3
"""Real camera + TensorRT record/replay check. Runs only explicitly on Orin.
Records to a new test directory; never stops an existing deployment process.
"""
import argparse
import hashlib
import json
import os
from pathlib import Path
import signal
import subprocess
import time


def rows(session):
    index = session / 'frames.txt'
    return index.read_text().splitlines() if index.exists() else []


def wait_frames(process, session, count, timeout=90):
    deadline = time.monotonic() + timeout
    while len(rows(session)) < count:
        if process.poll() is not None:
            raise RuntimeError(f'Inference exited early: {process.returncode}')
        if time.monotonic() > deadline:
            raise RuntimeError('Timed out awaiting real inference frames')
        time.sleep(.05)


def stop(process, signum=signal.SIGTERM):
    start = time.monotonic()
    process.send_signal(signum)
    code = process.wait(timeout=15)
    if code:
        raise RuntimeError(f'Nonzero stop exit: {code}')
    return time.monotonic() - start


def run_record(command, session, log, count, signum=signal.SIGTERM):
    with log.open('w') as output:
        process = subprocess.Popen(command + ['--record-session', str(session)], stdout=output, stderr=subprocess.STDOUT)
        try:
            wait_frames(process, session, count)
            return stop(process, signum)
        finally:
            if process.poll() is None:
                process.terminate()
                process.wait(timeout=15)


def load(session):
    return [json.loads((session / 'frames' / f'{stem}.json').read_text()) for stem in rows(session)]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--binary', type=Path, required=True)
    parser.add_argument('--config', type=Path, required=True)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--replay-input', type=Path, help='Verify deployment with an existing session; no camera access')
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=False)
    binary = args.binary.resolve()
    config = args.config.resolve()
    command = [str(binary), '--config', str(config)]
    # Refuse camera contention rather than alter an existing deployment.
    if subprocess.run(['pgrep', '-x', 'camera_risk'], stdout=subprocess.DEVNULL).returncode == 0:
        raise RuntimeError('An existing camera_risk process is running; preserve its launch parameters first')
    if args.replay_input:
        replayed = args.output / 'replayed'
        with (args.output/'replay.log').open('w') as log:
            process = subprocess.Popen(command + ['--replay-session', str(args.replay_input.resolve()), '--record-session', str(replayed)], stdout=log, stderr=subprocess.STDOUT)
            try:
                running_hash = hashlib.sha256(Path(f'/proc/{process.pid}/exe').read_bytes()).hexdigest()
                assert process.wait(timeout=120) == 0, 'Deployment replay failed'
            finally:
                if process.poll() is None:
                    process.terminate()
                    process.wait(timeout=15)
        original, restored = load(args.replay_input), load(replayed)
        assert [r['frame_index'] for r in original] == [r['frame_index'] for r in restored]
        result = {'binary': str(binary), 'binary_sha256': hashlib.sha256(binary.read_bytes()).hexdigest(),
                  'running_binary_sha256': running_hash, 'replayed_frames': len(restored)}
        assert result['binary_sha256'] == running_hash
        (args.output/'report.json').write_text(json.dumps(result, indent=2)+'\n')
        print(json.dumps(result, indent=2))
        return
    recorded, replayed = args.output / 'recorded', args.output / 'replayed'
    term_seconds = run_record(command, recorded, args.output / 'record.log', 12)
    original = load(recorded)
    with (args.output / 'replay.log').open('w') as log:
        process = subprocess.Popen(command + ['--replay-session', str(recorded), '--record-session', str(replayed)], stdout=log, stderr=subprocess.STDOUT)
        try:
            # Verify the running executable, not just the file on disk.
            running_hash = hashlib.sha256(Path(f'/proc/{process.pid}/exe').read_bytes()).hexdigest()
            if process.wait(timeout=120):
                raise RuntimeError('Replay failed')
        finally:
            if process.poll() is None:
                process.terminate()
                process.wait(timeout=15)
    restored = load(replayed)
    assert len(original) == len(restored) >= 12, 'EOF dropped frames'
    max_pose_error = 0.0
    matching_tracks = matching_decisions = 0
    for before, after, stem in zip(original, restored, rows(recorded)):
        assert before['frame_index'] == after['frame_index']
        assert before['capture'] == after['source_capture']
        if 'can' in before['result']:
            expected_can = dict(before['result']['can'], source='replay')
            assert expected_can == after['result']['can'], 'Recorded CAN state changed during replay'
        assert before['video_frame_index'] == after['video_frame_index']
        matching_tracks += before['result']['track_ids'] == after['result']['track_ids']
        matching_decisions += before['result']['risk']['decision'] == after['result']['risk']['decision']
        for stage in ('raw_keypoints', 'isolated_keypoints'):
            for a, b in zip(before['stages'][stage], after['stages'][stage]):
                for ka, kb in zip(a, b):
                    max_pose_error = max(max_pose_error, *(abs(x-y) for x, y in zip(ka, kb)))
    # H.264 is lossy: report algorithm differences, do not require bitwise equality.
    probe = json.loads(subprocess.check_output(['ffprobe', '-v', 'error', '-select_streams', 'v:0', '-show_entries', 'stream=codec_name,nb_frames,width,height', '-of', 'json', str(recorded/'frames.mp4')]))['streams'][0]
    assert probe['codec_name'] == 'h264' and int(probe['nb_frames']) == len(original)
    # Exercise cancellation while replay has queued input, and reuse afterwards.
    interrupt_seconds = run_record(command + ['--replay-session', str(recorded)], args.output/'interrupted', args.output/'interrupt.log', 1, signal.SIGINT)
    interrupted_count = json.loads((args.output/'interrupted/complete.json').read_text())['frame_count']
    assert 0 < interrupted_count < len(original), 'Cancellation did not stop before EOF'
    result = {'platform': os.uname().nodename, 'binary': str(binary), 'binary_sha256': hashlib.sha256(binary.read_bytes()).hexdigest(),
              'running_binary_sha256': running_hash, 'recorded_frames': len(original), 'replayed_frames': len(restored),
              'frames_with_people': sum(bool(row['result']['track_ids']) for row in original),
              'max_raw_pose_error_px': max_pose_error, 'matching_track_frames': matching_tracks, 'matching_decision_frames': matching_decisions,
              'video_bytes': (recorded/'frames.mp4').stat().st_size, 'uncompressed_bgr_bytes': len(original)*int(probe['width'])*int(probe['height'])*3, 'sigterm_seconds': term_seconds, 'replay_sigint_seconds': interrupt_seconds, 'interrupted_frames': interrupted_count,
              'can_statuses': sorted({r['result'].get('can', {}).get('status', 'NOT_RECORDED') for r in original}),
              'can_received': max(r['result'].get('can', {}).get('received', 0) for r in original),
              'can_0x391_by_channel': [max((r['result'].get('can', {}).get('channels', [{}, {}])[ch].get('received', 0) for r in original)) for ch in range(2)]}
    assert result['binary_sha256'] == running_hash
    (args.output/'report.json').write_text(json.dumps(result, indent=2)+'\n')
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()