#!/usr/bin/env python3
"""
Initialize story structure for V4 workflow.

Creates:
- Story directory with story.md
- Take 001 directory
- Task queue initialized
- Empty feedback.jsonl
- Generation log initialized
- Session ID (reused across takes)
- Pipeline log (JSONL)
- CLAUDE.md with context
"""

import sys
import json
from pathlib import Path
from datetime import datetime, timezone

# Add lib to path
sys.path.insert(0, str(Path(__file__).parent))
from lib.queue import save_queue
from lib.generation_log import setup_generation_log
from lib import (
    get_session_id,
    update_current_symlink,
    log_init,
    save_metadata
)


def initialize_task_queue(story_name: str, take: str, scene_count: int = 9, skip_references: bool = False) -> list:
    """
    Create complete task queue for new generation.

    Starts with scene_count scenes (default 9). If script stage generates 16 scenes
    or user requests it, additional tasks will be added dynamically.

    Args:
        story_name: Story name
        take: Take number
        scene_count: Number of scenes (9 or 16, default 9)
        skip_references: Whether to skip reference search stage

    Returns:
        List of task dictionaries (~30-40 tasks for 9 scenes)
    """
    tasks = []
    task_id = 1

    # Stage 0: Init
    tasks.append({
        "id": task_id,
        "stage": "00-init",
        "action": "init_story",
        "status": "completed",  # Already done by running this script
        "result": "initialized",
        "iteration": None,
        "parent_iteration": None,
        "message": None,
        "feedback": None,
        "error": None
    })
    task_id += 1

    # Stage 1: Story (already provided during init)
    tasks.append({
        "id": task_id,
        "stage": "01-story",
        "action": "await_feedback",
        "status": "completed",  # Story text already provided
        "deps": [1],
        "message": "Story text provided",
        "iteration": "001",
        "parent_iteration": None,
        "result": "01-story/001/story.md",
        "feedback": None,
        "error": None
    })
    story_feedback_id = task_id
    task_id += 1

    # Stage 2: Understanding [REQUIRED]
    tasks.append({
        "id": task_id,
        "stage": "02-understanding",
        "action": "generate",
        "status": "pending",
        "deps": [story_feedback_id],
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "message": None,
        "feedback": None,
        "error": None
    })
    task_id += 1

    tasks.append({
        "id": task_id,
        "stage": "02-understanding",
        "action": "await_feedback",
        "status": "pending",
        "deps": [task_id - 1],
        "message": "Review understanding.md",
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "feedback": None,
        "error": None
    })
    understanding_feedback_id = task_id
    task_id += 1

    # Stage 3: Script [REQUIRED]
    tasks.append({
        "id": task_id,
        "stage": "03-script",
        "action": "generate",
        "status": "pending",
        "deps": [understanding_feedback_id],
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "message": None,
        "feedback": None,
        "error": None
    })
    task_id += 1

    tasks.append({
        "id": task_id,
        "stage": "03-script",
        "action": "await_feedback",
        "status": "pending",
        "deps": [task_id - 1],
        "message": "Review script.json (scene count will be determined here)",
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "feedback": None,
        "error": None
    })
    script_feedback_id = task_id
    task_id += 1

    # Stage 4: Assets [REQUIRED]
    tasks.append({
        "id": task_id,
        "stage": "04-assets",
        "action": "generate",
        "status": "pending",
        "deps": [script_feedback_id],
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "message": None,
        "feedback": None,
        "error": None
    })
    task_id += 1

    tasks.append({
        "id": task_id,
        "stage": "04-assets",
        "action": "await_feedback",
        "status": "pending",
        "deps": [task_id - 1],
        "message": "Review assets.json",
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "feedback": None,
        "error": None
    })
    assets_feedback_id = task_id
    task_id += 1

    # Stage 5: References [OPTIONAL]
    if not skip_references:
        tasks.append({
            "id": task_id,
            "stage": "05-references",
            "action": "check_skip",
            "status": "pending",
            "deps": [assets_feedback_id],
            "message": "Claude: Are assets generic enough to skip reference search?",
            "iteration": None,
            "parent_iteration": None,
            "result": None,
            "feedback": None,
            "error": None
        })
        task_id += 1

        tasks.append({
            "id": task_id,
            "stage": "05-references",
            "action": "search_or_skip",
            "status": "pending",
            "deps": [task_id - 1],
            "iteration": None,
            "parent_iteration": None,
            "result": None,
            "message": None,
            "feedback": None,
            "error": None
        })
        references_complete_id = task_id
        task_id += 1
    else:
        references_complete_id = assets_feedback_id

    # Stage 6: Characteristic [REQUIRED]
    tasks.append({
        "id": task_id,
        "stage": "06-characteristic",
        "action": "generate",
        "status": "pending",
        "deps": [references_complete_id],
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "message": None,
        "feedback": None,
        "error": None
    })
    task_id += 1

    tasks.append({
        "id": task_id,
        "stage": "06-characteristic",
        "action": "await_feedback",
        "status": "pending",
        "deps": [task_id - 1],
        "message": "Review characteristic shot",
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "feedback": None,
        "error": None
    })
    characteristic_feedback_id = task_id
    task_id += 1

    # Stage 7: Grid [CRITICAL CHECKPOINT]
    tasks.append({
        "id": task_id,
        "stage": "07-grids",
        "action": "generate",
        "status": "pending",
        "deps": [characteristic_feedback_id],
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "message": None,
        "feedback": None,
        "error": None
    })
    task_id += 1

    tasks.append({
        "id": task_id,
        "stage": "07-grids",
        "action": "await_feedback",
        "status": "pending",
        "deps": [task_id - 1],
        "message": "CRITICAL: Review grid before proceeding (defines all characters, cannot change after)",
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "feedback": None,
        "error": None
    })
    grid_feedback_id = task_id
    task_id += 1

    # Stage 8: Script adaptation
    tasks.append({
        "id": task_id,
        "stage": "08-script-adapted",
        "action": "generate",
        "status": "pending",
        "deps": [grid_feedback_id],
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "message": None,
        "feedback": None,
        "error": None
    })
    script_adapted_id = task_id
    task_id += 1

    # Stage 9: Frame extraction (batch)
    tasks.append({
        "id": task_id,
        "stage": "09-frames",
        "action": "extract_batch",
        "frame_count": scene_count,
        "status": "pending",
        "deps": [script_adapted_id],
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "message": None,
        "feedback": None,
        "error": None
    })
    frame_batch_id = task_id
    task_id += 1

    # Stage 9: Frame feedback (one task per frame)
    frame_feedback_ids = []
    for i in range(scene_count):
        tasks.append({
            "id": task_id,
            "stage": "09-frames",
            "action": "await_feedback",
            "frame_num": i + 1,
            "status": "pending",
            "deps": [frame_batch_id],
            "message": f"Review frame {i+1}/{scene_count}",
            "iteration": None,
            "parent_iteration": None,
            "result": None,
            "feedback": None,
            "error": None
        })
        frame_feedback_ids.append(task_id)
        task_id += 1

    # Stage 10: Video generation (batch, depends on all frame approvals)
    tasks.append({
        "id": task_id,
        "stage": "10-videos",
        "action": "generate_batch",
        "video_count": scene_count,
        "api": "kling-direct",  # or "kling-pollo", "veo3-laozhang", "veo3-pollo"
        "status": "pending",
        "deps": frame_feedback_ids,
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "message": None,
        "feedback": None,
        "error": None
    })
    video_batch_id = task_id
    task_id += 1

    # Stage 10: Video feedback (one task per video)
    video_feedback_ids = []
    for i in range(scene_count):
        tasks.append({
            "id": task_id,
            "stage": "10-videos",
            "action": "await_feedback",
            "scene_num": i + 1,
            "status": "pending",
            "deps": [video_batch_id],
            "message": f"Review video {i+1}/{scene_count}",
            "iteration": None,
            "parent_iteration": None,
            "result": None,
            "feedback": None,
            "error": None
        })
        video_feedback_ids.append(task_id)
        task_id += 1

    # Stage 11: Captions
    tasks.append({
        "id": task_id,
        "stage": "11-captions",
        "action": "generate",
        "status": "pending",
        "deps": video_feedback_ids,
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "message": None,
        "feedback": None,
        "error": None
    })
    captions_id = task_id
    task_id += 1

    # Stage 12: Assembly
    tasks.append({
        "id": task_id,
        "stage": "12-assembly",
        "action": "assemble",
        "status": "pending",
        "deps": [captions_id],
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "message": None,
        "feedback": None,
        "error": None
    })
    task_id += 1

    tasks.append({
        "id": task_id,
        "stage": "12-assembly",
        "action": "await_feedback",
        "status": "pending",
        "deps": [task_id - 1],
        "message": "Review final video",
        "iteration": None,
        "parent_iteration": None,
        "result": None,
        "feedback": None,
        "error": None
    })

    return tasks


