#!/usr/bin/env python3
"""
Microdrama Pipeline - Automated satirical news video generation.

Full pipeline that:
1. Fetches article from panorama.pub
2. Uses Claude to analyze and create script
3. Searches Yandex for reference images
4. Uses Claude to select best references
5. Generates frames with NanoBanana
6. Generates video clips (Veo3 + Ken Burns)
7. Assembles final video

Usage:
    python microdrama_pipeline.py "https://panorama.pub/news/article-slug"
    python microdrama_pipeline.py --task-dir /path/to/existing/task  # Resume
"""

import os
import sys
import json
import argparse
import subprocess
import time
import base64
import requests
from pathlib import Path
from datetime import datetime
from typing import Optional, Dict, Any, List

# Load .env
from dotenv import load_dotenv
env_path = Path(__file__).parent.parent.parent / ".env"
load_dotenv(env_path)

# Import local modules
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))
from laozhang_proxy import call_nanobanana_laozhang, call_veo3_laozhang

# Import microdrama styles
from microdrama_styles import (
    select_style_for_content,
    get_style_config,
    get_kind_config,
    get_style_prompt_context,
    get_frame_style_hints,
    format_style_for_display,
    VISUAL_STYLES,
    VIDEO_KINDS
)

# Constants
BASE_DIR = Path("/mnt/video/microdrama")
SCRIPTS_DIR = Path(__file__).parent
PROMPTS_DIR = SCRIPTS_DIR.parent / "prompts" / "stages"

# LaoZhang API for Claude
LAOZHANG_CLAUDE_URL = "https://api.laozhang.ai/v1/chat/completions"


def setup_logging(task_dir: Path, task_id: str):
    """
    Setup real-time logging to task-specific file.

    Creates logger that writes to both:
    - stdout (for terminal)
    - {task_dir}/pipeline.log (for real-time debugging)
    """
    import logging

    log_file = task_dir / "pipeline.log"

    # Create logger for this task
    logger = logging.getLogger(f"microdrama_{task_id}")
    logger.setLevel(logging.INFO)
    logger.propagate = False

    # Remove existing handlers
    logger.handlers.clear()

    # File handler - real-time writes with flush
    file_handler = logging.FileHandler(log_file, mode='w', encoding='utf-8')
    file_handler.setLevel(logging.INFO)

    # Console handler
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)

    # Format: timestamp [task_id] message
    formatter = logging.Formatter(
        f'%(asctime)s [{task_id}] %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(logging.Formatter('%(message)s'))

    logger.addHandler(file_handler)
    logger.addHandler(console_handler)

    return logger


def save_generation_log(task_dir: Path, data: Dict[str, Any]):
    """Save detailed generation log as JSON."""
    log_path = task_dir / "generation.log"
    with open(log_path, "w") as f:
        json.dump(data, f, ensure_ascii=False, indent=2, default=str)


# ============================================================================
# INTERACTIVE MODE
# ============================================================================

def analyze_story_for_understanding(article_text: str, api_key: str) -> Dict[str, Any]:
    """
    First pass: Claude analyzes article and explains what's funny.
    Returns understanding for user approval before script generation.
    """
    prompt = f"""Analyze this satirical news article. Your task is to UNDERSTAND THE JOKE, not create a script yet.

ARTICLE:
{article_text}

Respond in JSON format:
{{
  "joke_summary": "One sentence: what's the core joke/absurdity?",
  "funny_elements": [
    "List 3-5 specific funny details or absurdities in the article"
  ],
  "story_angle": "What narrative angle would work best for a video?",
  "suggested_storytellers": [
    {{
      "type": "news_anchor|victim|witness|expert|unexpected",
      "description": "Who this would be and why they work for this story",
      "example": "e.g., 'The zoo director defending his import-substituted panda'"
    }}
  ],
  "suggested_styles": [
    {{
      "style": "style_key",
      "why": "Why this visual style fits the joke"
    }}
  ],
  "tone_recommendation": "deadpan serious, absurdist, melodrama, etc."
}}

Be specific about WHY things are funny. The goal is to show understanding before we create the script."""

    response = call_claude(prompt, api_key)

    # Parse JSON from response
    json_start = response.find("{")
    json_end = response.rfind("}") + 1
    if json_start >= 0 and json_end > json_start:
        return json.loads(response[json_start:json_end])
    return {"error": "Failed to parse understanding", "raw": response}


def save_interactive_questions(understanding: Dict, available_styles: List[str], task_dir: Path) -> Path:
    """
    Save interactive questions to a file for user to answer.
    Returns path to the questions file.
    """
    questions_file = task_dir / "interactive_questions.md"

    content = """# Interactive Questions - Please Answer Below

## Story Understanding

**JOKE SUMMARY:**
{summary}

**FUNNY ELEMENTS:**
{elements}

**SUGGESTED STORY ANGLE:**
{angle}

**RECOMMENDED TONE:** {tone}

---

## YOUR ANSWERS (edit below)

### Question 1: Is the understanding correct?
(Leave empty if correct, or write corrections)

ANSWER_1:

### Question 2: Who should tell this story?

SUGGESTED STORYTELLERS:
{storytellers}

(Enter number 1-{num_storytellers}, or describe custom storyteller)

ANSWER_2: 1

### Question 3: Visual style?

SUGGESTED STYLES:
{styles}

AVAILABLE STYLES: {available}

(Enter number or style name)

ANSWER_3: 1

---

After editing, run:
python microdrama_pipeline.py --task-dir {task_dir} --from-stage refs
""".format(
        summary=understanding.get('joke_summary', 'N/A'),
        elements='\n'.join(f"  {i}. {e}" for i, e in enumerate(understanding.get('funny_elements', []), 1)),
        angle=understanding.get('story_angle', 'N/A'),
        tone=understanding.get('tone_recommendation', 'N/A'),
        storytellers='\n'.join(
            f"  {i}. {st['type'].upper()}: {st['description']}"
            for i, st in enumerate(understanding.get('suggested_storytellers', []), 1)
        ),
        num_storytellers=len(understanding.get('suggested_storytellers', [])),
        styles='\n'.join(
            f"  {i}. {st['style']}: {st['why']}"
            for i, st in enumerate(understanding.get('suggested_styles', []), 1)
        ),
        available=', '.join(available_styles[:10]),
        task_dir=task_dir
    )

    with open(questions_file, 'w') as f:
        f.write(content)

    return questions_file


def load_interactive_answers(task_dir: Path, understanding: Dict) -> Dict:
    """Load answers from the interactive questions file."""
    questions_file = task_dir / "interactive_questions.md"
    if not questions_file.exists():
        return {}

    with open(questions_file) as f:
        content = f.read()

    answers = {}

    # Parse ANSWER_1 (correction)
    if "ANSWER_1:" in content:
        answer1 = content.split("ANSWER_1:")[1].split("\n")[0].strip()
        answers['correction'] = answer1 if answer1 else None

    # Parse ANSWER_2 (storyteller)
    if "ANSWER_2:" in content:
        answer2 = content.split("ANSWER_2:")[1].split("\n")[0].strip()
        storytellers = understanding.get('suggested_storytellers', [])
        if answer2.isdigit() and 1 <= int(answer2) <= len(storytellers):
            answers['storyteller'] = storytellers[int(answer2) - 1]
        elif answer2:
            answers['storyteller'] = {"type": "custom", "description": answer2}

    # Parse ANSWER_3 (style)
    if "ANSWER_3:" in content:
        answer3 = content.split("ANSWER_3:")[1].split("\n")[0].strip()
        styles = understanding.get('suggested_styles', [])
        if answer3.isdigit() and 1 <= int(answer3) <= len(styles):
            answers['style'] = styles[int(answer3) - 1]['style']
        elif answer3:
            answers['style'] = answer3

    # Parse VISUAL CONCEPT (optional multi-line)
    if "VISUAL CONCEPT:" in content:
        # Get everything after VISUAL CONCEPT: until next --- or end
        concept_text = content.split("VISUAL CONCEPT:")[1]
        # Stop at --- or end of file
        if "---" in concept_text:
            concept_text = concept_text.split("---")[0]
        answers['visual_concept'] = concept_text.strip()

    return answers


def interactive_phase1_cli(understanding: Dict, available_styles: List[str]) -> Dict[str, str]:
    """
    Interactive CLI for Phase 1: Understanding + Choices.
    Returns user's choices for storyteller and style.
    NOTE: This requires stdin - use save_interactive_questions() for file-based mode.
    """
    print("\n" + "=" * 60)
    print("PHASE 1: STORY UNDERSTANDING")
    print("=" * 60)

    print(f"\n📖 JOKE SUMMARY:")
    print(f"   {understanding.get('joke_summary', 'N/A')}")

    print(f"\n😂 FUNNY ELEMENTS:")
    for i, elem in enumerate(understanding.get('funny_elements', []), 1):
        print(f"   {i}. {elem}")

    print(f"\n🎬 SUGGESTED STORY ANGLE:")
    print(f"   {understanding.get('story_angle', 'N/A')}")

    print(f"\n🎭 TONE: {understanding.get('tone_recommendation', 'N/A')}")

    # Question 1: Is understanding correct?
    print("\n" + "-" * 40)
    print("QUESTION 1: Is this understanding correct?")
    print("-" * 40)
    correction = input("Press Enter if correct, or type corrections: ").strip()

    # Question 2: Choose storyteller
    print("\n" + "-" * 40)
    print("QUESTION 2: Who should tell this story?")
    print("-" * 40)
    storytellers = understanding.get('suggested_storytellers', [])
    for i, st in enumerate(storytellers, 1):
        print(f"  {i}. {st['type'].upper()}: {st['description']}")
        print(f"     Example: {st.get('example', '')}")
    print(f"  0. Custom (describe your own)")

    choice = input("\nChoose (1-{0} or describe custom): ".format(len(storytellers))).strip()
    if choice.isdigit() and 1 <= int(choice) <= len(storytellers):
        storyteller_choice = storytellers[int(choice) - 1]
    elif choice == "0" or not choice.isdigit():
        storyteller_choice = {"type": "custom", "description": choice if choice != "0" else input("Describe storyteller: ")}
    else:
        storyteller_choice = storytellers[0] if storytellers else {"type": "news_anchor", "description": "Default narrator"}

    # Question 3: Choose style
    print("\n" + "-" * 40)
    print("QUESTION 3: Visual style and tone?")
    print("-" * 40)
    suggested_styles = understanding.get('suggested_styles', [])
    for i, st in enumerate(suggested_styles, 1):
        print(f"  {i}. {st['style']}: {st['why']}")
    print(f"\n  Available styles: {', '.join(available_styles[:10])}...")

    style_input = input("\nChoose number or type style name: ").strip()
    if style_input.isdigit() and 1 <= int(style_input) <= len(suggested_styles):
        style_choice = suggested_styles[int(style_input) - 1]['style']
    elif style_input in available_styles:
        style_choice = style_input
    else:
        style_choice = suggested_styles[0]['style'] if suggested_styles else "documentary"

    return {
        "correction": correction,
        "storyteller": storyteller_choice,
        "style": style_choice
    }


def format_script_for_review(script_data: Dict) -> str:
    """Format script scenes for user review."""
    lines = []
    lines.append("\n" + "=" * 60)
    lines.append("PHASE 2: SCRIPT REVIEW")
    lines.append("=" * 60)

    lines.append(f"\n🎬 Style: {script_data.get('visual_style', 'N/A')}")
    lines.append(f"📺 Kind: {script_data.get('video_kind', 'N/A')}")

    storyteller = script_data.get('storyteller', {})
    lines.append(f"\n🎭 STORYTELLER: {storyteller.get('type', 'N/A')}")
    lines.append(f"   {storyteller.get('description', '')}")

    lines.append("\n" + "-" * 40)
    lines.append("SCENES:")
    lines.append("-" * 40)

    for scene in script_data.get('scenes', []):
        scene_id = scene.get('scene_id', '?')
        scene_type = scene.get('type', '?')
        method = scene.get('method', 'veo3')

        lines.append(f"\n[Scene {scene_id}] {scene_type.upper()} ({method})")

        # Captions
        if 'caption_first' in scene:
            lines.append(f"  📝 First: \"{scene['caption_first']}\"")
            lines.append(f"  📝 Last:  \"{scene.get('caption_last', '')}\"")
        elif 'caption' in scene:
            lines.append(f"  📝 Caption: \"{scene['caption']}\"")

        # Frames
        frames = scene.get('frames', {})
        lines.append(f"  🖼️  First: {frames.get('first', '')[:60]}...")
        lines.append(f"  🖼️  Last:  {frames.get('last', '')[:60]}...")

        # Characters
        chars = scene.get('characters', [])
        if chars:
            lines.append(f"  👥 Characters: {', '.join(chars)}")

    return "\n".join(lines)


def interactive_phase2_cli(script_data: Dict) -> Dict[str, Any]:
    """
    Interactive CLI for Phase 2: Script Review.
    Returns approval or modifications.
    """
    print(format_script_for_review(script_data))

    print("\n" + "-" * 40)
    print("REVIEW THE SCRIPT ABOVE")
    print("-" * 40)
    print("Options:")
    print("  [Enter] - Approve and continue")
    print("  [m]     - Modify (provide feedback)")
    print("  [r]     - Regenerate with new instructions")

    choice = input("\nYour choice: ").strip().lower()

    if choice == 'm':
        feedback = input("What should be changed? ").strip()
        return {"action": "modify", "feedback": feedback}
    elif choice == 'r':
        feedback = input("New instructions for script: ").strip()
        return {"action": "regenerate", "feedback": feedback}
    else:
        return {"action": "approve"}


def call_claude(prompt: str, api_key: str, model: str = "claude-sonnet-4-20250514", images: List[Dict] = None) -> str:
    """Call Claude via LaoZhang API."""
    messages = [{"role": "user", "content": []}]

    # Add images if provided
    if images:
        for img in images:
            messages[0]["content"].append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:{img['mime_type']};base64,{img['base64']}"
                }
            })

    # Add text prompt
    messages[0]["content"].append({"type": "text", "text": prompt})

    response = requests.post(
        LAOZHANG_CLAUDE_URL,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        },
        json={
            "model": model,
            "messages": messages,
            "max_tokens": 8000
        },
        timeout=120
    )

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

    return response.json()["choices"][0]["message"]["content"]


