#!/usr/bin/env python3
"""
Generate videos from raw frames using Kling v2.6 Direct API.

Uses direct Kling API access (not Pollo proxy) for image-to-video generation with audio.
"""

import os
import sys
import json
import time
import base64
import requests
import jwt
import subprocess
import uuid
from pathlib import Path
from datetime import datetime
from dotenv import load_dotenv

# Load .env from repository root
load_dotenv(Path(__file__).parent.parent.parent / ".env")

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

# Kling API Configuration (Singapore endpoint)
KLING_API_BASE = os.getenv("KLING_API_BASE", "https://api-singapore.klingai.com")
KLING_ACCESS_KEY = os.getenv("KLING_ACCESS_KEY")
KLING_SECRET_KEY = os.getenv("KLING_SECRET_KEY")


def image_to_base64(image_path: str) -> str:
    """Convert image file to base64 string (without data URI prefix)."""
    with open(image_path, 'rb') as f:
        return base64.b64encode(f.read()).decode('utf-8')


def generate_jwt_token(access_key: str, secret_key: str) -> str:
    """
    Generate JWT token for Kling API authentication.

    Token format:
    - Header: HS256 algorithm
    - Payload: iss (access_key), exp (30 min), nbf (now - 5s)
    - Signature: HMAC with secret_key

    Args:
        access_key: Kling API Access Key
        secret_key: Kling API Secret Key

    Returns:
        JWT token string
    """
    headers = {
        "alg": "HS256",
        "typ": "JWT"
    }

    payload = {
        "iss": access_key,
        "exp": int(time.time()) + 1800,  # 30 minutes validity
        "nbf": int(time.time()) - 5  # Not before: 5 seconds ago
    }

    return jwt.encode(payload, secret_key, algorithm="HS256", headers=headers)


def call_kling_direct(
    prompt: str,
    image_path: str,
    duration: int = 5,
    generate_audio: bool = True,
    negative_prompt: str = ""
) -> dict:
    """
    Call Kling v2.6 API directly for image-to-video generation.

    Based on official Kling API documentation for image2video endpoint.

    Args:
        prompt: Video generation prompt (max 2500 chars)
        image_path: Path to reference image
        duration: Video duration in seconds ("5" or "10")
        generate_audio: Generate audio for the video (v2.6+ only)
        negative_prompt: Negative prompt (max 2500 chars)

    Returns:
        API response with task_id
    """
    if not KLING_ACCESS_KEY:
        raise ValueError("KLING_ACCESS_KEY environment variable not set")
    if not KLING_SECRET_KEY:
        raise ValueError("KLING_SECRET_KEY environment variable not set")

    # Generate JWT token for authentication
    token = generate_jwt_token(KLING_ACCESS_KEY, KLING_SECRET_KEY)

    # Encode image to base64 (no data URI prefix)
    print(f"  Encoding image: {image_path}")
    image_base64 = image_to_base64(image_path)

    # Build request payload according to Kling API docs
    payload = {
        "model_name": "kling-v2-6",
        "mode": "pro",  # Required for audio generation
        "image": image_base64,
        "prompt": prompt[:2500],  # Max 2500 chars
        "duration": str(duration),  # Must be string: "5" or "10"
        "sound": "on" if generate_audio else "off",  # v2.6+ only
    }

    if negative_prompt:
        payload["negative_prompt"] = negative_prompt[:2500]

    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }

    print(f"  Calling Kling API...")
    print(f"  Model: kling-v2-6")
    print(f"  Duration: {duration}s")
    print(f"  Audio: {'on' if generate_audio else 'off'}")

    response = requests.post(
        f"{KLING_API_BASE}/v1/videos/image2video",
        headers=headers,
        json=payload,
        timeout=60
    )

    if response.status_code != 200:
        raise RuntimeError(f"Kling API error: {response.status_code} - {response.text[:500]}")

    result = response.json()

    # Check for API error
    if result.get("code") != 0:
        error_msg = result.get("message", "Unknown error")
        raise RuntimeError(f"Kling API error: {error_msg}")

    # Extract task_id from response
    data = result.get("data", {})
    task_id = data.get("task_id")
    if not task_id:
        raise RuntimeError(f"No task_id in response: {result}")

    task_status = data.get("task_status")
    print(f"  Task created: {task_id}")
    print(f"  Initial status: {task_status}")

    return {
        "task_id": task_id,
        "task_status": task_status,
        "raw_response": result
    }