def create_claude_md(story_path: Path, story_name: str, take: str, session_id: str) -> None:
    """
    Create CLAUDE.md file from template or generate simple version.

    Args:
        story_path: Path to story directory
        story_name: Story name
        take: Take number
        session_id: Claude session ID
    """
    template_path = Path(__file__).parent.parent / "templates" / "CLAUDE.md.template"

    if not template_path.exists():
        # Create a simple CLAUDE.md without template
        claude_md = f"""# Microdrama V4 - {story_name}

## Current State

- **Take:** {take}
- **Stage:** 00-init (queue initialized)
- **Session ID:** {session_id}

## Files

Each take directory contains:

- **task_queue.json** - Task queue (editable JSON array, tracks workflow state)
- **feedback.jsonl** - Feedback log (append-only JSONL)
- **generation.log** - Verbose text log

Story-level:
- **pipeline.log** - Structured JSONL log (all takes)
- **session.txt** - Claude session ID (reused across takes)

## Next Step

Run stage 2 (understanding generation).
"""
        claude_path = story_path / "CLAUDE.md"
        claude_path.write_text(claude_md)
        print(f"Created CLAUDE.md")
        return

    try:
        template = template_path.read_text()

        # Fill in template variables
        claude_md = template.format(
            story_name=story_name,
            current_take=take,
            current_stage="00-init (queue initialized)",
            session_id=session_id,
            feedback_summary="None yet - this is the first take."
        )

        # Write CLAUDE.md
        claude_path = story_path / "CLAUDE.md"
        claude_path.write_text(claude_md)
        print(f"Created CLAUDE.md")
    except Exception as e:
        print(f"Warning: Could not create CLAUDE.md: {e}")


