#!/usr/bin/env python3
"""
Stage 13: Video Assembly

Assembles final video from scene clips with timed captions.

- Reads selected videos from pipeline.json (10-videos)
- Reads selected captions from pipeline.json (12-captions)
- Normalizes clips to consistent format
- Adds timed captions using FFmpeg drawtext
- Concatenates all clips into final video
- Saves to 13-assembly/{iteration}/final.mp4

Commands:
- assemble: Assemble final video
- status_assembly: Show assembly status
"""

import sys
import json
import subprocess
from pathlib import Path
from datetime import datetime, timezone
from typing import Optional, Dict, List, Any

# Add parent to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from lib import (
    save_metadata,
    ensure_dir,
    setup_generation_log,
)
from lib.path_utils_v3 import get_next_iteration
from lib.pipeline import load_pipeline, save_pipeline


def get_video_duration(path: Path) -> float:
    """Get video duration in seconds using ffprobe."""
    cmd = [
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        str(path)
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode == 0:
        try:
            return float(result.stdout.strip())
        except ValueError:
            return 0.0
    return 0.0


def get_video_resolution(path: Path) -> tuple:
    """Get video resolution (width, height) using ffprobe."""
    cmd = [
        "ffprobe", "-v", "error",
        "-select_streams", "v:0",
        "-show_entries", "stream=width,height",
        "-of", "csv=s=x:p=0",
        str(path)
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode == 0:
        try:
            w, h = result.stdout.strip().split('x')
            return (int(w), int(h))
        except:
            return (0, 0)
    return (0, 0)


def create_streaming_variants(input_path: Path, logger=None) -> dict:
    """
    Create streaming-optimized variants of the final video.

    Creates:
    - final_web.mp4: Web-optimized with faststart (medium quality)
    - final_telegram.mp4: Smaller file for Telegram/mobile (lower bitrate)

    Returns:
        Dict with variant paths and sizes
    """
    variants = {}
    output_dir = input_path.parent

    # Web-optimized version (faststart, good quality)
    web_path = output_dir / "final_web.mp4"
    if logger:
        logger.info("Creating web-optimized version (faststart)...")

    web_cmd = [
        "ffmpeg", "-y", "-i", str(input_path),
        "-c:v", "libx264", "-preset", "medium", "-crf", "23",
        "-c:a", "aac", "-b:a", "128k",
        "-profile:v", "high", "-level", "4.1",
        "-movflags", "+faststart",
        "-pix_fmt", "yuv420p",
        str(web_path)
    ]

    result = subprocess.run(web_cmd, capture_output=True, text=True, timeout=600)
    if result.returncode == 0 and web_path.exists():
        size_mb = web_path.stat().st_size / 1024 / 1024
        variants["web"] = {"path": str(web_path), "size_mb": round(size_mb, 2)}
        if logger:
            logger.info(f"  final_web.mp4: {size_mb:.1f}MB")
    else:
        if logger:
            error_detail = result.stderr or result.stdout
            logger.warning(f"Failed to create web variant: {error_detail[:200] if error_detail else 'no output'}")

    # Telegram-optimized version (smaller, lower bitrate)
    telegram_path = output_dir / "final_telegram.mp4"
    if logger:
        logger.info("Creating Telegram-optimized version (smaller)...")

    telegram_cmd = [
        "ffmpeg", "-y", "-i", str(input_path),
        "-c:v", "libx264", "-preset", "slow", "-crf", "26",
        "-c:a", "aac", "-b:a", "96k",
        "-movflags", "+faststart",
        "-pix_fmt", "yuv420p",
        str(telegram_path)
    ]

    result = subprocess.run(telegram_cmd, capture_output=True, text=True, timeout=600)
    if result.returncode == 0 and telegram_path.exists():
        size_mb = telegram_path.stat().st_size / 1024 / 1024
        variants["telegram"] = {"path": str(telegram_path), "size_mb": round(size_mb, 2)}
        if logger:
            logger.info(f"  final_telegram.mp4: {size_mb:.1f}MB")
    else:
        if logger:
            error_detail = result.stderr or result.stdout
            logger.warning(f"Failed to create telegram variant: {error_detail[:200] if error_detail else 'no output'}")

    return variants


def build_drawtext_filter(
    text: str,
    start_time: float,
    end_time: float,
    font_size: int = 42,
    font_color: str = "white",
    bg_color: str = "black@0.7",
    position: str = "bottom",
) -> str:
    """Build a single drawtext filter string."""
    # Escape special characters for FFmpeg drawtext filter
    escaped_text = (
        text.replace("\\", "\\\\")
        .replace("'", "\u2019")  # Apostrophe → right single quotation mark
        .replace(":", "\\:")
        .replace("%", "\\%")
    )

    # Position mapping
    if position == "bottom":
        y_pos = "h-th-80"
    elif position == "top":
        y_pos = "80"
    else:
        y_pos = "(h-th)/2"

    return (
        f"drawtext=text='{escaped_text}':"
        f"fontsize={font_size}:"
        f"fontcolor={font_color}:"
        f"x=(w-tw)/2:"
        f"y={y_pos}:"
        f"box=1:"
        f"boxcolor={bg_color}:"
        f"boxborderw=15:"
        f"fontfile=/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf:"
        f"enable='between(t,{start_time},{end_time})'"
    )


def add_all_captions(
    input_path: Path,
    output_path: Path,
    captions: List[Dict],
    font_size: int = 42,
    font_color: str = "white",
    bg_color: str = "black@0.7",
    target_width: int = 1072,
    target_height: int = 1920,
    target_sample_rate: int = 44100,
    scene_num: int = None,
    trim_duration: float = None,
    trim_start: float = None,
) -> bool:
    """Add all captions to video in a single FFmpeg pass.

    Also normalizes resolution and sample rate for consistent concat.
    Uses scale-to-fill + center crop to avoid black bars.

    Args:
        scene_num: If set, overlay scene number badge in top-right corner (debug mode)
        trim_duration: If set, trim clip to this many seconds
        trim_start: If set, skip this many seconds from the beginning
    """
    if not captions:
        return False

    # Build combined filter with all drawtext filters chained
    filters = []

    # First: scale to fill + center crop (no black bars)
    scale_crop_filter = (
        f"scale={target_width}:{target_height}:"
        f"force_original_aspect_ratio=increase,"
        f"crop={target_width}:{target_height}"
    )
    filters.append(scale_crop_filter)

    # Add scene number overlay for review mode
    if scene_num is not None:
        filters.append(_build_scene_number_filter(scene_num))

    # Then: add caption overlays
    for ct in captions:
        text = ct.get("text", "")
        start = ct.get("start", 0.5)
        end = ct.get("end", 4.5)
        position = ct.get("position", "bottom")

        if text.strip():
            filters.append(build_drawtext_filter(
                text, start, end,
                font_size=font_size,
                font_color=font_color,
                bg_color=bg_color,
                position=position
            ))

    if len(filters) <= 1:  # Only scale filter, no captions
        return False

    # Chain filters with commas
    filter_str = ",".join(filters)

    cmd = ["ffmpeg", "-y"]

    if trim_start is not None:
        cmd.extend(["-ss", str(trim_start)])

    cmd.extend([
        "-i", str(input_path),
        "-vf", filter_str,
        "-c:v", "libx264", "-preset", "fast", "-crf", "18",
        "-c:a", "aac", "-b:a", "192k", "-ar", str(target_sample_rate),
        "-pix_fmt", "yuv420p",
    ])

    if trim_duration is not None:
        cmd.extend(["-t", str(trim_duration)])

    cmd.append(str(output_path))

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        if result.returncode != 0:
            return False
        return True
    except Exception as e:
        return False


def _build_scene_number_filter(scene_num: int) -> str:
    """Build FFmpeg drawtext filter for scene number badge in top-right corner."""
    return (
        f"drawtext=text='{scene_num}':"
        f"fontsize=48:"
        f"fontcolor=white:"
        f"x=w-tw-30:"
        f"y=30:"
        f"box=1:"
        f"boxcolor=black@0.5:"
        f"boxborderw=12:"
        f"fontfile=/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
    )


def normalize_clip(
    input_path: Path,
    output_path: Path,
    target_width: int = 1072,
    target_height: int = 1920,
    target_sample_rate: int = 44100,
    scene_num: int = None,
    trim_duration: float = None,
    trim_start: float = None,
) -> bool:
    """Normalize clip to consistent format for concatenation.

    Uses scale-to-fill + center crop to avoid black bars.
    Target 1072x1920 works for both Kling (1072x1928) and Veo3 (1080x1920).

    Args:
        scene_num: If set, overlay scene number badge in top-right corner (debug mode)
        trim_duration: If set, trim clip to this many seconds
        trim_start: If set, skip this many seconds from the beginning
    """
    # Scale to fill target (may overflow), then center crop
    # This avoids black bars by cropping excess instead of padding
    filters = [
        f"scale={target_width}:{target_height}:"
        f"force_original_aspect_ratio=increase,"
        f"crop={target_width}:{target_height}"
    ]

    # Add scene number overlay for review mode
    if scene_num is not None:
        filters.append(_build_scene_number_filter(scene_num))

    filter_str = ",".join(filters)

    cmd = ["ffmpeg", "-y"]

    if trim_start is not None:
        cmd.extend(["-ss", str(trim_start)])

    cmd.extend([
        "-i", str(input_path),
        "-vf", filter_str,
        "-c:v", "libx264", "-preset", "fast", "-crf", "18",
        "-c:a", "aac", "-b:a", "192k", "-ar", str(target_sample_rate),
        "-pix_fmt", "yuv420p",
    ])

    if trim_duration is not None:
        cmd.extend(["-t", str(trim_duration)])

    cmd.append(str(output_path))

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        if result.returncode != 0:
            return False
        return True
    except Exception as e:
        return False


def _create_title_card_clip(
    image_path: Path,
    output_path: Path,
    duration: float = 3.0,
    target_width: int = 1072,
    target_height: int = 1920,
    target_fps: int = 24,
    target_sample_rate: int = 44100,
) -> Path | None:
    """Create a video clip from a static image with a slow zoom effect.

    Used for episode title cards. Generates a silent video clip from a JPEG/PNG.
    FPS must match scene clips (Kling=24fps) to avoid audio desync during concat.
    """
    # Scale image to high res, then zoompan with slow zoom-in (Ken Burns effect)
    # Using 4x upscale to minimize pixel-rounding jitter in zoompan
    fps = target_fps
    total_frames = int(duration * fps)
    filter_str = (
        f"scale={target_width * 4}:{target_height * 4}:flags=lanczos,"
        f"zoompan=z='1+0.06*on/{total_frames}':x='iw/2-(iw/zoom/2)':"
        f"y='ih/2-(ih/zoom/2)':d={total_frames}:s={target_width}x{target_height}:fps={fps},"
        f"format=yuv420p"
    )

    cmd = [
        "ffmpeg", "-y",
        "-loop", "1",
        "-i", str(image_path),
        "-f", "lavfi", "-i", f"anullsrc=r={target_sample_rate}:cl=stereo",
        "-vf", filter_str,
        "-c:v", "libx264", "-preset", "fast", "-crf", "18",
        "-c:a", "aac", "-b:a", "192k",
        "-r", str(fps),
        "-t", str(duration),
        "-shortest",
        "-pix_fmt", "yuv420p",
        str(output_path)
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        if result.returncode == 0 and output_path.exists():
            return output_path
        return None
    except Exception:
        return None


def concatenate_videos(video_paths: list, output_path: Path) -> bool:
    """Concatenate videos using FFmpeg concat demuxer with stream copy.

    All clips must have matching fps/codec/resolution before calling this.
    Title card clips are generated at matching fps to avoid audio desync.
    """

    # Create concat file
    concat_file = output_path.parent / "concat.txt"
    with open(concat_file, 'w') as f:
        for path in video_paths:
            f.write(f"file '{path.resolve()}'\n")

    cmd = [
        "ffmpeg", "-y",
        "-f", "concat", "-safe", "0",
        "-i", str(concat_file),
        "-c", "copy",
        str(output_path)
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
        if result.returncode != 0:
            return False
        return True
    except Exception as e:
        return False


def get_captions_for_scene(take_path: Path, scene: str, caption_iter: str) -> List[Dict]:
    """
    Load caption lines for a scene from 12-captions.

    Returns list of {text, start, end, position}
    """
    if not caption_iter:
        return []

    caption_file = take_path / "12-captions" / scene / caption_iter / "caption.json"
    if not caption_file.exists():
        return []

    with open(caption_file) as f:
        caption_data = json.load(f)

    lines = caption_data.get("lines", [])

    # Filter out empty lines
    return [line for line in lines if line.get("text", "").strip()]


def assemble(
    take_path: Path,
    no_captions: bool = False,
    final: bool = False
) -> dict:
    """
    Assemble final video from scene clips.

    Args:
        take_path: Path to take directory
        no_captions: Skip caption overlay
        final: If True, produce clean output without debug scene numbers.
               If False (default), overlay scene numbers in top-right corner for review.

    Returns:
        Result dict with status and output path
    """
    logger = setup_generation_log(take_path)
    logger.section("Stage 13: Assemble Video")

    # Load pipeline
    pipeline = load_pipeline(take_path)
    video_selections = pipeline.get("10-videos", {})
    caption_selections = pipeline.get("12-captions", {})
    scene_trims = pipeline.get("scene_trims", {})
    scene_trim_starts = pipeline.get("scene_trim_starts", {})

    if not isinstance(video_selections, dict):
        return {
            "status": "error",
            "error": "No videos selected. Complete Stage 10 first."
        }

    # Get scene count from video selections
    scene_count = len(video_selections)
    if scene_count == 0:
        return {
            "status": "error",
            "error": "No videos selected."
        }

    # Create assembly output directory
    assembly_dir = take_path / "13-assembly"
    ensure_dir(assembly_dir)

    # Get next iteration
    iteration = get_next_iteration(assembly_dir)
    iter_path = assembly_dir / iteration
    ensure_dir(iter_path)

    temp_dir = iter_path / "temp"
    ensure_dir(temp_dir)

    logger.info(f"Assembly iteration: {iteration}")
    if final:
        logger.info("Mode: FINAL (clean output, no scene numbers)")
    else:
        logger.info("Mode: REVIEW (scene numbers in top-right corner)")

    # Process each scene
    processed_clips = []
    scene_info = []
    total_duration = 0.0
    caption_count = 0

    # Prepend intro title card if available (from 11-titles stage or legacy path)
    title_selections = pipeline.get("11-titles", {}) if isinstance(pipeline.get("11-titles"), dict) else {}
    intro_title_path = None
    if title_selections.get("intro"):
        intro_path = take_path / "11-titles" / "intro" / f"{title_selections['intro']}.jpg"
        if intro_path.exists():
            intro_title_path = intro_path
    if not intro_title_path:
        # Legacy fallback: episode_title_card.jpg at take root
        legacy_path = take_path / "episode_title_card.jpg"
        if legacy_path.exists():
            intro_title_path = legacy_path

    if intro_title_path:
        title_durations = pipeline.get("title_durations", {})
        intro_duration = title_durations.get("intro", 3.0)
        title_clip = _create_title_card_clip(
            intro_title_path, temp_dir / "title_card_intro.mp4",
            duration=intro_duration, target_width=1072, target_height=1920,
        )
        if title_clip:
            processed_clips.append(title_clip)
            total_duration += intro_duration
            logger.info(f"Prepended intro title card ({intro_duration}s): {intro_title_path.name}")
            scene_info.append({
                "scene": "title_card_intro",
                "video_iteration": None,
                "caption_iteration": None,
                "duration": intro_duration,
                "caption_lines": 0,
            })

    # Sort scenes by custom order or numeric default
    scene_order = pipeline.get("scene_order", None)
    if scene_order:
        # Use custom order — ONLY include listed scenes (unlisted are skipped)
        sorted_scenes = [s for s in scene_order if s in video_selections]
    else:
        sorted_scenes = sorted(video_selections.keys(), key=lambda x: int(x.split('-')[0]))

    for scene in sorted_scenes:
        video_iter = video_selections.get(scene)
        if not video_iter:
            logger.warning(f"{scene}: No video selected - skipping")
            continue

        # Skip interpolation end-frame scenes (no independent video)
        if video_iter == "interpolation-end":
            logger.info(f"{scene}: Interpolation end-frame scene - skipping")
            continue

        # Get video path
        video_path = take_path / "10-videos" / scene / video_iter / "video.mp4"
        if not video_path.exists():
            logger.warning(f"{scene}: Video not found: {video_path}")
            continue

        # Get video duration and apply trim if configured
        duration = get_video_duration(video_path)
        trim_start = scene_trim_starts.get(scene)
        if trim_start is not None:
            trim_start = float(trim_start)
        trim_duration = scene_trims.get(scene)
        if trim_start is not None or trim_duration is not None:
            effective = duration
            if trim_start:
                effective -= trim_start
            if trim_duration is not None:
                trim_duration = float(trim_duration)
                effective = min(effective, trim_duration)
            logger.info(f"{scene}: Trim start={trim_start or 0}s, duration={trim_duration or 'full'}s → {effective:.1f}s")
            duration = effective
        total_duration += duration

        scene_num = int(scene.split('-')[0])

        # Get captions for this scene
        caption_iter = caption_selections.get(scene) if isinstance(caption_selections, dict) else None
        captions = []
        if not no_captions and caption_iter:
            captions = get_captions_for_scene(take_path, scene, caption_iter)

        # Scene number for review mode (None in final mode)
        debug_scene_num = scene_num if not final else None

        if captions:
            # Apply all captions in single FFmpeg pass
            logger.info(f"{scene}: {len(captions)} caption line(s) [{duration:.1f}s]")
            caption_count += len(captions)

            for idx, ct in enumerate(captions):
                text = ct.get("text", "")
                start = ct.get("start", 0.5)
                end = ct.get("end", duration - 0.5)
                logger.info(f"  Line {idx + 1}: '{text[:40]}...' ({start:.1f}-{end:.1f}s)")

            output_captioned = temp_dir / f"{scene}_captioned.mp4"

            if add_all_captions(video_path, output_captioned, captions, scene_num=debug_scene_num, trim_duration=trim_duration, trim_start=trim_start):
                processed_clips.append(output_captioned)
            else:
                # Fallback: normalize original without captions
                logger.warning(f"  Caption overlay failed, using original")
                output_normalized = temp_dir / f"{scene}_normalized.mp4"
                if normalize_clip(video_path, output_normalized, scene_num=debug_scene_num, trim_duration=trim_duration, trim_start=trim_start):
                    processed_clips.append(output_normalized)
                else:
                    processed_clips.append(video_path)
        else:
            # No captions, just normalize
            output_normalized = temp_dir / f"{scene}_normalized.mp4"
            logger.info(f"{scene}: No caption [{duration:.1f}s]")

            if normalize_clip(video_path, output_normalized, scene_num=debug_scene_num, trim_duration=trim_duration, trim_start=trim_start):
                processed_clips.append(output_normalized)
            else:
                processed_clips.append(video_path)

        scene_info.append({
            "scene": scene,
            "video_iteration": video_iter,
            "caption_iteration": caption_iter,
            "duration": duration,
            "caption_lines": len(captions)
        })

    # Append outro title card if available (from 11-titles stage)
    if title_selections.get("outro"):
        outro_path = take_path / "11-titles" / "outro" / f"{title_selections['outro']}.jpg"
        if outro_path.exists():
            title_durations = pipeline.get("title_durations", {})
            outro_duration = title_durations.get("outro", 3.0)
            outro_clip = _create_title_card_clip(
                outro_path, temp_dir / "title_card_outro.mp4",
                duration=outro_duration, target_width=1072, target_height=1920,
            )
            if outro_clip:
                processed_clips.append(outro_clip)
                total_duration += outro_duration
                logger.info(f"Appended outro title card ({outro_duration}s): {outro_path.name}")
                scene_info.append({
                    "scene": "title_card_outro",
                    "video_iteration": None,
                    "caption_iteration": None,
                    "duration": outro_duration,
                    "caption_lines": 0,
                })

    if not processed_clips:
        return {
            "status": "error",
            "error": "No clips to assemble"
        }

    # Concatenate all clips
    output_path = iter_path / "final.mp4"
    logger.info(f"Concatenating {len(processed_clips)} clips...")

    if not concatenate_videos(processed_clips, output_path):
        return {
            "status": "error",
            "error": "Failed to concatenate videos"
        }

    # Get final video stats
    final_duration = get_video_duration(output_path)
    final_resolution = get_video_resolution(output_path)
    size_mb = output_path.stat().st_size / 1024 / 1024

    logger.info(f"Final video: {final_duration:.1f}s, {final_resolution[0]}x{final_resolution[1]}, {size_mb:.1f}MB")

    # Create streaming-optimized variants
    logger.info("Creating streaming-optimized variants...")
    variants = create_streaming_variants(output_path, logger)

    # Save metadata
    metadata = {
        "stage": "13-assembly",
        "iteration": iteration,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "scene_count": len(processed_clips),
        "scenes": scene_info,
        "caption_lines_total": caption_count,
        "no_captions": no_captions,
        "final": final,
        "final_duration": final_duration,
        "final_resolution": f"{final_resolution[0]}x{final_resolution[1]}",
        "final_size_mb": round(size_mb, 2),
        "output_file": "final.mp4",
        "variants": variants
    }

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

    # Update pipeline
    pipeline["13-assembly"] = iteration
    save_pipeline(take_path, pipeline)

    return {
        "status": "success",
        "iteration": iteration,
        "output_path": str(output_path),
        "duration": final_duration,
        "resolution": f"{final_resolution[0]}x{final_resolution[1]}",
        "size_mb": size_mb,
        "scene_count": len(processed_clips),
        "caption_count": caption_count,
        "variants": variants
    }


def status_assembly(take_path: Path) -> dict:
    """
    Get assembly status.

    Returns:
        Dict with assembly info
    """
    pipeline = load_pipeline(take_path)
    assembly_iter = pipeline.get("13-assembly")

    assembly_dir = take_path / "13-assembly"
    iterations = []

    if assembly_dir.exists():
        for item in sorted(assembly_dir.iterdir()):
            if item.is_dir() and item.name.isdigit():
                meta_file = item / "metadata.json"
                if meta_file.exists():
                    with open(meta_file) as f:
                        meta = json.load(f)
                    iterations.append({
                        "iteration": item.name,
                        "duration": meta.get("final_duration", 0),
                        "size_mb": meta.get("final_size_mb", 0),
                        "scene_count": meta.get("scene_count", 0),
                        "caption_count": meta.get("caption_lines_total", 0),
                        "created_at": meta.get("created_at", "")
                    })
                else:
                    iterations.append({"iteration": item.name})

    return {
        "status": "success",
        "selected": assembly_iter,
        "iterations": iterations
    }