def fetch_article(url: str, task_dir: Path, logger) -> Dict[str, Any]:
    """Fetch article from panorama.pub using requests."""
    logger.info(f"Fetching article: {url}")

    # For now, just extract basic info from URL
    # TODO: Use Playwright for full scraping
    slug = url.split("/")[-1]

    article = {
        "url": url,
        "slug": slug,
        "title": "",  # Will be filled by Claude
        "body": "",   # User should provide or we scrape
    }

    # Save article info
    article_path = task_dir / "article.json"
    with open(article_path, "w") as f:
        json.dump(article, f, ensure_ascii=False, indent=2)

    return article


def analyze_article(
    article_text: str,
    api_key: str,
    task_dir: Path,
    logger,
    visual_style: str = None,
    video_kind: str = None,
    user_feedback: str = None
) -> Dict[str, Any]:
    """Use Claude to analyze article and create script structure."""
    logger.info("  Sending article to Claude for analysis...")

    # Get style context if specified
    if visual_style and video_kind:
        style_context = get_style_prompt_context(visual_style, video_kind)
        style_hints = get_frame_style_hints(visual_style)
    else:
        # List available styles for Claude to choose
        style_context = f"""## Visual Style Selection
Choose an appropriate visual style and video kind based on the article content.

Available styles: {', '.join(VISUAL_STYLES.keys())}
Available kinds: {', '.join(VIDEO_KINDS.keys())}

Output your selection in the JSON with "visual_style" and "video_kind" fields."""
        style_hints = {}

    # User feedback section
    feedback_section = ""
    if user_feedback:
        feedback_section = f"""
## USER FEEDBACK (CRITICAL - FOLLOW THIS!)
{user_feedback}

IMPORTANT: The user's feedback takes PRIORITY over default structure!
"""

    # Check for new v2 scene format flag (experimental)
    use_v2_format = os.environ.get("MICRODRAMA_V2", "0") == "1"

    if use_v2_format:
        prompt = f"""Analyze this article and create a FUNNY, EMOTIONAL video script.

ARTICLE:
{article_text}

{style_context}
{feedback_section}

Create a JSON script for a 45-60 second video. KEY: Make it FUNNY and EMOTIONALLY engaging!

## NEW SCENE FORMAT (V2)
Each scene = ONE MOMENT with ONE FRAME. More scenes, shorter duration each.

REQUIREMENTS:

1. TONE: COMEDY and EMOTION are priority!
   - Exaggerate reactions, facial expressions
   - Use comedic timing - pauses before punchlines
   - Build emotional tension before release
   - Let the absurdity speak for itself

2. SCENES: Create 10-14 SHORT scenes (3-5 seconds each)
   - Each scene = one visual moment with one caption
   - Total video length: 45-60 seconds
   - Transitions between scenes create the flow

3. SCENE FORMAT (NEW):
   - "caption": Single caption for this moment (Russian)
   - "duration": 3-5 seconds
   - "transition_to_next": "cut" | "dissolve" | "zoom_in" | "zoom_out" | "morph"
   - "frame": Single frame description (what we see at this moment)
   - "emotion": Dominant emotion (joy, fear, confusion, triumph, etc.)

4. STORYTELLER: The person telling the story - make them EXPRESSIVE!
   - Strong emotional reactions
   - Facing camera/audience (NEVER back to camera!)
   - Gestures, facial expressions described

5. CHARACTERS: Full person descriptions with orientation
   - ALWAYS specify which way they face
   - Include emotional state and body language
   - "facing camera", "looking at audience", etc.

6. LOCATIONS: Include people in settings
   - "lecture hall with attentive audience"
   - "launch pad with busy engineers"

7. NO CHILDREN - all adults 25-45 years old

8. CAPTIONS in natural Russian, no anglicisms

Output JSON:
{{
  "visual_style": "documentary",
  "video_kind": "comedy",
  "format_version": "v2",
  "storyteller": {{
    "type": "expert|witness|victim|narrator",
    "description": "Full person description with typical expression",
    "reference_query": "Yandex search query",
    "is_real_person": true/false,
    "is_cartoon": false
  }},
  "characters": {{
    "char_id": {{
      "description": "Full person, clothing, typical expression",
      "reference_query": "Real name + фото портрет",
      "is_real_person": true/false,
      "is_cartoon": false
    }}
  }},
  "objects": {{ ... }},
  "locations": {{
    "loc_id": {{
      "description": "Location WITH people present",
      "reference_query": "..."
    }}
  }},
  "scenes": [
    {{
      "scene_id": 1,
      "type": "hook",
      "method": "veo3",
      "caption": "Космонавта спросили: в чём подвиг?",
      "duration": 4,
      "transition_to_next": "cut",
      "characters": ["lebedev"],
      "objects": [],
      "locations": ["lecture_hall"],
      "frame": "lebedev at podium, facing audience, eyebrows raised in surprise, audience visible behind",
      "emotion": "amused surprise",
      "sound_hints": "audience murmur"
    }},
    {{
      "scene_id": 2,
      "type": "setup",
      "method": "veo3",
      "caption": "Вы видели 10-этажный дом?",
      "duration": 3,
      "transition_to_next": "zoom_out",
      "characters": ["lebedev"],
      "objects": [],
      "locations": ["lecture_hall"],
      "frame": "lebedev gesturing upward with both hands, knowing smile, audience nodding",
      "emotion": "playful teaching",
      "sound_hints": "thoughtful hmm"
    }}
  ]
}}"""
    else:
        prompt = f"""Analyze this satirical news article and create a video script structure.

ARTICLE:
{article_text}

{style_context}
{feedback_section}

Create a JSON script for a 45-60 second satirical parody video with these requirements:

1. NARRATOR (storyteller): Choose WHO tells the story - can be creative!
   Options:
   - News anchor (classic approach)
   - The victim/protagonist themselves (e.g., if story is about Cheburashka - HE can narrate!)
   - A witness to events
   - An expert/academic
   - Something unexpected and fun

   The narrator appears in opening (hook) and closing (punchline) scenes with emotional reactions.

2. CHARACTERS: List all characters with BOTH descriptions:
   - Use SEMANTIC character IDs that describe the role (e.g., "son", "father", "historian")
   - NOT generic IDs like "teenage_boy" or "man_1" - use story-relevant names!
   - "description": Generic visual for AI generation (no real names)
   - "reference_query": Yandex search with REAL NAME for reference photos
   - "is_real_person": true/false (CRITICAL for scene method selection!)
   - "is_cartoon": true/false (for animated characters like Cheburashka)

   Example for Dugin:
   - description: "Elderly bearded Russian philosopher"
   - reference_query: "Александр Дугин фото портрет"
   - is_real_person: true

   Example for Cheburashka:
   - description: "Brown furry creature with big ears"
   - reference_query: "Чебурашка мультфильм 2023"
   - is_real_person: false
   - is_cartoon: true

3. LOCATIONS: Describe settings INCLUDING background people!
   - Locations should include typical people present (audience, crowd, workers)
   - These are NOT story characters - they're part of the setting
   - Example: "lecture_hall" → "University lecture hall filled with adult students in seats"
   - Example: "launch_pad" → "Baikonur launch pad with engineers and technicians working"
   - This ensures frames show realistic populated settings, not empty rooms!

4. SCENE METHOD: ALWAYS use method: "veo3"
   - ALL scenes use veo3 method - Kling will handle famous people automatically
   - Every scene MUST have frames.first AND frames.last (DIFFERENT descriptions!)
   - No kenburns method - it's disabled

5. ABSOLUTE RESTRICTIONS:
   - NO CHILDREN! All people must be adults (25-45 years old)
   - Never mention children, kids, babies in any frame description
   - Use "visitors", "crowd", "audience" instead of "children"
   - Veo3 will REJECT frames containing minors!

6. SCENES: Create 5-7 scenes following this structure:
   - Scene 1: HOOK - Narrator discovers/presents news
   - Scene 2-3: SETUP - Establish situation
   - Scene 4-5: DEVELOPMENT - Build absurdity
   - Scene 6: CLIMAX - Peak absurdity
   - Scene 7: PUNCHLINE - Narrator reaction

7. For EVERY scene specify:
   - method: "veo3" (ALWAYS - no other methods!)
   - frames.first AND frames.last (MUST be DIFFERENT! Show progression!)
   - caption_first AND caption_last (DIFFERENT text for each!)
   - CAPTIONS (CRITICAL for storytelling!):
     * ALWAYS use "caption_first" and "caption_last" - different text!
     * Captions tell the story - viewer should understand without audio!
     * Each caption should advance the narrative
   - sound_hints (for veo3 scenes only)
   - ACTION (CRITICAL for veo3 scenes!):
     * Specify PHYSICAL MOVEMENT that happens during the scene - NOT just expression changes!
     * 8 seconds is LONG - need substantial body movement, not just "smile appears"
     * BAD: "Expression changes from neutral to sad" - TOO WEAK for 8 seconds!
     * GOOD: "Son walks across room to father, hands over the book, steps back"
     * GOOD: "Father stands up from chair, walks to bookshelf, pulls out book, turns around"
     * GOOD: "Son enters room from doorway, approaches father at desk, sits down next to him"
     * Characters should MOVE through space: walk, turn, sit, stand, reach, hand over
     * First frame shows STARTING position (person far if approaching, standing if will sit)
   - OBJECT_PLACEMENT (when scene has objects):
     * For EACH frame, specify exactly WHERE each object is and WHO holds it
     * This prevents objects appearing in wrong hands/places
     * Use "object_placement" field with "first" and "last" sub-objects
     * Example: first: {{"book": "in son's hands"}}, last: {{"book": "now in father's hands"}}

8. WHERE TO USE CHARACTER IDs vs NORMAL WORDS:

   CHARACTER IDs (like "son", "father", "book") are for TECHNICAL fields:
   - frames.first / frames.last - for image generator to match references
   - object_placement - for specifying object positions

   NORMAL WORDS are for HUMAN-READABLE fields:
   - "action" field - English, normal words: "Son enters room, approaches father at bookshelf"
   - "caption_first" / "caption_last" - Russian for viewers: "Отец протягивает книгу"

   WRONG: "teenage_son входит в комнату, подходит к father" (mixed!)
   CORRECT action: "Son enters room, walks to father at bookshelf, stops beside him"
   CORRECT caption: "Сын подходит к отцу"

9. FRAME DESCRIPTIONS (use character IDs here!):
   - In frames.first / frames.last, use EXACT character/object IDs for reference matching
   - Example: "father standing at bookshelf, son entering from doorway far in background"
   - Example: "son holding book, extending it toward father"
   - STARTING POSITIONS: First frame must show LOGICAL starting position:
     * If action is "approaches" → first frame: character FAR from target
     * If action is "sits down" → first frame: character STANDING
     * If action is "hands over" → first frame: character holding object

   CHARACTER ORIENTATION - CRITICAL!
   - ALWAYS specify which way character is FACING: "facing camera", "facing audience", "profile view"
   - For speakers/presenters: MUST face audience/camera, NOT turn back!
   - BAD: "lebedev at podium" (unclear orientation - might face wrong way!)
   - GOOD: "lebedev at podium, facing audience, addressing the crowd"
   - GOOD: "lebedev turned slightly toward camera, making eye contact"
   - When character addresses audience → "facing audience", "looking at audience"
   - When character is filmed → "facing camera", "looking into camera"
   - Only use "back to camera" when dramatically appropriate!

   FULL PERSON DESCRIPTIONS:
   - Describe WHOLE PERSON, not just parts: "cosmonaut in spacesuit" NOT "spacesuit helmet"
   - Include body, posture, clothing: "young man in white spacesuit, confident stance"
   - BAD: "helmet visible" - confusing, might generate just a helmet!
   - GOOD: "cosmonaut visible through capsule window, face determined"

10. CAPTIONS (Russian for viewers!):
   - ALL captions in natural Russian
   - NO character IDs in captions - use normal words (сын, отец, книга)
   - Captions tell the story - viewer should understand without audio
   - NO ANGLICISMS! Use proper Russian words:
     * плакаты (NOT баннеры)
     * контент → содержание/материал
     * топ → лучший/вершина
     * лайк → одобрение
     * Use native Russian equivalents wherever possible

   BANNER/PLACARD TEXT - CRITICAL!
   - If scene shows people holding banners with text, the EXACT text must be specified!
   - Put banner text in quotes in frame description: holding placards with "Юра, мы всё посчитали!"
   - Text on banners should match the STORY, not be random slogans!
   - Example: Scientists holding "Юра, ты вернёшься!" (matching the story's quote)

11. VISUAL STYLE ADAPTATION:
   Frame descriptions should match the visual style:
   - For "melodrama": Include emotional lighting, dramatic close-ups, exaggerated expressions
   - For "soviet_cartoon": Reference classic Soyuzmultfilm aesthetic, soft colors, nostalgic charm
   - For "documentary": Natural lighting, observational style, authentic moments
   - For "news_broadcast": Studio lighting, professional framing, news anchor poses
   - For "thriller": Shadows, tension, suspenseful atmosphere
   Each frame description should explicitly mention the visual style elements!

Output ONLY valid JSON matching this structure:
{{
  "visual_style": "style_key from available styles",
  "video_kind": "kind_key from available kinds",
  "storyteller": {{
    "type": "news_anchor|victim|witness|expert|unexpected",
    "description": "Visual description for AI generation",
    "reference_query": "Yandex search query for reference",
    "is_real_person": false,
    "is_cartoon": false
  }},
  "characters": {{
    "char_id": {{
      "description": "FULL PERSON description: age, clothing, posture, expression - NOT just face/head!",
      "reference_query": "REAL NAME + фото портрет",
      "is_real_person": true,
      "is_cartoon": false
    }}
  }},
  "objects": {{
    "obj_id": {{
      "description": "...",
      "reference_query": "Yandex search query"
    }}
  }},
  "locations": {{
    "loc_id": {{
      "description": "Location WITH typical people present (audience, workers, crowd)",
      "reference_query": "Yandex search query for location reference"
    }}
  }},
  "scenes": [
    {{
      "scene_id": 1,
      "type": "hook",
      "method": "veo3",
      "caption_first": "СРОЧНО: Начало истории",
      "caption_last": "...продолжение мысли!",
      "characters": ["storyteller"],
      "objects": [],
      "locations": ["studio"],
      "action": "Narrator leans forward in chair, picks up paper from desk, looks at it, then looks up at camera with knowing expression",
      "frames": {{
        "first": "Narrator sitting back in chair, arms on armrests, looking thoughtfully at camera",
        "last": "Narrator leaned forward, holding paper, direct eye contact with camera, slight smile"
      }},
      "sound_hints": "paper rustle, gasp"
    }},
    {{
      "scene_id": 2,
      "type": "setup",
      "method": "veo3",
      "caption_first": "Начало действия",
      "caption_last": "Действие завершено",
      "characters": ["person_a", "person_b"],
      "objects": ["important_object"],
      "locations": ["location_id"],
      "action": "Person A hands the object to Person B who accepts it",
      "object_placement": {{
        "first": {{"important_object": "in Person A's hands, being offered"}},
        "last": {{"important_object": "now in Person B's hands, received"}}
      }},
      "frames": {{
        "first": "Person A extending object toward Person B, anticipation",
        "last": "Person B now holding object, Person A stepping back"
      }},
      "sound_hints": "object handling sounds"
    }}
  ]
}}"""

    # Save prompt for debugging
    prompts_dir = task_dir / "prompts"
    prompts_dir.mkdir(parents=True, exist_ok=True)
    with open(prompts_dir / "stage2_analysis.txt", "w") as f:
        f.write(prompt)

    response = call_claude(prompt, api_key)

    # Save response
    with open(prompts_dir / "stage2_response.txt", "w") as f:
        f.write(response)

    # Extract JSON from response
    try:
        # Find JSON in response
        json_start = response.find("{")
        json_end = response.rfind("}") + 1
        if json_start >= 0 and json_end > json_start:
            script_data = json.loads(response[json_start:json_end])
        else:
            raise ValueError("No JSON found in response")
    except json.JSONDecodeError as e:
        logger.error(f"Failed to parse Claude response as JSON: {e}")
        logger.error(f"Response: {response[:500]}")
        raise

    # Save analysis
    analysis_path = task_dir / "analysis" / "script_draft.json"
    analysis_path.parent.mkdir(parents=True, exist_ok=True)
    with open(analysis_path, "w") as f:
        json.dump(script_data, f, ensure_ascii=False, indent=2)

    logger.info(f"Script draft saved: {analysis_path}")
    return script_data