def init_story(story_name: str, story_text: str, base_path: Path = None) -> Path:
    """
    Initialize story structure.

    Args:
        story_name: Story slug (e.g., "water-walking-dog")
        story_text: Full story text
        base_path: Base path (default: /mnt/video/microdrama-v3/)

    Returns:
        Path to take directory
    """
    if base_path is None:
        base_path = Path("/mnt/video/microdrama-v3")

    # Create directory name with date prefix: YYYYMMDD-slug
    today = datetime.now().strftime("%Y%m%d")
    story_dirname = f"{today}-{story_name}"

    # Create story directory
    story_path = base_path / story_dirname
    story_path.mkdir(parents=True, exist_ok=True)

    # Write story.md
    story_file = story_path / "story.md"
    with open(story_file, 'w') as f:
        f.write(story_text)

    print(f"Created story directory: {story_path}")
    print(f"Created story.md")

    # Generate or load session ID (one per story, reused across takes)
    session_file = story_path / "session.txt"
    if session_file.exists():
        session_id = session_file.read_text().strip()
        print(f"Reusing session ID: {session_id}")
    else:
        session_id = get_session_id(story_name)
        session_file.write_text(session_id)
        print(f"Created session ID: {session_id}")

    # Create pipeline.log (JSONL structured log) if it doesn't exist
    pipeline_log = story_path / "pipeline.log"
    if not pipeline_log.exists():
        pipeline_log.touch()
        print(f"Created pipeline.log")

    # Create take 001
    take_path = story_path / "001"
    take_path.mkdir(parents=True, exist_ok=True)

    # Create 01-story directory and copy story
    story_stage_path = take_path / "01-story" / "001"
    story_stage_path.mkdir(parents=True, exist_ok=True)

    story_copy = story_stage_path / "story.md"
    with open(story_copy, 'w') as f:
        f.write(story_text)

    # Create metadata for 01-story
    story_metadata = {
        "stage": "01-story",
        "take": "001",
        "iteration": "001",
        "parent_iteration": None,
        "source": "../../../story.md",
        "timestamp": datetime.now(timezone.utc).isoformat()
    }
    save_metadata(story_stage_path / "metadata.json", story_metadata)

    # Create current symlink for 01-story
    story_stage_dir = take_path / "01-story"
    update_current_symlink(story_stage_dir, "001")

    print(f"Created take directory: {take_path}")
    print(f"Created 01-story/001/story.md with metadata")

    # Initialize task queue
    task_queue = initialize_task_queue(story_name, "001")
    save_queue(take_path, task_queue)
    print(f"Initialized task queue: {len(task_queue)} tasks")

    # Initialize empty feedback.jsonl
    feedback_file = take_path / "feedback.jsonl"
    feedback_file.touch()
    print(f"Created feedback.jsonl")

    # Initialize generation log
    logger = setup_generation_log(take_path)
    logger.orchestrator_start(story_name, "001")
    logger.info(f"Story initialized: {story_name}")
    logger.info(f"Task queue created with {len(task_queue)} tasks")
    print(f"Created generation.log")

    # Create story-level current symlink pointing to this take
    update_current_symlink(story_path, "001")
    print(f"Created current -> 001 symlink")

    # Log initialization to pipeline.log (manual write for v4 date-prefixed dirs)
    log_entry = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "event": "init",
        "action": "init_take",
        "take": "001",
        "previous_take": None,
        "story_name": story_name
    }
    with open(pipeline_log, 'a') as f:
        f.write(json.dumps(log_entry) + '\n')
    print(f"Logged init to pipeline.log")

    # Create CLAUDE.md with context (optional, won't fail if template missing)
    create_claude_md(story_path, story_name, "001", session_id)

    print()
    print("=" * 60)
    print(f"Story '{story_name}' initialized successfully!")
    print("=" * 60)
    print(f"Story path: {story_path}")
    print(f"Take path: {take_path}")
    print(f"Session ID: {session_id}")
    print(f"Task queue: {len(task_queue)} tasks initialized")
    print()
    print(f"Next step: Run stage 2 (understanding generation)")
    print()

    return take_path


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python init_story_v4.py <story-name> <story-text>")
        print()
        print("Example:")
        print("  python init_story_v4.py water-walking-dog \"Dave was a keen duck hunter...\"")
        sys.exit(1)

    story_name = sys.argv[1]
    story_text = sys.argv[2]

    init_story(story_name, story_text)