def poll_kling_status(task_id: str, poll_interval: int = 10, timeout: int = 600) -> dict:
    """
    Poll Kling API for task completion.

    Args:
        task_id: Task ID from initial request
        poll_interval: Seconds between polls
        timeout: Maximum seconds to wait

    Returns:
        Final task result with video URL
    """
    if not KLING_ACCESS_KEY:
        raise ValueError("KLING_ACCESS_KEY environment variable not set")
    if not KLING_SECRET_KEY:
        raise ValueError("KLING_SECRET_KEY environment variable not set")

    # Generate JWT token for authentication
    token = generate_jwt_token(KLING_ACCESS_KEY, KLING_SECRET_KEY)

    status_url = f"{KLING_API_BASE}/v1/videos/image2video/{task_id}"
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }

    start_time = time.time()

    while time.time() - start_time < timeout:
        response = requests.get(status_url, headers=headers, timeout=30)

        if response.status_code != 200:
            raise RuntimeError(f"Status check error: {response.status_code} - {response.text[:500]}")

        result = response.json()

        # Check for API error
        if result.get("code") != 0:
            error_msg = result.get("message", "Unknown error")
            raise RuntimeError(f"API error: {error_msg}")

        data = result.get("data", {})
        task_status = data.get("task_status")

        print(f"  Status: {task_status}")

        # Task statuses: submitted, processing, succeed, failed
        if task_status == "succeed":
            print(f"  Video generation completed!")
            return data
        elif task_status == "failed":
            error_msg = data.get("task_status_msg", "Unknown error")
            raise RuntimeError(f"Video generation failed: {error_msg}")

        # Still processing (submitted or processing)
        time.sleep(poll_interval)

    raise TimeoutError(f"Video generation timed out after {timeout}s")


def download_video(video_url: str, output_path: str):
    """Download video from URL."""
    print(f"  Downloading: {video_url}")
    response = requests.get(video_url, stream=True, timeout=120)
    response.raise_for_status()

    with open(output_path, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)

    print(f"  Saved: {output_path}")


def build_video_prompt(scene_data: dict, style_info: dict, feedback: str = None, previous_feedback: str = None) -> str:
    """
    Build video prompt from scene data and style.

    Args:
        scene_data: Scene from adapted script
        style_info: Style information from script
        feedback: Optional current user feedback to incorporate
        previous_feedback: Optional accumulated feedback from previous iterations

    Returns:
        Video generation prompt
    """
    # If feedback provided, use Claude to regenerate the entire prompt
    if feedback:
        print("  Using Claude to regenerate prompt with feedback...")
        return regenerate_prompt_with_claude(scene_data, style_info, feedback, previous_feedback)

    lines = []

    # Start with ACTION description - what HAPPENS in the video (MOST IMPORTANT)
    action = scene_data.get("action", "")
    what_is_shown = scene_data.get("what_is_shown", "")

    # Combine action with visual setting - ACTION FIRST
    if action:
        lines.append(f"ACTION: {action}.")
        lines.append(f"SETTING: {what_is_shown}")
    else:
        lines.append(what_is_shown)

    # Add camera movement
    camera = scene_data.get("camera", "static")
    if camera and camera != "static":
        lines.append(f"\nCamera: {camera}")

    # Add style
    visual_style = style_info.get("visual", "cinematic")
    tone = style_info.get("tone", "dramatic")
    lines.append(f"\nStyle: {visual_style}, {tone}, photorealistic")

    # Add audio description if present
    audio_ambient = scene_data.get("audio_ambient", "")
    if audio_ambient:
        lines.append(f"\nAudio: {audio_ambient}")

    # Add dialogue if present - DON'T include character variable names
    dialogue = scene_data.get("dialogue")
    if dialogue and dialogue.get("text"):
        # Skip character variable name like "biker_2"
        delivery = dialogue.get('delivery', 'natural')
        text = dialogue.get('text')
        lines.append(f"\nDialogue spoken with {delivery} delivery: \"{text}\"")

    # Add negative prompts
    lines.append("\n\nNO text overlays, NO subtitles, NO watermarks, NO distorted faces or hands")

    return "\n".join(lines)