def search_yandex_images(query: str, output_dir: Path, prefix: str, count: int = 3, logger=None) -> List[Path]:
    """Search Yandex and download reference images."""
    if logger:
        logger.info(f"Searching Yandex: {query}")

    cmd = [
        sys.executable,
        str(SCRIPTS_DIR / "yandex_image_search.py"),
        query,
        "-o", str(output_dir),
        "-n", str(count),
        "--size", "large",
        "--download",
        "-p", prefix
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)

        if result.returncode != 0:
            if logger:
                logger.warning(f"  Yandex search failed: {result.stderr[:200]}")
            return []
    except subprocess.TimeoutExpired:
        if logger:
            logger.warning(f"  Yandex search timed out for: {prefix}")
        return []
    except Exception as e:
        if logger:
            logger.warning(f"  Yandex search error: {e}")
        return []

    # Find downloaded files
    downloaded = list(output_dir.glob(f"{prefix}_*.jpg")) + list(output_dir.glob(f"{prefix}_*.webp")) + list(output_dir.glob(f"{prefix}_*.png"))
    return sorted(downloaded)


def select_best_reference(images: List[Path], description: str, api_key: str, logger) -> Optional[Path]:
    """Use Claude to select the best reference image."""
    if not images:
        return None

    if len(images) == 1:
        return images[0]

    logger.info(f"Selecting best reference from {len(images)} images...")

    # Prepare images for Claude
    image_data = []
    for img_path in images[:5]:  # Max 5 images
        with open(img_path, "rb") as f:
            img_bytes = f.read()
        img_b64 = base64.b64encode(img_bytes).decode()
        mime = "image/jpeg"
        if img_path.suffix.lower() == ".png":
            mime = "image/png"
        elif img_path.suffix.lower() == ".webp":
            mime = "image/webp"
        image_data.append({"base64": img_b64, "mime_type": mime, "path": str(img_path)})

    prompt = f"""I need to select the best reference image for this character/object:

DESCRIPTION: {description}

Look at the {len(image_data)} images provided and select the ONE that best matches the description.

Consider:
- How well it matches the description
- Image quality and clarity
- Suitability for news parody video style

Respond with ONLY the number (1, 2, 3, etc.) of the best image."""

    try:
        response = call_claude(prompt, api_key, images=image_data)
        # Extract number from response
        for word in response.split():
            if word.isdigit():
                idx = int(word) - 1
                if 0 <= idx < len(images):
                    logger.info(f"Selected image {idx + 1}: {images[idx].name}")
                    return images[idx]
    except Exception as e:
        logger.warning(f"Failed to select with Claude: {e}")

    # Default to first image
    return images[0]


