"""Local recording/capture intake and encoded-file checks. No network requests."""

import argparse
from array import array
from datetime import datetime, timezone
from fractions import Fraction
import hashlib
import json
import math
from pathlib import Path
import shutil
import subprocess
import sys
import wave


def require(condition, message):
    if not condition:
        raise ValueError(message)


def sha(path):
    with Path(path).open('rb') as source:
        digest = hashlib.file_digest(source, 'sha256')
    return digest.hexdigest()


def write(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, allow_nan=False) + '\n')


def run(command):
    return subprocess.check_output(command, stderr=subprocess.PIPE)


def probe(path):
    return json.loads(run(['ffprobe', '-v', 'error', '-count_frames', '-show_streams', '-show_format', '-of', 'json', str(path)]))


def versions():
    return {'python': sys.version, 'ffmpeg': run(['ffmpeg', '-version']).decode().splitlines()[0],
            'ffprobe': run(['ffprobe', '-version']).decode().splitlines()[0]}


def timestamps(path):
    return [float(frame['best_effort_timestamp_time']) for frame in json.loads(run([
        'ffprobe', '-v', 'error', '-select_streams', 'v:0', '-show_frames',
        '-show_entries', 'frame=best_effort_timestamp_time', '-of', 'json', str(path)]))['frames']]


def inspect_recording(source, output, audio_stream=None, threshold_dbfs=-35):
    source, output = Path(source), Path(output)
    output.mkdir(parents=True, exist_ok=True)
    original_hash = sha(source)
    metadata = probe(source)
    origin = float(metadata['format'].get('start_time', 0))
    tracks = []
    selected_pcm = None
    for stream in metadata['streams']:
        if stream['codec_type'] != 'audio':
            continue
        pcm = run(['ffmpeg', '-v', 'error', '-i', str(source), '-map', f"0:{stream['index']}",
                   '-ac', '1', '-ar', '16000', '-f', 's16le', '-'])
        samples = array('h', pcm)
        if sys.byteorder != 'little':
            samples.byteswap()
        peak = max((abs(sample) for sample in samples), default=0)
        digest = hashlib.sha256(pcm).hexdigest()
        tracks.append({'index': stream['index'], 'pcm_sha256': digest, 'sample_count': len(samples),
                       'peak_dbfs': 20 * math.log10(max(peak, 1) / 32768),
                       'source_offset_seconds': float(stream.get('start_time', origin)) - origin})
        if audio_stream == stream['index']:
            selected_pcm = pcm
    require(tracks, 'No audio track. Supply the authoritative take; do not infer speech from a script.')
    groups = {}
    for track in tracks:
        groups.setdefault(track['pcm_sha256'], []).append(track['index'])
    if audio_stream is None:
        require(len(groups) == 1, 'Distinct audio tracks: inspect them and select --audio-stream explicitly.')
        audio_stream = tracks[0]['index']
        selected_pcm = run(['ffmpeg', '-v', 'error', '-i', str(source), '-map', f'0:{audio_stream}', '-ac', '1', '-ar', '16000', '-f', 's16le', '-'])
    require(selected_pcm is not None, 'Selected audio stream does not exist')
    selected = next(track for track in tracks if track['index'] == audio_stream)
    require(selected['peak_dbfs'] > -65, 'Selected track is silent/near-silent; review track selection instead of making a transcript.')
    with wave.open(str(output / 'analysis.wav'), 'wb') as audio:
        audio.setnchannels(1)
        audio.setsampwidth(2)
        audio.setframerate(16000)
        audio.writeframes(selected_pcm)
    samples = array('h', selected_pcm)
    if sys.byteorder != 'little':
        samples.byteswap()
    envelope = []
    for start in range(0, len(samples), 320):
        window = samples[start:start + 320]
        rms = math.sqrt(sum(sample * sample for sample in window) / len(window))
        envelope.append({'source_seconds': start / 16000 + selected['source_offset_seconds'],
                         'rms_dbfs': 20 * math.log10(max(rms, 1) / 32768)})
    report = {'schema_version': 1, 'source_sha256': original_hash, 'source_probe': metadata,
              'time_origin': 'seconds relative to format start_time; original untrimmed recording',
              'source_format_start_seconds': origin, 'tracks': tracks, 'duplicate_pcm_groups': list(groups.values()),
              'selected_audio_stream': audio_stream, 'audio_to_source_offset_seconds': selected['source_offset_seconds'],
              'analysis_wav_sha256': sha(output / 'analysis.wav'), 'envelope_window_seconds': .02,
              'threshold_dbfs': threshold_dbfs, 'threshold_policy': 'Triage threshold only; listen and inspect the video. Noise is not speech or proof of silence.',
              'envelope': envelope, 'versions': versions(), 'visual_tail_review': 'required, not inferred from audio',
              'speaker_count': 'not inferred from track count'}
    require(sha(source) == original_hash, 'Source changed during intake')
    write(output / 'intake.json', report)
    return report


