#!/usr/bin/env python3
"""Look at and listen to a take: a contact sheet of N frames, plus a transcript and a description of
the soundtrack (an audio model hears it so the cut does not rest on guesswork).

    uv run --with requests python tools/trailer/qa.py <clip.mp4> [frames=6]
    uv run --with requests python tools/trailer/qa.py --ask <media> "<question>"   # any file, any question
"""
import base64, json, os, subprocess, sys, tempfile

import requests

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from gen import API, H, OUT

def ask(media, question, model=os.environ.get('QA_LISTENER', 'google/gemini-2.5-pro')):
    wav = tempfile.mktemp(suffix='.wav')
    subprocess.run(['ffmpeg', '-v', 'error', '-y', '-i', media, '-ac', '1', '-ar', '16000', wav], check=True)
    r = requests.post(f'{API}/chat/completions', headers=H(), json={'model': model, 'messages': [{'role': 'user', 'content': [
        {'type': 'text', 'text': question},
        {'type': 'input_audio', 'input_audio': {'data': base64.b64encode(open(wav, 'rb').read()).decode(), 'format': 'wav'}}]}]},
        timeout=600)
    return r.json()['choices'][0]['message']['content'] if r.ok else r.text[:500]


if sys.argv[1] == '--ask':
    sys.exit(print(ask(sys.argv[2], sys.argv[3])))
clip = sys.argv[1]
n = int(sys.argv[2]) if len(sys.argv) > 2 else 6
name = os.path.splitext(os.path.basename(clip))[0]
dur = float(subprocess.check_output(['ffprobe', '-v', 'error', '-show_entries', 'format=duration',
                                     '-of', 'csv=p=0', clip]).decode())
info = subprocess.check_output(['ffprobe', '-v', 'error', '-show_entries', 'stream=codec_type,width,height,r_frame_rate',
                                '-of', 'compact', clip]).decode().strip()
print(f'{name}: {dur:.2f}s  {info}')
sheet = os.path.join(OUT, 'qa', name + '.jpg')
os.makedirs(os.path.dirname(sheet), exist_ok=True)
tmp = tempfile.mkdtemp()
for i in range(n):
    t = dur * (i + 0.5) / n
    subprocess.run(['ffmpeg', '-v', 'error', '-y', '-ss', f'{t:.2f}', '-i', clip, '-frames:v', '1',
                    '-vf', 'scale=-2:640', f'{tmp}/{i:02d}.jpg'], check=True)
subprocess.run(['magick', *sorted(f'{tmp}/{f}' for f in os.listdir(tmp)), '+append', sheet], check=True)
print('sheet', sheet)
if 'audio' not in info:
    sys.exit('no audio stream')
wav = f'{tmp}/a.wav'
subprocess.run(['ffmpeg', '-v', 'error', '-y', '-i', clip, '-ac', '1', '-ar', '16000', wav], check=True)
r = requests.post(f'{API}/chat/completions', headers=H(), json={
    'model': os.environ.get('QA_MODEL', 'google/gemini-2.5-flash'),
    'messages': [{'role': 'user', 'content': [
        {'type': 'text', 'text': 'Transcribe every spoken word in this clip verbatim with rough timestamps. '
         'Then, in two short lines: describe the voice (accent, age, texture) and list every non-speech '
         'sound you hear (music? ambience? effects?).'},
        {'type': 'input_audio', 'input_audio': {'data': base64.b64encode(open(wav, 'rb').read()).decode(),
                                                'format': 'wav'}}]}]}, timeout=300)
print(r.json()['choices'][0]['message']['content'] if r.ok else r.text[:500])