def regenerate_prompt_with_claude(scene_data: dict, style_info: dict, feedback: str, previous_feedback: str = None) -> str:
    """
    Use Claude Code CLI to regenerate prompt incorporating user feedback.

    Args:
        scene_data: Scene from adapted script
        style_info: Style information from script
        feedback: Current user feedback about what to change
        previous_feedback: Accumulated feedback from previous iterations

    Returns:
        Regenerated video prompt
    """
    # Build context for Claude
    original_action = scene_data.get("action", "")
    what_is_shown = scene_data.get("what_is_shown", "")
    dialogue = scene_data.get("dialogue", {})
    audio = scene_data.get("audio_ambient", "")
    camera = scene_data.get("camera", "static")

    # Build feedback section
    feedback_section = f"MOST RECENT FEEDBACK (apply this):\n{feedback}"
    if previous_feedback:
        feedback_section += f"\n\nPREVIOUS FEEDBACK (already addressed, for context):\n{previous_feedback}"

    claude_prompt = f"""You are writing a video generation prompt for Kling AI video model.

SCENE CONTEXT:
- Action: {original_action}
- Visual setting: {what_is_shown}
- Camera: {camera}
- Audio: {audio}
- Dialogue: {dialogue.get('text', 'none') if dialogue else 'none'}
- Style: {style_info.get('visual', 'cinematic')}, {style_info.get('tone', 'dramatic')}

{feedback_section}

IMPORTANT RULES:
1. Describe the ACTION step-by-step: what happens from start to finish
2. NEVER use character variable names like "biker_2" - instead describe them (e.g., "the tall biker", "biker with slicked hair")
3. Describe physical actions clearly and BRIEFLY - if feedback says "short", make it minimal
4. If feedback says "no dialogue", remove all dialogue
5. Keep it under 500 characters for Kling model
6. Include camera movement if specified
7. End with audio description
8. Add negative prompts: NO text overlays, NO subtitles, NO watermarks
9. Pay attention to scale/intensity words: "SHORT", "SMALL", "BRIEF" mean minimize the action

Write ONLY the final video prompt, nothing else."""

    # Generate unique session ID
    session_id = str(uuid.uuid4())

    try:
        # Call Claude Code CLI
        result = subprocess.run(
            [
                "claude",
                "--session-id", session_id,
                "-p", claude_prompt
            ],
            capture_output=True,
            text=True,
            timeout=120
        )

        if result.returncode != 0:
            print(f"  Warning: Claude returned code {result.returncode}")
            print(f"  Stderr: {result.stderr}")
            # Fallback to basic prompt building
            return build_basic_prompt_with_feedback(scene_data, style_info, feedback)

        # Extract prompt from Claude output
        output = result.stdout.strip()
        if output:
            return output
        else:
            # Fallback
            return build_basic_prompt_with_feedback(scene_data, style_info, feedback)

    except subprocess.TimeoutExpired:
        print("  Warning: Claude timed out, using fallback")
        return build_basic_prompt_with_feedback(scene_data, style_info, feedback)
    except Exception as e:
        print(f"  Warning: Claude error: {e}, using fallback")
        return build_basic_prompt_with_feedback(scene_data, style_info, feedback)