def gather_references(script_data: Dict, task_dir: Path, api_key: str, logger) -> Dict[str, Path]:
    """Search Yandex for all needed references and select best ones."""
    refs_dir = task_dir / "references" / "search"
    refs_dir.mkdir(parents=True, exist_ok=True)

    selected_refs = {}

    # Search for storyteller
    storyteller = script_data.get("storyteller", {})
    if storyteller.get("reference_query"):
        images = search_yandex_images(
            storyteller["reference_query"],
            refs_dir,
            "storyteller",
            logger=logger
        )
        best = select_best_reference(images, storyteller.get("description", ""), api_key, logger)
        if best:
            selected_refs["storyteller"] = best

    # Search for characters
    for char_id, char_info in script_data.get("characters", {}).items():
        if char_info.get("reference_query"):
            images = search_yandex_images(
                char_info["reference_query"],
                refs_dir,
                char_id,
                logger=logger
            )
            best = select_best_reference(images, char_info.get("description", ""), api_key, logger)
            if best:
                selected_refs[char_id] = best

    # Search for objects
    for obj_id, obj_info in script_data.get("objects", {}).items():
        if obj_info.get("reference_query"):
            images = search_yandex_images(
                obj_info["reference_query"],
                refs_dir,
                obj_id,
                logger=logger
            )
            best = select_best_reference(images, obj_info.get("description", ""), api_key, logger)
            if best:
                selected_refs[obj_id] = best

    # Search for locations
    for loc_id, loc_info in script_data.get("locations", {}).items():
        if loc_info.get("reference_query"):
            images = search_yandex_images(
                loc_info["reference_query"],
                refs_dir,
                loc_id,
                logger=logger
            )
            best = select_best_reference(images, loc_info.get("description", ""), api_key, logger)
            if best:
                selected_refs[loc_id] = best

    # Save reference mapping
    refs_mapping = {k: str(v) for k, v in selected_refs.items()}
    with open(task_dir / "references" / "selected_refs.json", "w") as f:
        json.dump(refs_mapping, f, indent=2)

    logger.info(f"Selected {len(selected_refs)} reference images")
    return selected_refs