def import_asr(intake_path, asr_path, output, engine, revision, proper_nouns):
    intake = json.loads(Path(intake_path).read_text())
    raw = json.loads(Path(asr_path).read_text())
    offset = intake['audio_to_source_offset_seconds']
    duration = float(intake['source_probe']['format']['duration'])
    envelope = intake['envelope']
    words, uncertainties = [], []
    for segment in raw['segments']:
        require(segment.get('words'), 'Word estimates are required; segment lengths are not word timing.')
        for word in segment['words']:
            start, end = float(word['start']) + offset, float(word['end']) + offset
            require(math.isfinite(start) and math.isfinite(end) and 0 <= start < end <= duration, 'ASR word outside original source time')
            text = word.get('word', word.get('text', '')).strip()
            identifier = f'word-{len(words):04d}'
            words.append({'id': identifier, 'text': text, 'start_seconds': start, 'end_seconds': end,
                          'raw_asr_start_seconds': word['start'], 'raw_asr_end_seconds': word['end'],
                          'evidence_kind': 'asr-estimate', 'probability': word.get('probability')})
            onset = [sample for sample in envelope if start <= sample['source_seconds'] < min(end, start + .12)]
            if not onset or max(sample['rms_dbfs'] for sample in onset) < intake['threshold_dbfs']:
                active = next((sample['source_seconds'] for sample in envelope if start <= sample['source_seconds'] < end and sample['rms_dbfs'] >= intake['threshold_dbfs']), None)
                uncertainties.append({'id': f'{identifier}-onset', 'word_id': identifier, 'kind': 'timestamp-over-low-energy',
                                      'raw_start_seconds': start, 'candidate_acoustic_onset_seconds': active,
                                      'status': 'unresolved', 'decision': 'Listen/inspect. Do not replace raw timing with an energy threshold automatically.'})
            if text.lower().strip('.,:!?') in proper_nouns or word.get('probability', 1) < .7:
                uncertainties.append({'id': f'{identifier}-word', 'word_id': identifier, 'kind': 'spelling-or-low-confidence',
                                      'status': 'unresolved', 'decision': 'Confirm against the recording and authorized context during consolidated review.'})
    require(words, 'Empty ASR transcript')
    report = {'schema_version': 1, 'source_sha256': intake['source_sha256'], 'intake_sha256': sha(intake_path),
              'raw_asr_sha256': sha(asr_path), 'engine': engine, 'model_revision': revision,
              'time_origin': intake['time_origin'], 'words': words, 'uncertainties': uncertainties,
              'emphasis_policy': 'Semantic emphasis is an editorial judgment, not measured vocal stress.',
              'review_decisions': [], 'human_verified': False}
    write(output, report)
    return report


