#!/usr/bin/env python3
import json
import time
import urllib.request
from pathlib import Path

PROJECT = Path('/data/video-pipeline/Youtube-Video-Maker/Project/08062026-003-Le-Hoan')
ENDPOINT = 'http://192.168.1.104:8002/probe'
MASTER_AUDIO = PROJECT / 'audio/master_narration.wav'
TIMELINE = PROJECT / 'audio/master_narration_timeline.json'
OUTPUT = PROJECT / 'audio/master_audio_probe.json'
MANIFEST = PROJECT / 'logs/step9_probe_audio_manifest.json'
PART_GLOBS = [
    PROJECT / 'audio/render_parts/part_*.wav',
    PROJECT / 'audio/parts/part_*.wav',
    PROJECT / 'audio/part_*.wav',
]


def log(msg):
    print(time.strftime('%Y-%m-%d %H:%M:%S ') + msg, flush=True)


def probe(path):
    req = urllib.request.Request(
        ENDPOINT,
        data=json.dumps({'input_path': str(path)}).encode('utf-8'),
        headers={'Content-Type': 'application/json'},
        method='POST',
    )
    with urllib.request.urlopen(req, timeout=300) as resp:
        raw = json.loads(resp.read().decode('utf-8'))
    fmt = raw.get('format', {})
    audio_stream = next((s for s in raw.get('streams', []) if s.get('codec_type') == 'audio'), {})
    return {
        'audio_path': str(path),
        'duration_sec': float(fmt.get('duration') or audio_stream.get('duration') or 0),
        'sample_rate': int(audio_stream.get('sample_rate') or 0),
        'channels': audio_stream.get('channels'),
        'codec_name': audio_stream.get('codec_name'),
        'format_name': fmt.get('format_name'),
        'bit_rate': int(fmt.get('bit_rate') or audio_stream.get('bit_rate') or 0),
        'size_bytes': int(fmt.get('size') or 0),
        'raw_probe': raw,
    }


def part_id(path):
    return path.stem


def main():
    if not MASTER_AUDIO.exists():
        raise SystemExit(f'Missing master audio: {MASTER_AUDIO}')

    timeline_expected = None
    timeline_path = None
    if TIMELINE.exists():
        timeline_doc = json.loads(TIMELINE.read_text(encoding='utf-8'))
        timeline_expected = timeline_doc.get('total_duration')
        timeline_path = str(TIMELINE)

    log(f'PROBE master {MASTER_AUDIO}')
    master = probe(MASTER_AUDIO)

    parts_paths = []
    for pattern in PART_GLOBS:
        parts_paths.extend(sorted(pattern.parent.glob(pattern.name)))
    seen = set()
    unique_parts = []
    for p in parts_paths:
        if p not in seen:
            seen.add(p)
            unique_parts.append(p)

    parts = []
    for p in unique_parts:
        log(f'PROBE part {p}')
        info = probe(p)
        info['part_id'] = part_id(p)
        parts.append(info)

    notes = []
    diff = None
    if timeline_expected is not None:
        diff = master['duration_sec'] - float(timeline_expected)
        if abs(diff) > 0.01:
            notes.append(f'Master probe duration differs from Step 8 timeline by {diff:.3f}s; use probe duration as source of truth.')
        else:
            notes.append('Master probe duration matches Step 8 timeline within 0.01s.')
    if not parts:
        notes.append('No part_XX.wav files found yet; only master audio was probed.')

    result = {
        'project_id': PROJECT.name,
        'step': 9,
        'endpoint': ENDPOINT,
        'source_of_truth': 'worker_probe_duration',
        'master_audio_path': str(MASTER_AUDIO),
        'master_duration_sec': master['duration_sec'],
        'master_probe': master,
        'timeline_path': timeline_path,
        'timeline_expected_duration_sec': timeline_expected,
        'timeline_vs_probe_diff_sec': diff,
        'parts': parts,
        'notes': notes,
    }
    OUTPUT.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding='utf-8')

    manifest = {
        'project_id': PROJECT.name,
        'step': 9,
        'status': 'completed',
        'endpoint': ENDPOINT,
        'output_probe_json': str(OUTPUT),
        'master_audio_path': str(MASTER_AUDIO),
        'master_duration_sec': master['duration_sec'],
        'parts_count': len(parts),
        'notes': notes,
    }
    MANIFEST.parent.mkdir(parents=True, exist_ok=True)
    MANIFEST.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding='utf-8')
    log('DONE ' + json.dumps({'status': 'completed', 'master_duration_sec': master['duration_sec'], 'parts': len(parts)}, ensure_ascii=False))


if __name__ == '__main__':
    main()