def build_character_sheet_prompt(char_info: Dict, has_reference: bool) -> str:
    """Build personalized character sheet prompt based on character info."""
    description = char_info.get("description", "character")
    is_cartoon = char_info.get("is_cartoon", False)
    sheet_hints = char_info.get("sheet_hints", "")  # Custom hints from script (clothes, props, etc.)

    # Style note based on character type
    if is_cartoon:
        style_note = "matching the cartoon/animated style of the reference"
    else:
        style_note = "in a photorealistic style matching the reference"

    if has_reference:
        prompt = f"""Create a professional character reference sheet based strictly on the uploaded
reference image. Use a clean, neutral plain background and present the sheet
as a technical model turnaround {style_note}.

CHARACTER: {description}
"""
    else:
        prompt = f"""Create a professional character reference sheet of: {description}
Use a clean, neutral plain background and present the sheet as a technical model
turnaround in a photographic style.

"""

    # Add custom hints if provided (clothes, accessories, props, etc.)
    if sheet_hints:
        prompt += f"""
SPECIAL REQUIREMENTS:
{sheet_hints}

"""

    # Standard layout instructions
    prompt += """Arrange the composition into two horizontal rows.
Top row: four full-body standing views placed side-by-side in this order:
front view, left profile view (facing left), right profile view (facing right), back view.
Bottom row: three highly detailed close-up portraits aligned beneath the full-body row
in this order: front portrait, left profile portrait (facing left), right profile portrait (facing right).

Maintain perfect identity consistency across every panel. Keep the subject in a
relaxed A-pose and with consistent scale and alignment between views, accurate
anatomy, and clear silhouette; ensure even spacing and clean panel separation,
with uniform framing and consistent head height across the full-body lineup and
consistent facial scale across the portraits. Lighting should be consistent across
all panels (same direction, intensity, and softness), with natural, controlled
shadows that preserve detail without dramatic mood shifts. Output a crisp,
print-ready reference sheet look, sharp details."""

    return prompt


def build_object_sheet_prompt(obj_info: Dict) -> str:
    """Build personalized object sheet prompt based on object info."""
    description = obj_info.get("description", "object")
    sheet_hints = obj_info.get("sheet_hints", "")  # Custom hints (specific details to show)

    prompt = f"""Create a professional object reference sheet based on the uploaded reference image.
Use a clean, neutral plain background and present the sheet as a technical product turnaround.

OBJECT: {description}
"""

    if sheet_hints:
        prompt += f"""
SPECIAL REQUIREMENTS:
{sheet_hints}

"""

    prompt += """Arrange the composition showing the object from multiple angles:
- Main view (largest, centered) showing the object clearly
- Three smaller views around it showing: front angle, side angle, detail close-up
- If applicable, show any text, labels, or important details in a separate close-up panel

Maintain perfect consistency across all views - same object, same condition, same textures.
Use even studio lighting across all panels. Output a crisp, detailed reference sheet."""

    return prompt


def build_location_sheet_prompt(loc_info: Dict) -> str:
    """Build personalized location sheet prompt based on location info."""
    description = loc_info.get("description", "location")
    sheet_hints = loc_info.get("sheet_hints", "")

    # Check if location description mentions people (audience, students, engineers, etc.)
    people_keywords = ["audience", "students", "crowd", "people", "engineers", "technicians",
                       "workers", "visitors", "staff", "seated", "filled with", "academics",
                       "аудитория", "студенты", "толпа", "инженеры", "рабочие", "зрители",
                       "слушатели", "посетители", "сотрудники"]
    has_people = any(kw.lower() in description.lower() for kw in people_keywords)

    prompt = f"""Create a reference sheet showing this location from multiple angles:

LOCATION: {description}
"""

    if sheet_hints:
        prompt += f"""
SPECIAL REQUIREMENTS:
{sheet_hints}

"""

    prompt += """Layout:
- Main view (large, top): Wide establishing shot of the location
- Three smaller views below: Different angles/details of the same location

"""

    if has_people:
        prompt += """IMPORTANT: This location includes people as part of the setting!
Include the described people (audience, students, workers, etc.) in ALL views.
These are background/setting people, not main characters - show them naturally
populating the space. The location should feel alive and realistic.
"""
    else:
        prompt += """CRITICAL: This is a LOCATION reference - NO PEOPLE in the image!
Show only the environment, architecture, furniture, lighting, and atmosphere.
"""

    prompt += """Use a clean, neutral presentation style suitable for video production reference.
Maintain consistent lighting and color palette across all views."""

    return prompt


