#!/usr/bin/env python3
"""
Generate captions_overlay from script dialogue and video metadata.

Reads dialogue from 08-script-adapted/current/script.json
Reads actual video durations from 10-videos/scene_XX/current/metadata.json
Creates captions_overlay with proper timing for assembly.

Usage:
    ./venv/bin/python scripts/generate_captions_v3.py "truck-stop-revenge" --take 003
"""

import sys
import json
from pathlib import Path
from typing import List, Dict, Any

# Add lib to path
sys.path.insert(0, str(Path(__file__).parent))
from lib import (
    get_story_path,
    get_take_path,
    save_metadata,
    load_metadata
)


def split_dialogue_into_lines(text: str, max_chars_per_line: int = 50) -> List[str]:
    """
    Split dialogue into multiple caption lines.

    Args:
        text: Full dialogue text
        max_chars_per_line: Maximum characters per line

    Returns:
        List of caption lines
    """
    # If text fits in one line, return as-is
    if len(text) <= max_chars_per_line:
        return [text]

    # Split into sentences first
    sentences = []
    for s in text.replace("! ", "!|").replace(". ", ".|").replace("? ", "?|").split("|"):
        sentences.append(s.strip())

    lines = []
    current_line = ""

    for sentence in sentences:
        # If sentence alone is too long, split at natural break points
        if len(sentence) > max_chars_per_line:
            # Try to split at commas or conjunctions
            parts = sentence.replace(", ", ",|").split("|")
            for part in parts:
                if len(current_line) + len(part) + 1 <= max_chars_per_line:
                    current_line = (current_line + " " + part).strip()
                else:
                    if current_line:
                        lines.append(current_line)
                    current_line = part.strip()
        else:
            # Try to add sentence to current line
            if len(current_line) + len(sentence) + 1 <= max_chars_per_line:
                current_line = (current_line + " " + sentence).strip()
            else:
                if current_line:
                    lines.append(current_line)
                current_line = sentence

    # Add remaining text
    if current_line:
        lines.append(current_line)

    return lines


def estimate_caption_timing(
    scene_duration: float,
    dialogue_text: str,
    line_index: int,
    total_lines: int,
    has_action_before_dialogue: bool = True
) -> tuple:
    """
    Estimate when caption should appear within scene.

    Args:
        scene_duration: Actual scene duration in seconds
        dialogue_text: Text being spoken (this line)
        line_index: Index of this line (0-based)
        total_lines: Total number of lines
        has_action_before_dialogue: If True, delay caption start

    Returns:
        (start_time, end_time) in seconds from scene start
    """
    # Estimate speech duration for this line: ~150 words per minute = 2.5 words/sec
    word_count = len(dialogue_text.split())
    line_speech_duration = (word_count / 2.5) * 1.3  # 30% buffer

    # Calculate dialogue window (time available for all dialogue)
    if has_action_before_dialogue:
        start_offset = min(1.5, scene_duration * 0.25)
        dialogue_window = scene_duration - start_offset - 0.5
    else:
        start_offset = 0.2
        dialogue_window = scene_duration - 0.5

    # Split dialogue window evenly across lines
    time_per_line = dialogue_window / total_lines

    # Calculate this line's timing
    start_time = start_offset + (line_index * time_per_line)
    end_time = start_time + min(line_speech_duration, time_per_line - 0.2)

    # Ensure end_time doesn't exceed scene duration
    end_time = min(end_time, scene_duration - 0.2)

    return (round(start_time, 1), round(end_time, 1))


def should_delay_caption(scene_data: dict) -> bool:
    """
    Determine if caption should be delayed (action happens before speech).

    Returns True if scene has significant action before dialogue.
    """
    action = scene_data.get("action", "").lower()

    # Actions that happen before speaking
    action_first_keywords = [
        "turn", "look", "walk", "approach", "lean",
        "grab", "flip", "push", "enter", "exit",
        "see", "notice", "react", "stare"
    ]

    for keyword in action_first_keywords:
        if keyword in action:
            return True

    # If action explicitly mentions speaking first
    speaking_first_keywords = ["speak", "say", "shout", "deliver", "punchline"]
    for keyword in speaking_first_keywords:
        if keyword in action:
            return False

    # Default: assume brief action before speech
    return True


