#!/usr/bin/env python3
"""
Stage 4: Assets Identification

Reads script from 03-script/{inputs["03-script"]}/script.json, calls Claude to identify all
characters, locations, and objects needed, and determines which need reference images.

Creates 04-assets/XXX/assets.json.

V5 format: uses inputs dict, feedback history chain, no current symlinks.
"""

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

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

from lib import (
    save_metadata,
    ensure_dir,
    setup_generation_log,
    log_event,
    claude_env
)
from lib.path_utils_v3 import get_next_iteration
from lib.feedback_history_v5 import get_feedback_history_for_claude


def execute(take_path: Path, story_name: str, inputs: dict,
            feedback: str = None, refined_from: str = None) -> dict:
    """
    Execute Stage 4: Assets identification.

    Args:
        take_path: Path to take directory
        story_name: Story name (slug)
        inputs: Parent stage iterations, e.g. {"03-script": "001"}
        feedback: User feedback text (if any)
        refined_from: Iteration being refined (if refinement)

    Returns:
        Result dict with status, iteration, error (if any)
    """
    logger = setup_generation_log(take_path)
    logger.section("Stage 4: Assets Identification")

    # Read script from 03-script/{inputs["03-script"]}
    script_iter = inputs.get("03-script")
    if not script_iter:
        return {"status": "error", "error": "No 03-script iteration in inputs"}

    script_iter_path = take_path / "03-script" / script_iter
    script_file = script_iter_path / "script.json"

    if not script_file.exists():
        return {
            "status": "error",
            "error": f"Script file not found: {script_file}"
        }

    # Validate script is valid JSON
    try:
        script_data = json.loads(script_file.read_text())
        logger.info(f"Read script: {len(script_data.get('scenes', []))} scenes")
    except json.JSONDecodeError as e:
        return {
            "status": "error",
            "error": f"Invalid JSON in script file: {str(e)}"
        }

    # Determine iteration number
    stage_dir = take_path / "04-assets"
    ensure_dir(stage_dir)
    iteration = get_next_iteration(stage_dir)

    iter_path = stage_dir / iteration
    ensure_dir(iter_path)

    logger.info(f"Creating iteration {iteration}")

    # Save feedback if provided
    if feedback:
        (iter_path / "feedback.txt").write_text(feedback)
        logger.info("Saved feedback.txt")

    # Load session ID
    session_file = take_path.parent / "session.txt"
    if not session_file.exists():
        return {
            "status": "error",
            "error": "Session file not found"
        }

    session_id = session_file.read_text().strip()
    logger.info(f"Session ID: {session_id}")

    # Feedback history from refinement chain
    feedback_history = get_feedback_history_for_claude(
        take_path, "04-assets", refined_from
    )

    # Build prompt
    prompt_template_path = Path(__file__).parent.parent.parent / "prompts" / "stage_assets.txt"

    if not prompt_template_path.exists():
        return {
            "status": "error",
            "error": f"Prompt template not found: {prompt_template_path}"
        }

    prompt = prompt_template_path.read_text()

    # Inject existing reference filenames so Claude uses matching IDs
    story_path = take_path.parent
    refs_path = story_path / "references"
    if refs_path.exists():
        existing_refs = []
        for category in ["people", "locations", "objects"]:
            cat_path = refs_path / category
            if cat_path.exists():
                if category == "people":
                    # People are directories
                    for d in sorted(cat_path.iterdir()):
                        if d.is_dir():
                            existing_refs.append(f"  - {category}/{d.name}/ (character ID: {d.name})")
                else:
                    # Locations/objects are files
                    for f in sorted(cat_path.glob("*.*")):
                        if f.suffix.lower() in (".jpg", ".jpeg", ".png", ".webp"):
                            existing_refs.append(f"  - {category}/{f.name} (asset ID: {f.stem})")
        if existing_refs:
            refs_section = (
                "\n\n## Existing Reference Images\n\n"
                "The project already has these reference images. **You MUST use the exact "
                "asset IDs shown below** when referring to these assets in assets.json. "
                "If a script asset matches an existing reference, use that ID. "
                "Only create new IDs for assets that have no existing reference.\n\n"
                + "\n".join(existing_refs)
                + "\n"
            )
            prompt += refs_section
            logger.info(f"Injected {len(existing_refs)} existing reference paths into prompt")

    # Inject frame index from completed episodes
    from lib.frame_index import load_frame_index, format_frame_index_for_prompt
    frame_index = load_frame_index(story_path)
    if frame_index.get("frames"):
        frame_section = format_frame_index_for_prompt(frame_index)
        prompt += frame_section
        logger.info(f"Injected frame index ({len(frame_index['frames'])} frames) into prompt")

    # Add feedback
    if feedback_history:
        prompt += f"\n\n## Previous Feedback\n\n{feedback_history}\n"
    elif feedback:
        prompt += f"\n\n## User Feedback\n\n{feedback}\n"

    # Prepend output directory so Claude knows where to write files
    # cwd is story_path for session continuity, so Claude needs explicit path
    relative_iter_path = iter_path.relative_to(story_path)
    prompt = f"## Output Directory\n\nWrite ALL output files to: {relative_iter_path}/\n\n{prompt}"

    # Save prompt
    (iter_path / "prompt.txt").write_text(prompt)
    logger.info(f"Saved prompt")

    # Call Claude
    logger.info("Calling Claude...")
    print("\n" + "=" * 60)
    print(f"Stage 4: Assets Identification (iteration {iteration})")
    print(f"Session ID: {session_id}")
    print("=" * 60)

    try:
        result = subprocess.run(
            [
                "claude",
                "--resume", session_id,
                "--allowedTools", "Read,Write,Glob,Grep",
                "-p", prompt
            ],
            capture_output=True,
            text=True,
            cwd=str(story_path),
            timeout=600,  # 10 minutes
            env=claude_env()
        )

        if result.returncode != 0:
            error_detail = result.stderr or result.stdout
            error_msg = f"Claude returned code {result.returncode}: {error_detail}"
            logger.error(error_msg)
            return {
                "status": "error",
                "error": error_msg
            }

        logger.info("Claude completed successfully")

    except subprocess.TimeoutExpired:
        error_msg = "Claude timed out after 10 minutes"
        logger.error(error_msg)
        return {
            "status": "error",
            "error": error_msg
        }

    except Exception as e:
        error_msg = f"Exception calling Claude: {str(e)}"
        logger.error(error_msg)
        return {
            "status": "error",
            "error": error_msg
        }

    # Validate output files
    assets_file = iter_path / "assets.json"
    if not assets_file.exists():
        error_msg = "Expected output file not created: assets.json"
        logger.error(error_msg)
        return {
            "status": "error",
            "error": error_msg
        }

    logger.info(f"Created assets.json ({assets_file.stat().st_size} bytes)")

    # Validate assets.json is valid JSON
    try:
        assets_data = json.loads(assets_file.read_text())

        # Check required fields
        if "needs_references" not in assets_data:
            logger.warning("assets.json missing 'needs_references' field")

        logger.info(f"Assets: {len(assets_data.get('characters', []))} characters, "
                   f"{len(assets_data.get('locations', []))} locations, "
                   f"{len(assets_data.get('objects', []))} objects")

        needs_refs = assets_data.get('needs_references', False)
        logger.info(f"Needs reference images: {needs_refs}")

    except json.JSONDecodeError as e:
        error_msg = f"Invalid JSON in assets.json: {str(e)}"
        logger.error(error_msg)
        return {
            "status": "error",
            "error": error_msg
        }

    # Check for user message
    user_message_file = iter_path / "user_message.txt"
    user_message = None
    if user_message_file.exists():
        user_message = user_message_file.read_text().strip()
        logger.info(f"Created user_message.txt ({user_message_file.stat().st_size} bytes)")
    else:
        logger.warning("user_message.txt not created by Claude")
        user_message = "Asset identification complete. Please review the assets.json file."

    # Create metadata (V5 format)
    metadata = {
        "stage": "04-assets",
        "iteration": iteration,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "based_on": inputs,
        "refined_from": refined_from,
        "rating": None,
        "result": f"04-assets/{iteration}/assets.json",
        "model": "claude-sonnet-4.5",
        "session_id": session_id,
        "needs_references": assets_data.get("needs_references", False),
    }

    save_metadata(iter_path / "metadata.json", metadata)
    logger.info("Saved metadata")

    # Log event
    log_event(
        story_path, "stage_complete",
        {"action": "assets_identified", "iteration": iteration,
         "result": f"04-assets/{iteration}/assets.json",
         "needs_references": assets_data.get("needs_references", False)},
        take=take_path.name, stage="04-assets", iteration=iteration
    )
    logger.info("Logged to pipeline.log")

    print("\n" + "=" * 60)
    print(f"Assets identified successfully (iteration {iteration})")
    print(f"Location: {iter_path / 'assets.json'}")
    print(f"Needs references: {assets_data.get('needs_references', False)}")
    print("=" * 60)

    return {
        "status": "success",
        "iteration": iteration,
        "result": f"04-assets/{iteration}/assets.json",
        "user_message": user_message,
        "needs_references": assets_data.get("needs_references", False)
    }