def generate_character_sheets(
    script_data: Dict,
    selected_refs: Dict[str, Path],
    task_dir: Path,
    api_key: str,
    logger
) -> Dict[str, Path]:
    """
    Generate character reference sheets for consistency across scenes.

    Uses NanoBanana to create turnaround sheets from Yandex reference images.
    Sheets provide multiple angles and expressions for better consistency.

    ALL characters get sheets (including real people) because:
    - We generate frames with NanoBanana for all characters
    - Sheets ensure consistency across scenes
    - Later we may use Kling for video which handles real people
    """
    sheets_dir = task_dir / "references" / "sheets"
    prompts_dir = task_dir / "references" / "prompts"
    sheets_dir.mkdir(parents=True, exist_ok=True)
    prompts_dir.mkdir(parents=True, exist_ok=True)

    character_sheets = {}

    def generate_single_sheet(entity_id: str, entity_info: Dict, is_object: bool = False, is_location: bool = False) -> Optional[Path]:
        """Generate a single character, object, or location sheet with personalized prompt."""
        sheet_path = sheets_dir / f"sheet_{entity_id}.jpg"
        prompt_path = prompts_dir / f"sheet_{entity_id}_prompt.txt"

        if sheet_path.exists():
            logger.info(f"  Sheet exists: {sheet_path.name}")
            return sheet_path

        ref_path = selected_refs.get(entity_id)
        ref_data = None
        ref_mime = None
        has_reference = False

        if ref_path and ref_path.exists():
            with open(ref_path, "rb") as f:
                ref_data = f.read()
            ref_mime = "image/jpeg"
            if ref_path.suffix.lower() == ".png":
                ref_mime = "image/png"
            elif ref_path.suffix.lower() == ".webp":
                ref_mime = "image/webp"
            has_reference = True

        # Build personalized prompt based on entity type and info
        if is_location:
            prompt = build_location_sheet_prompt(entity_info)
            logger.info(f"  Generating location sheet: {entity_id}")
        elif is_object:
            prompt = build_object_sheet_prompt(entity_info)
            logger.info(f"  Generating object sheet: {entity_id}")
        else:
            prompt = build_character_sheet_prompt(entity_info, has_reference)
            logger.info(f"  Generating character sheet: {entity_id}")

        # Save prompt for debugging/iteration
        with open(prompt_path, "w") as f:
            f.write(prompt)
            if ref_path:
                f.write(f"\n\n# Reference image: {ref_path}")

        try:
            sheet_bytes, sheet_mime = call_nanobanana_laozhang(
                prompt=prompt,
                image_data=ref_data,
                mime_type=ref_mime or "image/jpeg",
                api_key=api_key,
                aspect_ratio="16:9"  # Wide for turnaround layout
            )

            if sheet_bytes:
                with open(sheet_path, "wb") as f:
                    f.write(sheet_bytes)
                logger.info(f"    Saved: {sheet_path.name}")
                return sheet_path
            else:
                logger.warning(f"    Failed to generate sheet for {entity_id}")
                return ref_path  # Fall back to original reference
        except Exception as e:
            logger.warning(f"    Error generating sheet for {entity_id}: {e}")
            return ref_path  # Fall back to original reference

    # Process storyteller (ALL storytellers get sheets now)
    storyteller = script_data.get("storyteller", {})
    if storyteller:
        result = generate_single_sheet("storyteller", storyteller)
        if result:
            character_sheets["storyteller"] = result

    # Process ALL characters (including real people - they need sheets for frame generation)
    for char_id, char_info in script_data.get("characters", {}).items():
        result = generate_single_sheet(char_id, char_info)
        if result:
            character_sheets[char_id] = result

    # Process objects
    for obj_id, obj_info in script_data.get("objects", {}).items():
        result = generate_single_sheet(obj_id, obj_info, is_object=True)
        if result:
            character_sheets[obj_id] = result

    # Process locations
    for loc_id, loc_info in script_data.get("locations", {}).items():
        result = generate_single_sheet(loc_id, loc_info, is_location=True)
        if result:
            character_sheets[loc_id] = result

    # Save sheet mapping
    sheets_mapping = {k: str(v) for k, v in character_sheets.items()}
    with open(sheets_dir / "character_sheets.json", "w") as f:
        json.dump(sheets_mapping, f, indent=2)

    logger.info(f"  Generated {len([p for p in character_sheets.values() if 'sheet_' in p.name])} character sheets")
    return character_sheets


def finalize_script(script_data: Dict, selected_refs: Dict[str, Path], task_dir: Path, logger) -> Dict:
    """Finalize script with reference file paths."""
    # Add reference paths to script
    if "storyteller" in selected_refs:
        script_data["storyteller"]["reference"] = selected_refs["storyteller"].name

    for char_id in script_data.get("characters", {}):
        if char_id in selected_refs:
            script_data["characters"][char_id]["reference"] = selected_refs[char_id].name

    for obj_id in script_data.get("objects", {}):
        if obj_id in selected_refs:
            script_data["objects"][obj_id]["reference"] = selected_refs[obj_id].name

    for loc_id in script_data.get("locations", {}):
        if loc_id in selected_refs:
            script_data["locations"][loc_id]["reference"] = selected_refs[loc_id].name

    # Add metadata
    script_data["task_id"] = task_dir.name
    script_data["aspect_ratio"] = "9:16"
    # Keep visual_style and video_kind from Claude's analysis (or use defaults)
    if "visual_style" not in script_data:
        script_data["visual_style"] = "documentary"
    if "video_kind" not in script_data:
        script_data["video_kind"] = "satire_commentary"

    # Save final script
    script_dir = task_dir / "script"
    script_dir.mkdir(parents=True, exist_ok=True)
    script_path = script_dir / "script_v5.json"

    with open(script_path, "w") as f:
        json.dump(script_data, f, ensure_ascii=False, indent=2)

    logger.info(f"Final script saved: {script_path}")
    return script_data


def generate_frames(task_dir: Path, logger) -> bool:
    """Generate all frames using generate_all_frames.py."""
    logger.info("Generating frames...")

    cmd = [
        sys.executable,
        str(SCRIPTS_DIR / "generate_all_frames.py"),
        str(task_dir)
    ]

    result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)

    if result.returncode != 0:
        logger.error(f"Frame generation failed: {result.stderr[-500:]}")
        return False

    logger.info("Frame generation complete")
    return True


def generate_videos(task_dir: Path, logger) -> bool:
    """Generate Veo3 videos and Ken Burns clips."""
    logger.info("Generating video clips...")

    # Generate Veo3 clips
    veo3_cmd = [
        sys.executable,
        str(SCRIPTS_DIR / "generate_scene_videos.py"),
        str(task_dir),
        "--skip-existing"
    ]

    result = subprocess.run(veo3_cmd, capture_output=True, text=True, timeout=1800)
    if result.returncode != 0:
        logger.warning(f"Some Veo3 clips may have failed: {result.stderr[-300:]}")

    # DISABLED: Ken Burns - now using Kling for all scenes including famous people
    # kb_cmd = [
    #     sys.executable,
    #     str(SCRIPTS_DIR / "generate_kenburns.py"),
    #     str(task_dir)
    # ]
    #
    # result = subprocess.run(kb_cmd, capture_output=True, text=True, timeout=300)
    # if result.returncode != 0:
    #     logger.warning(f"Some Ken Burns clips may have failed: {result.stderr[-300:]}")

    return True


def assemble_video(task_dir: Path, logger) -> Optional[Path]:
    """Assemble final video."""
    logger.info("Assembling final video...")

    cmd = [
        sys.executable,
        str(SCRIPTS_DIR / "assemble_video.py"),
        str(task_dir),
        "--mode", "fade"
    ]

    result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)

    if result.returncode != 0:
        logger.error(f"Assembly failed: {result.stderr[-500:]}")
        return None

    output_path = task_dir / "assembly" / "final_video.mp4"
    if output_path.exists():
        logger.info(f"Final video: {output_path}")
        return output_path

    return None