def build_basic_prompt_with_feedback(scene_data: dict, style_info: dict, feedback: str) -> str:
    """
    Fallback: Build prompt incorporating feedback manually.
    """
    lines = []

    # Start with action
    action = scene_data.get("action", "")
    what_is_shown = scene_data.get("what_is_shown", "")

    if action:
        lines.append(f"ACTION: {action}. {what_is_shown}")
    else:
        lines.append(what_is_shown)

    # Camera
    camera = scene_data.get("camera", "static")
    if camera and camera != "static":
        lines.append(f"Camera: {camera}")

    # Style
    lines.append(f"Style: {style_info.get('visual', 'cinematic')}, {style_info.get('tone', 'dramatic')}, photorealistic")

    # Audio
    audio = scene_data.get("audio_ambient", "")
    if audio:
        lines.append(f"Audio: {audio}")

    # Incorporate feedback
    lines.append(f"\nIMPORTANT: {feedback}")

    # Negative prompts
    lines.append("\nNO text overlays, NO subtitles, NO watermarks, NO distorted faces or hands")

    return "\n".join(lines)


def generate_scene_video(
    story_name: str,
    scene_num: int,
    take: str = "current",
    feedback: str = None
) -> bool:
    """
    Generate video for a single scene using Kling v2.6 Direct API.

    Args:
        story_name: Story name
        scene_num: Scene number (1-9)
        take: Take number or "current"
        feedback: User feedback for iteration

    Returns:
        True if successful
    """
    # Get paths
    try:
        story_path = get_story_path(story_name)
        take_path = get_take_path(story_name, take)
    except Exception as e:
        print(f"Error: {e}")
        return False

    # Get take number
    if take == "current":
        current_link = story_path / "current"
        if not current_link.is_symlink():
            print(f"Error: No current take set")
            return False
        take_num = current_link.readlink().name
    else:
        take_num = take

    # Get raw frame
    raw_frame_path = take_path / "09-frames" / f"scene_{scene_num:02d}" / "raw_frame.jpg"
    if not raw_frame_path.exists():
        print(f"Error: Raw frame not found: {raw_frame_path}")
        return False

    # 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 False

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

    # Find scene in script
    scene_data = None
    for scene in script["scenes"]:
        if scene["number"] == scene_num:
            scene_data = scene
            break

    if not scene_data:
        print(f"Error: Scene {scene_num} not found in script")
        return False

    # Video output directory
    video_stage_path = take_path / "10-videos" / f"scene_{scene_num:02d}"
    ensure_dir(video_stage_path)

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

    iteration_num = max(existing_iterations) + 1 if existing_iterations else 1
    iteration = f"{iteration_num:03d}"
    iter_path = video_stage_path / iteration
    ensure_dir(iter_path)

    print(f"Generating video:")
    print(f"  Story: {story_name}")
    print(f"  Take: {take_num}")
    print(f"  Scene: {scene_num}")
    print(f"  Iteration: {iteration}")
    print(f"  Raw frame: {raw_frame_path}")
    if feedback:
        print(f"  With feedback: {feedback}")
    print()

    # Collect all previous feedback from existing iterations
    previous_feedback_list = []
    if existing_iterations and feedback:
        for iter_num in sorted(existing_iterations):
            prev_iter_path = video_stage_path / f"{iter_num:03d}"
            prev_feedback_file = prev_iter_path / "feedback.md"
            if prev_feedback_file.exists():
                prev_content = prev_feedback_file.read_text()
                # Extract just the feedback text (skip the "# User Feedback" header)
                feedback_text = prev_content.replace("# User Feedback", "").strip()
                if feedback_text:
                    previous_feedback_list.append(f"Iteration {iter_num:03d}: {feedback_text}")

    previous_feedback = "\n\n".join(previous_feedback_list) if previous_feedback_list else None

    if previous_feedback:
        print(f"  Previous feedback from {len(previous_feedback_list)} iteration(s)")

    # Build prompt (with feedback if provided)
    style_info = script.get("style", {})
    prompt = build_video_prompt(scene_data, style_info, feedback, previous_feedback)

    # Save prompt
    prompt_file = iter_path / "prompt.txt"
    prompt_file.write_text(prompt)
    print(f"✓ Saved prompt: {prompt_file}")

    if feedback:
        feedback_file = iter_path / "feedback.md"
        feedback_file.write_text(f"# User Feedback\n\n{feedback}\n")
        print(f"✓ Saved feedback: {feedback_file}")

    # Generate video
    try:
        # Kling v2.6 only supports 5 or 10 second durations
        # Logic: Always 5s, UNLESS dialogue length > 5s (then use 10s)
        duration = 5  # Default

        dialogue = scene_data.get("dialogue")
        if dialogue and dialogue.get("text"):
            # Estimate dialogue duration: ~150 words per minute = 2.5 words per second
            # Add buffer for natural speech pacing
            dialogue_text = dialogue.get("text", "")
            word_count = len(dialogue_text.split())
            estimated_duration = (word_count / 2.5) * 1.3  # 30% buffer for pacing

            if estimated_duration > 5:
                duration = 10
                print(f"  Note: Dialogue estimated at {estimated_duration:.1f}s, using 10s duration")
            else:
                print(f"  Note: Dialogue estimated at {estimated_duration:.1f}s, using 5s duration")
        else:
            print(f"  Note: No dialogue, using 5s duration")

        # Call Kling API
        result = call_kling_direct(
            prompt=prompt,
            image_path=str(raw_frame_path),
            duration=duration,
            generate_audio=True
        )

        task_id = result["task_id"]

        # Poll for completion
        print(f"\nPolling for completion...")
        video_result = poll_kling_status(task_id, poll_interval=10, timeout=600)

        # Extract video URL from task_result.videos[0].url
        task_result = video_result.get("task_result", {})
        videos = task_result.get("videos", [])
        if not videos:
            print(f"Error: No videos in response: {video_result}")
            return False

        video_info = videos[0]
        video_url = video_info.get("url")
        video_id = video_info.get("id")
        video_duration = video_info.get("duration")

        if not video_url:
            print(f"Error: No video URL in response")
            return False

        print(f"  Video ID: {video_id}")
        print(f"  Duration: {video_duration}s")

        # Download video
        output_path = iter_path / "video.mp4"
        download_video(video_url, str(output_path))

        # Save metadata
        metadata = {
            "scene": scene_num,
            "iteration": iteration,
            "timestamp": datetime.now().isoformat(),
            "source_raw_frame": "../../09-frames/scene_{:02d}/raw_frame.jpg".format(scene_num),
            "source_script": "../../08-script-adapted/current/script.json",
            "model": "kling-v2.6-direct",
            "api": "direct",
            "task_id": task_id,
            "video_url": video_url,
            "duration_requested": duration,
            "aspect_ratio": "9:16",
            "audio_generated": True,
            "prompt_file": "prompt.txt",
            "feedback": feedback,
            "scene_description": scene_data.get("what_is_shown", "")[:200]
        }

        metadata_file = iter_path / "metadata.json"
        save_metadata(metadata_file, metadata)
        print(f"✓ Saved metadata: {metadata_file}")

        # Update current symlink
        update_current_symlink(video_stage_path, iteration)
        print(f"✓ Updated current → {iteration}")

        print(f"\n{'='*60}")
        print(f"Scene {scene_num} video generation completed!")
        print(f"{'='*60}")
        print(f"Output: {output_path}")

        return True

    except Exception as e:
        print(f"\nError generating video: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    import argparse
    parser = argparse.ArgumentParser(
        description="Generate videos from raw frames using Kling v2.6 Direct API"
    )
    parser.add_argument("story_name", help="Story name")
    parser.add_argument("scene_num", type=int, help="Scene number (1-9)")
    parser.add_argument(
        "--take",
        default="current",
        help="Take number (default: current)"
    )
    parser.add_argument(
        "--feedback",
        help="User feedback for iteration"
    )

    args = parser.parse_args()

    success = generate_scene_video(
        args.story_name,
        args.scene_num,
        args.take,
        args.feedback
    )

    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