def generate_captions_overlay(story_name: str, take: str = "current") -> List[Dict[str, Any]]:
    """
    Generate captions_overlay from script dialogue and video metadata.

    Args:
        story_name: Story name
        take: Take number or "current"

    Returns:
        List of caption overlay definitions
    """
    # Get paths
    story_path = get_story_path(story_name)
    take_path = get_take_path(story_name, take)

    # Read adapted script
    script_path = take_path / "08-script-adapted" / "current" / "script.json"
    if not script_path.exists():
        print(f"Error: Adapted script not found: {script_path}")
        return []

    with open(script_path) as f:
        script = json.load(f)

    scenes = script.get("scenes", [])
    captions_overlay = []

    print(f"Generating captions for {story_name} (take: {take})")
    print(f"Found {len(scenes)} scenes\n")

    for scene in scenes:
        scene_num = scene["number"]
        dialogue = scene.get("dialogue")

        if not dialogue or not dialogue.get("text"):
            print(f"Scene {scene_num}: No dialogue - skipping")
            continue

        # Read video metadata to get actual duration
        video_metadata_path = take_path / "10-videos" / f"scene_{scene_num:02d}" / "current" / "metadata.json"

        if not video_metadata_path.exists():
            print(f"Scene {scene_num}: WARNING - Video metadata not found, using script duration")
            actual_duration = scene.get("duration", 5)
        else:
            video_metadata = load_metadata(video_metadata_path)
            actual_duration = video_metadata.get("duration_actual", scene.get("duration", 5))

        # Get dialogue text
        text = dialogue["text"]
        character = dialogue.get("character", "unknown")
        delivery = dialogue.get("delivery", "")

        # Split long dialogue into multiple lines
        lines = split_dialogue_into_lines(text, max_chars_per_line=45)

        # Determine if caption should be delayed
        delay_caption = should_delay_caption(scene)

        # Determine position (default: bottom, but can vary)
        position = "bottom"

        # Top position for certain scenes (like overhead shots, establishing shots)
        shot_type = scene.get("shot_type", "").lower()
        if "wide" in shot_type or "establishing" in shot_type:
            position = "top"

        print(f"Scene {scene_num}: \"{text}\"")
        print(f"  Duration: {actual_duration}s | Lines: {len(lines)} | Position: {position}")
        if delay_caption:
            print(f"  Delayed caption (action before speech)")

        # Create caption for each line
        for line_idx, line_text in enumerate(lines):
            # Calculate timing for this line
            start_time, end_time = estimate_caption_timing(
                actual_duration,
                line_text,
                line_index=line_idx,
                total_lines=len(lines),
                has_action_before_dialogue=delay_caption
            )

            caption = {
                "scene": scene_num,
                "text": line_text,
                "timing": f"{start_time}-{end_time}s",
                "position": position,
                "character": character,
                "delivery": delivery if line_idx == 0 else "",  # Only first line gets delivery info
                "actual_duration": actual_duration,
                "delay_applied": delay_caption,
                "line_number": line_idx + 1,
                "total_lines": len(lines)
            }

            captions_overlay.append(caption)

            print(f"  Line {line_idx + 1}/{len(lines)}: \"{line_text}\" ({start_time}-{end_time}s)")

        print()

    return captions_overlay


def save_captions_to_script(
    story_name: str,
    take: str,
    captions_overlay: List[Dict[str, Any]]
) -> bool:
    """
    Add captions_overlay to adapted script.json.

    Creates a new iteration of script-adapted with captions.
    """
    take_path = get_take_path(story_name, take)
    script_adapted_stage = take_path / "08-script-adapted"

    # Get current script
    current_script_path = script_adapted_stage / "current" / "script.json"
    with open(current_script_path) as f:
        script = json.load(f)

    # Add captions_overlay
    script["captions_overlay"] = captions_overlay

    # Get next iteration number
    existing_iterations = []
    for item in script_adapted_stage.iterdir():
        if item.is_dir() and item.name.isdigit() and len(item.name) == 3:
            existing_iterations.append(int(item.name))

    next_iteration = max(existing_iterations) + 1 if existing_iterations else 1
    iteration = f"{next_iteration:03d}"

    # Create new iteration
    from lib import ensure_dir, update_current_symlink
    iter_path = script_adapted_stage / iteration
    ensure_dir(iter_path)

    # Save script with captions
    output_path = iter_path / "script.json"
    with open(output_path, 'w') as f:
        json.dump(script, f, indent=2)

    # Save metadata
    metadata = {
        "stage": "08-script-adapted",
        "take": take,
        "iteration": iteration,
        "timestamp": __import__("datetime").datetime.utcnow().isoformat() + "Z",
        "source": f"../{'001' if next_iteration == 1 else f'{next_iteration-1:03d}'}/script.json",
        "action": "add_captions_overlay",
        "caption_count": len(captions_overlay)
    }

    metadata_path = iter_path / "metadata.json"
    save_metadata(metadata_path, metadata)

    # Update current symlink
    update_current_symlink(script_adapted_stage, iteration)

    print(f"\n✓ Saved script with captions: {output_path}")
    print(f"✓ Created iteration: {iteration}")
    print(f"✓ Updated current → {iteration}")

    return True


def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Generate captions overlay from script dialogue and video metadata"
    )
    parser.add_argument("story_name", help="Story name")
    parser.add_argument("--take", default="current", help="Take number (default: current)")
    parser.add_argument("--preview-only", action="store_true", help="Preview captions without saving")

    args = parser.parse_args()

    # Generate captions
    captions = generate_captions_overlay(args.story_name, args.take)

    if not captions:
        print("No captions generated (no dialogue in scenes)")
        return 1

    print("=" * 60)
    print(f"Generated {len(captions)} captions")
    print("=" * 60)

    if args.preview_only:
        print("\nPreview mode - not saving to script")
        print("\nCaptions overlay:")
        print(json.dumps(captions, indent=2))
        return 0

    # Save to script
    if save_captions_to_script(args.story_name, args.take, captions):
        print("\n✓ Captions added to adapted script")
        print("✓ Run assembly script to generate final video with captions")
        return 0
    else:
        print("\n✗ Failed to save captions")
        return 1


if __name__ == "__main__":
    sys.exit(main())