def run_pipeline(
    url: str,
    article_text: str = None,
    task_dir: Path = None,
    visual_style: str = None,
    video_kind: str = None,
    user_feedback: str = None,
    interactive: bool = False,
    slug: str = None,
    stop_after: str = "all",
    start_from: str = None
):
    """
    Run the full pipeline with optional stage control.

    Stages in order: script → refs → sheets → frames → videos → assembly

    Args:
        stop_after: Run up to this stage and stop ("script", "refs", "sheets", "frames", "videos", "assembly", "all")
        start_from: Resume from this stage (requires existing task_dir with previous stage outputs)
    """
    STAGE_ORDER = ["script", "refs", "sheets", "frames", "videos", "assembly"]

    def should_run_stage(stage: str) -> bool:
        """Check if a stage should run based on start_from and stop_after."""
        stage_idx = STAGE_ORDER.index(stage)

        # Check start_from
        if start_from:
            start_idx = STAGE_ORDER.index(start_from)
            if stage_idx < start_idx:
                return False

        # Check stop_after
        if stop_after != "all":
            stop_idx = STAGE_ORDER.index(stop_after)
            if stage_idx > stop_idx:
                return False

        return True
    api_key = os.getenv("LAOZHANG_KEY")
    if not api_key:
        print("Error: Set LAOZHANG_KEY environment variable")
        sys.exit(1)

    # Determine story slug
    if slug:
        story_slug = slug
    elif url:
        story_slug = url.split("/")[-1]
    else:
        story_slug = "manual"

    # Create story directory (shared between all tasks for this story)
    story_dir = BASE_DIR / story_slug
    story_dir.mkdir(parents=True, exist_ok=True)

    # Save source text in story directory if provided (shared between tasks)
    source_file = story_dir / "source.md"
    if article_text and not source_file.exists():
        with open(source_file, "w") as f:
            f.write(article_text)
        print(f"Saved source text: {source_file}")
    elif article_text is None and source_file.exists():
        # Load from existing source file
        with open(source_file) as f:
            article_text = f.read()
        print(f"Loaded source text from: {source_file}")

    # Create task directory
    if task_dir is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        task_id = f"{timestamp}_{story_slug[:20]}"
        task_dir = story_dir / task_id
    else:
        task_id = task_dir.name

    task_dir.mkdir(parents=True, exist_ok=True)

    # Create subdirectories
    for subdir in ["script", "references/search", "references/sheets", "analysis", "assembly"]:
        (task_dir / subdir).mkdir(parents=True, exist_ok=True)

    # Setup logging
    logger = setup_logging(task_dir, task_id)

    # Generation log data
    gen_log = {
        "task_id": task_id,
        "url": url,
        "started_at": datetime.now().isoformat(),
        "stages": {}
    }

    logger.info("=" * 60)
    logger.info("MICRODRAMA PIPELINE")
    logger.info(f"Task: {task_id}")
    logger.info(f"URL: {url or 'manual input'}")
    logger.info(f"Directory: {task_dir}")
    logger.info("=" * 60)

    try:
        # Stage 1: Fetch/load article
        logger.info("\n[STAGE 1] Loading article...")
        stage_start = time.time()

        if article_text:
            article = {"url": url, "body": article_text}
            with open(task_dir / "article.md", "w") as f:
                f.write(article_text)
            logger.info(f"  Article loaded from text ({len(article_text)} chars)")
        else:
            article = fetch_article(url, task_dir, logger)
            logger.warning("  Article fetching not fully implemented - provide article text with --article")

        gen_log["stages"]["stage1_article"] = {
            "duration_sec": time.time() - stage_start,
            "article_length": len(article.get("body", ""))
        }

        # INTERACTIVE PHASE 1: Understanding + Choices (file-based)
        if interactive and article.get("body"):
            questions_file = task_dir / "interactive_questions.md"
            understanding_file = task_dir / "analysis" / "understanding.json"

            # Check if we already have answers
            if questions_file.exists() and understanding_file.exists():
                # Load existing understanding and answers
                logger.info("\n[INTERACTIVE] Loading your answers...")
                with open(understanding_file) as f:
                    understanding = json.load(f)

                phase1_choices = load_interactive_answers(task_dir, understanding)

                if phase1_choices:
                    # Apply user corrections as feedback
                    if phase1_choices.get("correction"):
                        if user_feedback:
                            user_feedback += f"\n\nUSER CORRECTION: {phase1_choices['correction']}"
                        else:
                            user_feedback = f"USER CORRECTION: {phase1_choices['correction']}"

                    # Apply storyteller choice
                    storyteller_choice = phase1_choices.get("storyteller", {})
                    if storyteller_choice:
                        if user_feedback:
                            user_feedback += f"\n\nSTORYTELLER: {storyteller_choice.get('type', 'custom')}: {storyteller_choice.get('description', '')}"
                        else:
                            user_feedback = f"STORYTELLER: {storyteller_choice.get('type', 'custom')}: {storyteller_choice.get('description', '')}"

                    # Apply style choice
                    if phase1_choices.get("style"):
                        visual_style = phase1_choices["style"]

                    # Apply visual concept (most important!)
                    if phase1_choices.get("visual_concept"):
                        if user_feedback:
                            user_feedback += f"\n\nVISUAL CONCEPT (IMPORTANT - follow this structure!):\n{phase1_choices['visual_concept']}"
                        else:
                            user_feedback = f"VISUAL CONCEPT (IMPORTANT - follow this structure!):\n{phase1_choices['visual_concept']}"
                        logger.info(f"  Visual concept: {phase1_choices['visual_concept'][:100]}...")

                    gen_log["interactive_phase1"] = {
                        "understanding": understanding,
                        "choices": phase1_choices
                    }

                    logger.info(f"  Loaded choices: storyteller={storyteller_choice.get('type')}, style={visual_style}")
                else:
                    logger.warning("  No answers found in questions file - using defaults")
            else:
                # Generate understanding and save questions
                logger.info("\n[INTERACTIVE] Phase 1: Analyzing the story...")
                understanding = analyze_story_for_understanding(article["body"], api_key)

                # Save understanding
                with open(understanding_file, "w") as f:
                    json.dump(understanding, f, ensure_ascii=False, indent=2)

                # Save questions file
                questions_path = save_interactive_questions(understanding, list(VISUAL_STYLES.keys()), task_dir)

                # Stop and wait for user answers
                logger.info("\n" + "=" * 60)
                logger.info("INTERACTIVE: Questions saved!")
                logger.info(f"Edit your answers in: {questions_path}")
                logger.info("")
                logger.info("After answering, continue with:")
                logger.info(f"  python microdrama_pipeline.py --task-dir {task_dir} --from-stage refs -i")
                logger.info("=" * 60)

                save_generation_log(task_dir, gen_log)
                return task_dir

        # Stage 2: Analyze and create script
        # Determine if we should run Stage 2:
        # - Always run if not resuming (start_from is None)
        # - Run if resuming but have user_feedback (from interactive answers)
        # - Skip if resuming without feedback (use existing script)
        should_run_script_stage = (
            not start_from or  # Not resuming → run
            STAGE_ORDER.index(start_from) < STAGE_ORDER.index("refs") or  # Starting before refs → run
            user_feedback  # Have feedback from interactive → run to regenerate
        )

        if should_run_script_stage:
            logger.info("\n[STAGE 2] Analyzing article with Claude...")
            stage_start = time.time()

            script_data = analyze_article(
                article.get("body", ""),
                api_key,
                task_dir,
                logger,
                visual_style=visual_style,
                video_kind=video_kind,
                user_feedback=user_feedback
            )
            num_scenes = len(script_data.get("scenes", []))
            num_chars = len(script_data.get("characters", {}))
            logger.info(f"  Created script with {num_scenes} scenes, {num_chars} characters")

            gen_log["stages"]["stage2_analysis"] = {
                "duration_sec": time.time() - stage_start,
                "num_scenes": num_scenes,
                "num_characters": num_chars
            }

            # Check if we should stop after script
            if stop_after == "script":
                logger.info("\n" + "=" * 60)
                logger.info("STOPPED AFTER SCRIPT STAGE")
                logger.info(f"Script saved to: {task_dir / 'analysis' / 'script_draft.json'}")
                logger.info(f"To continue: python microdrama_pipeline.py --task-dir {task_dir} --from-stage refs")
                logger.info("=" * 60)
                save_generation_log(task_dir, gen_log)
                return task_dir

            # INTERACTIVE PHASE 2: Script Review (only if running interactively with stdin)
            if interactive and sys.stdin.isatty():
                logger.info("\n[INTERACTIVE] Phase 2: Script review...")
                phase2_result = interactive_phase2_cli(script_data)

                if phase2_result["action"] == "regenerate":
                    logger.info("  Regenerating script with new instructions...")
                    regen_feedback = phase2_result.get("feedback", "")
                    if user_feedback:
                        user_feedback = f"{user_feedback}\n\nUSER MODIFICATION: {regen_feedback}"
                    else:
                        user_feedback = f"USER MODIFICATION: {regen_feedback}"

                    script_data = analyze_article(
                        article.get("body", ""),
                        api_key,
                        task_dir,
                        logger,
                        visual_style=visual_style,
                        video_kind=video_kind,
                        user_feedback=user_feedback
                    )
                    phase2_result = interactive_phase2_cli(script_data)

                elif phase2_result["action"] == "modify":
                    mod_feedback = phase2_result.get("feedback", "")
                    logger.info(f"  Noted modification: {mod_feedback}")
                    gen_log["interactive_phase2_modification"] = mod_feedback

                gen_log["interactive_phase2"] = phase2_result
                logger.info("  Script approved, continuing...")
        else:
            # Load existing script when resuming without new feedback
            logger.info("\n[STAGE 2] Loading existing script (--from-stage)")
            script_file = task_dir / "script" / "script_v5.json"
            if not script_file.exists():
                script_file = task_dir / "analysis" / "script_draft.json"
            if script_file.exists():
                with open(script_file) as f:
                    script_data = json.load(f)
                logger.info(f"  Loaded existing script from: {script_file}")
            else:
                logger.error("Cannot resume: no script found")
                return None

        if not should_run_stage("refs"):
            logger.info("\n[STAGE 3] Skipping references (--from-stage)")
            # Load existing refs
            refs_file = task_dir / "references" / "selected_refs.json"
            if refs_file.exists():
                with open(refs_file) as f:
                    selected_refs = {k: Path(v) for k, v in json.load(f).items()}
            else:
                selected_refs = {}
        else:
            logger.info("\n[STAGE 3] Searching Yandex for references...")
            stage_start = time.time()

            selected_refs = gather_references(script_data, task_dir, api_key, logger)
            logger.info(f"  Found {len(selected_refs)} reference images")

            gen_log["stages"]["stage3_references"] = {
                "duration_sec": time.time() - stage_start,
                "num_references": len(selected_refs),
                "references": {k: str(v) for k, v in selected_refs.items()}
            }

            # Check stop after refs
            if stop_after == "refs":
                logger.info("\n" + "=" * 60)
                logger.info("STOPPED AFTER REFS STAGE")
                logger.info(f"To continue: python microdrama_pipeline.py --task-dir {task_dir} --from-stage sheets")
                logger.info("=" * 60)
                save_generation_log(task_dir, gen_log)
                return task_dir

        # Stage 3b: Generate character sheets
        if not should_run_stage("sheets"):
            logger.info("\n[STAGE 3b] Skipping sheets (--from-stage)")
            sheets_file = task_dir / "references" / "sheets" / "character_sheets.json"
            if sheets_file.exists():
                with open(sheets_file) as f:
                    character_sheets = {k: Path(v) for k, v in json.load(f).items()}
            else:
                character_sheets = {}
        else:
            logger.info("\n[STAGE 3b] Generating character sheets...")
            stage_start = time.time()

            character_sheets = generate_character_sheets(
                script_data,
                selected_refs,
                task_dir,
                api_key,
                logger
            )

            num_sheets = len([p for p in character_sheets.values() if "sheet_" in p.name])
            gen_log["stages"]["stage3b_sheets"] = {
                "duration_sec": time.time() - stage_start,
                "num_sheets": num_sheets,
                "sheets": {k: str(v) for k, v in character_sheets.items()}
            }

            # Check stop after sheets
            if stop_after == "sheets":
                logger.info("\n" + "=" * 60)
                logger.info("STOPPED AFTER SHEETS STAGE")
                logger.info(f"To continue: python microdrama_pipeline.py --task-dir {task_dir} --from-stage frames")
                logger.info("=" * 60)
                save_generation_log(task_dir, gen_log)
                return task_dir

        # Stage 4: Finalize script (now using character sheets)
        logger.info("\n[STAGE 4] Finalizing script...")
        stage_start = time.time()

        # Use character sheets instead of raw references
        final_script = finalize_script(script_data, character_sheets, task_dir, logger)

        gen_log["stages"]["stage4_script"] = {
            "duration_sec": time.time() - stage_start,
            "script_path": str(task_dir / "script" / "script_v5.json")
        }

        # Stage 5: Generate frames
        if should_run_stage("frames"):
            logger.info("\n[STAGE 5] Generating frames with NanoBanana...")
            stage_start = time.time()

            if not generate_frames(task_dir, logger):
                logger.error("  Frame generation failed")
                gen_log["stages"]["stage5_frames"] = {"error": "failed"}
                save_generation_log(task_dir, gen_log)
                return None

            gen_log["stages"]["stage5_frames"] = {
                "duration_sec": time.time() - stage_start,
                "status": "success"
            }

            # Check stop after frames
            if stop_after == "frames":
                logger.info("\n" + "=" * 60)
                logger.info("STOPPED AFTER FRAMES STAGE")
                logger.info(f"To continue: python microdrama_pipeline.py --task-dir {task_dir} --from-stage videos")
                logger.info("=" * 60)
                save_generation_log(task_dir, gen_log)
                return task_dir
        else:
            logger.info("\n[STAGE 5] Skipping frames (--from-stage)")

        # Stage 6: Generate videos
        if should_run_stage("videos"):
            logger.info("\n[STAGE 6] Generating video clips...")
            stage_start = time.time()

            if not generate_videos(task_dir, logger):
                logger.error("  Video generation failed")
                gen_log["stages"]["stage6_videos"] = {"error": "failed"}
                save_generation_log(task_dir, gen_log)
                return None

            gen_log["stages"]["stage6_videos"] = {
                "duration_sec": time.time() - stage_start,
                "status": "success"
            }

            # Check stop after videos
            if stop_after == "videos":
                logger.info("\n" + "=" * 60)
                logger.info("STOPPED AFTER VIDEOS STAGE")
                logger.info(f"To continue: python microdrama_pipeline.py --task-dir {task_dir} --from-stage assembly")
                logger.info("=" * 60)
                save_generation_log(task_dir, gen_log)
                return task_dir
        else:
            logger.info("\n[STAGE 6] Skipping videos (--from-stage)")

        # Stage 7: Assemble
        if should_run_stage("assembly"):
            logger.info("\n[STAGE 7] Assembling final video...")
            stage_start = time.time()

            output_path = assemble_video(task_dir, logger)
        else:
            logger.info("\n[STAGE 7] Skipping assembly (--from-stage)")
            output_path = None

        if output_path:
            size_mb = output_path.stat().st_size / (1024 * 1024)
            gen_log["stages"]["stage7_assembly"] = {
                "duration_sec": time.time() - stage_start,
                "output_path": str(output_path),
                "size_mb": size_mb
            }
            gen_log["completed_at"] = datetime.now().isoformat()
            gen_log["status"] = "success"

            logger.info("\n" + "=" * 60)
            logger.info("PIPELINE COMPLETE!")
            logger.info(f"Video: {output_path}")
            logger.info(f"Size: {size_mb:.1f} MB")
            logger.info("=" * 60)
        else:
            gen_log["stages"]["stage7_assembly"] = {"error": "failed"}
            gen_log["status"] = "failed"
            logger.error("Assembly failed")

        save_generation_log(task_dir, gen_log)
        return output_path

    except Exception as e:
        logger.exception(f"Pipeline failed: {e}")
        gen_log["error"] = str(e)
        gen_log["status"] = "error"
        save_generation_log(task_dir, gen_log)
        return None