def import_capture(source, output, useful_region, capture_kind):
    source, output = Path(source), Path(output)
    require(not output.exists(), 'Capture output already exists; use a new directory to preserve provenance.')
    metadata = probe(source)
    video = next(stream for stream in metadata['streams'] if stream['codec_type'] == 'video')
    points = timestamps(source)
    require(len(points) >= 2 and all(after > before for before, after in zip(points, points[1:])), 'Capture needs increasing actual encoded PTS')
    left, top, width, height = useful_region
    require(min(left, top) >= 0 and min(width, height) > 0 and left + width <= video['width'] and top + height <= video['height'], 'Useful captured region exceeds encoded pixels')
    output.mkdir(parents=True)
    original = output / ('source' + source.suffix)
    shutil.copy2(source, original)
    (output / 'frames').mkdir()
    run(['ffmpeg', '-v', 'error', '-xerror', '-i', str(original), '-map', '0:v:0', '-fps_mode', 'passthrough', '-start_number', '0', str(output / 'frames/%06d.png')])
    frames = sorted((output / 'frames').glob('*.png'))
    require(len(frames) == len(points), 'Capture frame/PTS count mismatch')
    intervals = [after - before for before, after in zip(points, points[1:])]
    report = {'schema_version': 1, 'kind': capture_kind, 'source_file': original.name, 'source_sha256': sha(original),
              'source_probe': metadata, 'useful_region': useful_region, 'region_evidence': 'Operator must inspect full moving capture and confirm useful pixels; container size alone is insufficient.',
              'encoded_start_pts_seconds': points[0], 'duration_seconds': points[-1] - points[0] + sorted(intervals)[len(intervals) // 2],
              'maximum_frame_interval_seconds': max(intervals), 'median_frame_interval_seconds': sorted(intervals)[len(intervals) // 2],
              'frames': [{'path': str(frame.relative_to(output)), 'source_seconds': moment - points[0], 'sha256': sha(frame)} for frame, moment in zip(frames, points)],
              'versions': versions(), 'audio_included_in_presentation': False}
    write(output / 'capture.json', report)
    return report


def validate_export(movie, width, height, fps, count):
    metadata = probe(movie)
    require(len(metadata['streams']) == 1, 'Presentation must contain exactly one silent video stream')
    video = metadata['streams'][0]
    require(video['codec_type'] == 'video' and video['codec_name'] == 'h264' and video['pix_fmt'] == 'yuv420p', 'Expected H.264/yuv420p video')
    require((video['width'], video['height']) == (width, height), 'Wrong encoded dimensions')
    require(Fraction(video['avg_frame_rate']) == fps and Fraction(video['r_frame_rate']) == fps, 'Wrong encoded frame rate')
    require(int(video['nb_read_frames']) == count, 'Wrong decoded frame count')
    require(abs(float(video['duration']) - count / fps) < 1 / fps, 'Wrong encoded duration')
    require(all(video.get(field) == value for field, value in {'color_space': 'bt709', 'color_range': 'tv', 'color_primaries': 'bt709', 'color_transfer': 'bt709'}.items()), 'Missing/wrong actual encoded BT.709 metadata')
    points = timestamps(movie)
    require(len(points) == count and all(abs(moment - index / fps) < .00001 for index, moment in enumerate(points)), 'Encoded PTS do not match the edit grid')
    run(['ffmpeg', '-v', 'error', '-xerror', '-i', str(movie), '-map', '0:v:0', '-f', 'null', '-'])
    return {'sha256': sha(movie), 'actual_probe': metadata, 'all_frame_pts_checked': len(points), 'full_decode': 'passed', 'meeting_compatibility': 'untested'}


def compare_frames(movie, frames, width, height, fps, count):
    from PIL import Image, ImageChops, ImageStat
    expected_paths = sorted(Path(frames).glob('*.png'))
    require(len(expected_paths) == count, 'Missing expected render frames')
    process = subprocess.Popen(['ffmpeg', '-v', 'error', '-xerror', '-i', str(movie), '-map', '0:v:0', '-fps_mode', 'passthrough', '-pix_fmt', 'rgb24', '-f', 'rawvideo', '-'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    maximum_mean = maximum_tile = 0
    try:
        for index, expected_path in enumerate(expected_paths):
            data = process.stdout.read(width * height * 3)
            require(len(data) == width * height * 3, f'Cannot decode frame {index}')
            actual = Image.frombytes('RGB', (width, height), data)
            with Image.open(expected_path) as source:
                expected = source.convert('RGB')
            require(expected.size == actual.size, 'Expected frame size mismatch')
            difference = ImageChops.difference(expected, actual)
            mean = sum(ImageStat.Stat(difference).mean) / 3
            tile = max(sum(ImageStat.Stat(difference.crop((left, top, min(width, left + 64), min(height, top + 64)))).mean) / 3 for top in range(0, height, 64) for left in range(0, width, 64))
            maximum_mean, maximum_tile = max(maximum_mean, mean), max(maximum_tile, tile)
            require(mean < 4 and tile < 18, f'Encoded frame {index} diverges: mean={mean:.3f}; 64px tile={tile:.3f}')
        require(not process.stdout.read(1) and process.wait() == 0, 'Extra frames or decode failure')
    finally:
        if process.poll() is None:
            process.kill()
        process.wait()
        process.stdout.close()
        process.stderr.close()
    return {'frames_compared': count, 'maximum_mean_rgb_error': maximum_mean, 'maximum_tile_rgb_error': maximum_tile,
            'tolerance': {'mean_rgb_exclusive': 4, 'tile_64px_mean_exclusive': 18},
            'expected_frames_sha256': hashlib.sha256(''.join(sha(path) for path in expected_paths).encode()).hexdigest()}


def encode(frames, output, fps):
    frames, output = Path(frames), Path(output)
    expected = sorted(frames.glob('*.png'))
    require(expected and not output.exists(), 'Need expected frames and a new output path')
    output.parent.mkdir(parents=True, exist_ok=True)
    command = ['ffmpeg', '-v', 'error', '-f', 'image2pipe', '-framerate', str(fps), '-i', '-', '-an',
               '-vf', 'scale=in_range=full:out_range=tv:out_color_matrix=bt709', '-c:v', 'libx264', '-preset', 'medium', '-crf', '18',
               '-pix_fmt', 'yuv420p', '-color_primaries', 'bt709', '-color_trc', 'bt709', '-colorspace', 'bt709', '-color_range', 'tv',
               '-bsf:v', 'h264_metadata=colour_primaries=1:transfer_characteristics=1:matrix_coefficients=1:video_full_range_flag=0',
               '-movflags', '+faststart', str(output)]
    process = subprocess.Popen(command, stdin=subprocess.PIPE)
    try:
        for frame in expected:
            process.stdin.write(frame.read_bytes())
    finally:
        process.stdin.close()
    require(process.wait() == 0, 'Encoding failed')
    return {'command': command, 'output_sha256': sha(output), 'versions': versions()}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--report', type=Path, required=True)
    commands = parser.add_subparsers(dest='command', required=True)
    inspect = commands.add_parser('inspect')
    inspect.add_argument('--source', type=Path, required=True)
    inspect.add_argument('--out', type=Path, required=True)
    inspect.add_argument('--audio-stream', type=int)
    inspect.add_argument('--threshold-dbfs', type=float, default=-35)
    asr = commands.add_parser('import-asr')
    for name in ['intake', 'asr', 'out', 'engine', 'revision']:
        asr.add_argument('--' + name, required=True)
    asr.add_argument('--proper-nouns', nargs='*', default=[])
    capture = commands.add_parser('capture')
    capture.add_argument('--source', type=Path, required=True)
    capture.add_argument('--out', type=Path, required=True)
    capture.add_argument('--useful-region', type=int, nargs=4, required=True)
    capture.add_argument('--kind', choices=['native-recording', 'sample-and-hold'], required=True)
    encoder = commands.add_parser('encode')
    encoder.add_argument('--frames', type=Path, required=True)
    encoder.add_argument('--out', type=Path, required=True)
    encoder.add_argument('--fps', type=int, required=True)
    verify = commands.add_parser('verify')
    verify.add_argument('--movie', type=Path, required=True)
    verify.add_argument('--frames', type=Path, required=True)
    for field in ['width', 'height', 'fps', 'count']:
        verify.add_argument('--' + field, type=int, required=True)
    arguments = parser.parse_args()
    report = {'status': 'running', 'started_at': datetime.now(timezone.utc).isoformat(), 'invocation': sys.argv,
              'tool_sha256': sha(__file__)}
    write(arguments.report, report)
    try:
        if arguments.command == 'inspect':
            result = inspect_recording(arguments.source, arguments.out, arguments.audio_stream, arguments.threshold_dbfs)
        elif arguments.command == 'import-asr':
            result = import_asr(arguments.intake, arguments.asr, arguments.out, arguments.engine, arguments.revision, [name.lower() for name in arguments.proper_nouns])
        elif arguments.command == 'capture':
            result = import_capture(arguments.source, arguments.out, arguments.useful_region, arguments.kind)
        elif arguments.command == 'encode':
            result = encode(arguments.frames, arguments.out, arguments.fps)
        else:
            report['movie_sha256'] = sha(arguments.movie)
            result = validate_export(arguments.movie, arguments.width, arguments.height, arguments.fps, arguments.count)
            result['decoded_pixel_checks'] = compare_frames(arguments.movie, arguments.frames, arguments.width, arguments.height, arguments.fps, arguments.count)
        report.update(status='passed', result=result, exit_status=0)
    except Exception as error:
        report.update(status='failed', error=str(error), exit_status=1)
        raise
    finally:
        report['finished_at'] = datetime.now(timezone.utc).isoformat()
        write(arguments.report, report)


if __name__ == '__main__':
    main()
