#!/usr/bin/env python3
"""
Caption videos using Qwen2.5-VL + Whisper for LTX-2 training
"""

import os
import torch
import json
import subprocess
import tempfile
from pathlib import Path
from tqdm import tqdm

# Qwen2.5-VL imports
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

# Whisper imports
from transformers import WhisperProcessor, WhisperForConditionalGeneration
import librosa

# === CONFIGURATION ===
VIDEO_DIR = "/path/to/your/videos"
OUTPUT_JSONL = "/path/to/output/dataset.jsonl"
TRIGGER_TOKEN = "your_trigger"  # e.g., "mystyle", "mychar"

QWEN_MODEL = "huihui-ai/Qwen2.5-VL-7B-Instruct-abliterated"
WHISPER_MODEL = "openai/whisper-large-v3"

CAPTION_PROMPT = """Describe this video in detail for AI training.
Focus on: actions, movements, camera angles, lighting, subjects.
Be descriptive and specific. Just provide the description."""


def extract_audio(video_path: Path, output_path: Path) -> bool:
    """Extract audio from video using ffmpeg"""
    cmd = [
        "ffmpeg", "-y", "-i", str(video_path),
        "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1",
        str(output_path)
    ]
    try:
        result = subprocess.run(cmd, capture_output=True, timeout=60)
        return result.returncode == 0 and output_path.exists()
    except Exception as e:
        print(f"  Audio extraction failed: {e}")
        return False


def transcribe_audio(whisper_model, whisper_processor, audio_path: Path, device: str) -> str:
    """Transcribe audio using Whisper"""
    try:
        audio, sr = librosa.load(str(audio_path), sr=16000)
        inputs = whisper_processor(audio, sampling_rate=16000, return_tensors="pt")
        inputs = {k: v.to(device=device, dtype=torch.float16) for k, v in inputs.items()}

        with torch.inference_mode():
            generated_ids = whisper_model.generate(
                inputs["input_features"],
                max_length=448,
                language="en",
                task="transcribe"
            )

        transcription = whisper_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
        return transcription.strip()
    except Exception as e:
        print(f"  Transcription failed: {e}")
        return ""


def caption_video(qwen_model, qwen_processor, video_path: Path, device: str) -> str:
    """Caption a video using Qwen2.5-VL"""
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "video",
                    "video": str(video_path),
                    "max_pixels": 360 * 420,
                    "nframes": 16,
                },
                {"type": "text", "text": CAPTION_PROMPT},
            ],
        }
    ]

    text = qwen_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs, video_kwargs = process_vision_info(messages, return_video_kwargs=True)

    inputs = qwen_processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
        **video_kwargs,
    )
    inputs = inputs.to(device)

    with torch.inference_mode():
        generated_ids = qwen_model.generate(
            **inputs,
            max_new_tokens=300,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
        )

    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = qwen_processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )

    return output_text[0].strip()


def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"

    video_dir = Path(VIDEO_DIR)
    video_files = sorted(list(video_dir.glob("**/*.mp4")))
    print(f"Found {len(video_files)} videos")

    # Load models
    print("Loading Whisper...")
    whisper_processor = WhisperProcessor.from_pretrained(WHISPER_MODEL)
    whisper_model = WhisperForConditionalGeneration.from_pretrained(
        WHISPER_MODEL, torch_dtype=torch.float16, device_map="auto"
    )

    print("Loading Qwen2.5-VL...")
    qwen_processor = AutoProcessor.from_pretrained(QWEN_MODEL)
    qwen_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        QWEN_MODEL, torch_dtype=torch.bfloat16, device_map="auto",
        attn_implementation="flash_attention_2",
    )

    results = []
    for video_path in tqdm(video_files, desc="Captioning"):
        # Transcribe audio
        speech_text = ""
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as tmp:
            tmp_path = Path(tmp.name)
            if extract_audio(video_path, tmp_path):
                speech_text = transcribe_audio(whisper_model, whisper_processor, tmp_path, device)

        # Caption video
        try:
            visual_caption = caption_video(qwen_model, qwen_processor, video_path, device)
        except Exception as e:
            print(f"  Caption failed: {e}")
            visual_caption = "video content"

        # Build LTX-2 format caption
        speech_line = speech_text if speech_text else "None"
        full_caption = f"[VISUAL]: {TRIGGER_TOKEN}. {visual_caption} [SPEECH]: {speech_line} [SOUNDS]: ambient sounds [TEXT]: None"

        results.append({
            "prompt": full_caption,
            "media_path": str(video_path.relative_to(video_dir))
        })

    # Save JSONL
    output_path = Path(OUTPUT_JSONL)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        for result in results:
            f.write(json.dumps(result) + "\n")

    print(f"Saved {len(results)} captions to {output_path}")


if __name__ == "__main__":
    main()