def main():
    parser = argparse.ArgumentParser(description="Microdrama Pipeline")
    parser.add_argument("url", nargs="?", help="Article URL from panorama.pub")
    parser.add_argument("--article", "-a", help="Article text file or text directly")
    parser.add_argument("--slug", help="Story slug for directory name (default: from URL or 'manual')")
    parser.add_argument("--task-dir", "-t", help="Existing task directory to resume")
    parser.add_argument("--style", "-s", help="Visual style (e.g., melodrama, soviet_cartoon)")
    parser.add_argument("--kind", "-k", help="Video kind (e.g., victim_story, absurdist)")
    parser.add_argument("--feedback", "-f", help="User feedback file or text to guide script generation")
    parser.add_argument("--interactive", "-i", action="store_true",
                        help="Interactive mode: ask questions before generating script")
    parser.add_argument("--stage", choices=["script", "refs", "sheets", "frames", "videos", "assembly", "all"],
                        default="all", help="Run up to this stage and stop (default: all)")
    parser.add_argument("--from-stage", choices=["script", "refs", "sheets", "frames", "videos", "assembly"],
                        help="Resume from this stage (requires --task-dir)")

    args = parser.parse_args()

    if args.task_dir:
        task_dir = Path(args.task_dir)
    else:
        task_dir = None

    # Load article text from file or use directly
    article_text = None
    if args.article:
        if os.path.exists(args.article):
            with open(args.article) as f:
                article_text = f.read()
        else:
            # Treat as direct text input
            article_text = args.article

    # Load feedback from file or use directly
    user_feedback = None
    if args.feedback:
        if os.path.exists(args.feedback):
            with open(args.feedback) as f:
                user_feedback = f.read()
        else:
            user_feedback = args.feedback

    # Check for existing story directory with source.md
    if args.slug and not article_text:
        source_file = BASE_DIR / args.slug / "source.md"
        if source_file.exists():
            with open(source_file) as f:
                article_text = f.read()
            print(f"Loaded existing source: {source_file}")

    # When resuming from task directory, load article from article.md
    if task_dir and not article_text:
        article_file = task_dir / "article.md"
        if article_file.exists():
            with open(article_file) as f:
                article_text = f.read()
            print(f"Loaded article from task directory: {article_file}")

    if not args.url and not article_text and not args.task_dir and not args.slug:
        print("Error: Provide URL, article text, task directory, or slug with existing source")
        parser.print_help()
        sys.exit(1)

    run_pipeline(
        args.url or "",
        article_text,
        task_dir,
        visual_style=args.style,
        video_kind=args.kind,
        user_feedback=user_feedback,
        interactive=args.interactive,
        slug=args.slug,
        stop_after=args.stage,
        start_from=getattr(args, 'from_stage', None)
    )


if __name__ == "__main__":
    main()
