#!/usr/bin/env python3
"""
Stage 5: References

Checks if reference images are needed (from assets.json).
If needed, automatically searches Yandex Images for each entity,
downloads candidates, and calls Claude CLI (same session as pipeline)
to review and select the best match.

Falls back to manual mode if Yandex credentials are missing.
"""

import json
import os
import sys
import subprocess
import shutil
import time
import re
from pathlib import Path
from datetime import datetime, timezone
from typing import Dict, List, Optional

from dotenv import load_dotenv

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

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.yandex_search import search_and_download, get_credentials


def _extract_query_from_description(desc: str, max_len: int = 50) -> str:
    """Extract a short, focused search query from description."""
    # First sentence usually names the place/thing
    first_sent = desc.split(".")[0].strip()
    # Remove parenthetical notes like "(Ayasofya)"
    first_sent = re.sub(r"\s*\([^)]*\)", "", first_sent)
    # Remove em-dashes and everything after (usually elaboration)
    for sep in [" — ", " – ", " - "]:
        if sep in first_sent:
            first_sent = first_sent.split(sep)[0].strip()
    # Remove time-of-day qualifiers (not useful for image search)
    first_sent = re.sub(r"\s+at\s+(blue hour|golden hour|sunset|sunrise|twilight|dawn|dusk)", "", first_sent)
    # Remove filler words to keep query short
    for filler in ["directly outside ", "viewed from the ", "leading toward "]:
        first_sent = first_sent.replace(filler, "")
    # Truncate at word boundary
    if len(first_sent) > max_len:
        first_sent = first_sent[:max_len].rsplit(" ", 1)[0]
    return first_sent


def _generate_search_queries(assets_data: dict) -> List[dict]:
    """
    Generate search queries for entities that need references.

    Returns list of {entity_type, entity_id, query, description}
    """
    queries = []

    for loc in assets_data.get("locations", []):
        if not loc.get("needs_reference"):
            continue
        entity_id = loc["id"]
        desc = loc.get("description", "")

        query = _extract_query_from_description(desc)
        if not query:
            query = entity_id.replace("_", " ").title()

        queries.append({
            "entity_type": "location",
            "entity_id": entity_id,
            "query": query,
            "description": desc,
        })

    for char in assets_data.get("characters", []):
        if not char.get("needs_reference"):
            continue
        entity_id = char["id"]
        desc = char.get("appearance", char.get("description", ""))
        queries.append({
            "entity_type": "character",
            "entity_id": entity_id,
            "query": f"{entity_id.replace('_', ' ')} portrait",
            "description": desc,
        })

    for obj in assets_data.get("objects", []):
        if not obj.get("needs_reference"):
            continue
        entity_id = obj["id"]
        desc = obj.get("description", "")
        query = _extract_query_from_description(desc) if desc else entity_id.replace("_", " ")
        queries.append({
            "entity_type": "object",
            "entity_id": entity_id,
            "query": query,
            "description": desc,
        })

    return queries


def _build_selection_prompt(iter_path: Path, queries: List[dict],
                            candidates_map: Dict[str, List[Path]]) -> str:
    """
    Build a prompt for Claude CLI to review downloaded images and select
    the best reference for each entity. Claude uses Read tool to view images.
    """
    lines = [
        "# Stage 5: Reference Image Selection",
        "",
        "I've downloaded candidate reference images from Yandex search.",
        "Your task: review each set of candidates and select the BEST one for each entity.",
        "",
        "For each entity below, use the Read tool to view the candidate images,",
        "then copy the best one to the output directory with the entity ID as filename.",
        "",
        f"## Output Directory: {iter_path}",
        "",
    ]

    for q in queries:
        entity_id = q["entity_id"]
        entity_type = q["entity_type"]
        description = q["description"]

        cands = candidates_map.get(entity_id, [])
        if not cands:
            continue

        lines.append(f"### {entity_type.upper()}: {entity_id}")
        lines.append(f"Description: {description[:300]}")
        lines.append(f"Candidates ({len(cands)} images):")
        for c in cands:
            lines.append(f"  - {c}")
        lines.append("")
        lines.append(f"Selection criteria for {entity_type}:")
        if entity_type == "location":
            lines.append("  - Correct place (matches the described location)")
            lines.append("  - Good composition, clear view of key architectural features")
            lines.append("  - High resolution, suitable as visual reference for AI video generation")
            lines.append("  - No watermarks, no heavy text overlays")
        elif entity_type == "character":
            lines.append("  - Clear face, good lighting, appropriate style")
        else:
            lines.append("  - Clear view, correct item, high quality")
        lines.append("")

    lines.extend([
        "## Instructions",
        "",
        "1. Read each candidate image using the Read tool",
        "2. For each entity, pick the SINGLE best image",
        "3. Copy it to the output directory as: {entity_id}.{ext}",
        "   Use the Bash tool: cp /path/to/best.jpg {iter_path}/{entity_id}.jpg",
        "4. After processing all entities, write a JSON file:",
        f"   {iter_path}/selections.json",
        "   Format: {{\"entity_id\": \"selected_filename.ext\", ...}}",
        "",
        "Work through ALL entities. Be selective — pick the image that best",
        "matches the description and would work as a reference for video generation.",
    ])

    return "\n".join(lines)


def _check_existing_references(take_path: Path, entity_ids: List[str]) -> Dict[str, Path]:
    """Check if project already has references for some entities."""
    project_path = take_path.parent
    found = {}

    # Check project references/locations/
    locations_dir = project_path / "references" / "locations"
    if locations_dir.exists():
        for img in locations_dir.iterdir():
            if not img.is_file():
                continue
            stem = img.stem.lower()
            for eid in entity_ids:
                if eid in stem or stem in eid:
                    found[eid] = img

    return found


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

    If references not needed → empty iteration.
    If references needed → auto-search Yandex, download, Claude selects best.
    Falls back to manual if Yandex credentials unavailable.
    """
    logger = setup_generation_log(take_path)
    logger.section("Stage 5: References")

    # Read assets to check if references needed
    assets_iter = inputs.get("04-assets")
    if not assets_iter:
        return {"status": "error", "error": "No 04-assets iteration in inputs"}

    assets_file = take_path / "04-assets" / assets_iter / "assets.json"
    if not assets_file.exists():
        return {"status": "error", "error": f"Assets file not found: {assets_file}"}

    try:
        assets_data = json.loads(assets_file.read_text())
    except json.JSONDecodeError as e:
        return {"status": "error", "error": f"Invalid JSON in assets.json: {e}"}

    needs_references = assets_data.get("needs_references", False)

    # Create iteration directory
    stage_dir = take_path / "05-references"
    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}")
    logger.info(f"References needed: {needs_references}")

    story_path = take_path.parent

    if not needs_references:
        # No references needed — empty iteration
        metadata = {
            "stage": "05-references",
            "iteration": iteration,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "based_on": inputs,
            "refined_from": refined_from,
            "rating": None,
            "result": f"05-references/{iteration}/",
            "needs_references": False,
            "references": {},
        }
        save_metadata(iter_path / "metadata.json", metadata)

        log_event(
            story_path, "stage_complete",
            {"action": "references_not_needed", "iteration": iteration},
            take=take_path.name, stage="05-references", iteration=iteration
        )

        msg = "No reference images needed. Stage completed."
        print(f"\n05-references/{iteration}/ — no references needed")
        return {
            "status": "success",
            "iteration": iteration,
            "result": f"05-references/{iteration}/",
            "user_message": msg,
        }

    # --- References needed: auto-search ---
    folder_id, api_key = get_credentials()
    can_search = bool(folder_id and api_key)

    if not can_search:
        logger.warning("Yandex credentials not found — manual mode")
        metadata = {
            "stage": "05-references",
            "iteration": iteration,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "based_on": inputs,
            "refined_from": refined_from,
            "rating": None,
            "result": f"05-references/{iteration}/",
            "needs_references": True,
            "references": {},
        }
        save_metadata(iter_path / "metadata.json", metadata)
        msg = (f"References needed but YANDEX_FOLDER_ID/YANDEX_API_KEY not set. "
               f"Add reference images manually to 05-references/{iteration}/.")
        print(f"\n{msg}")
        return {
            "status": "success",
            "iteration": iteration,
            "result": f"05-references/{iteration}/",
            "user_message": msg,
        }

    # Generate search queries from assets
    queries = _generate_search_queries(assets_data)
    logger.info(f"Entities needing references: {len(queries)}")

    # Check for existing project-level references
    entity_ids = [q["entity_id"] for q in queries]
    existing = _check_existing_references(take_path, entity_ids)
    if existing:
        logger.info(f"Found {len(existing)} existing references in project")

    # Search directory for raw candidates
    search_dir = iter_path / "search"
    ensure_dir(search_dir)

    # Track stats
    selected_refs = {}
    candidates_map = {}  # entity_id → list of downloaded paths
    search_stats = {"searched": 0, "downloaded": 0, "selected": 0, "existing": 0}

    print(f"\nSearching references for {len(queries)} entities...")

    # Phase 1: Copy existing + search & download candidates
    queries_needing_selection = []

    for q in queries:
        entity_id = q["entity_id"]
        entity_type = q["entity_type"]

        print(f"\n  [{entity_type.upper()}] {entity_id}")

        # Skip if already have a reference
        if entity_id in existing:
            existing_path = existing[entity_id]
            dest = iter_path / f"{entity_id}{existing_path.suffix}"
            shutil.copy2(existing_path, dest)
            selected_refs[entity_id] = dest.name
            search_stats["existing"] += 1
            print(f"    Using existing: {existing_path.name}")
            continue

        # Search and download
        query = q["query"]
        print(f"    Query: {query}")
        search_stats["searched"] += 1

        candidates = search_and_download(
            query=query,
            output_dir=search_dir / entity_id,
            prefix=entity_id,
            count=15,
            folder_id=folder_id,
            api_key=api_key,
        )

        search_stats["downloaded"] += len(candidates)
        print(f"    Downloaded: {len(candidates)} candidates")

        if not candidates:
            logger.warning(f"No images found for {entity_id}")
            continue

        if len(candidates) == 1:
            # Only one candidate — use it directly
            dest = iter_path / f"{entity_id}{candidates[0].suffix}"
            shutil.copy2(candidates[0], dest)
            selected_refs[entity_id] = dest.name
            search_stats["selected"] += 1
            print(f"    Auto-selected (only candidate): {dest.name}")
        else:
            candidates_map[entity_id] = candidates
            queries_needing_selection.append(q)

        time.sleep(0.5)

    # Phase 2: Call Claude CLI to select best from candidates
    if queries_needing_selection and candidates_map:
        print(f"\nCalling Claude to select best references for {len(queries_needing_selection)} entities...")

        # Load session ID
        session_file = story_path / "session.txt"
        if not session_file.exists():
            logger.error("Session file not found — selecting first candidates")
            for q in queries_needing_selection:
                eid = q["entity_id"]
                cands = candidates_map.get(eid, [])
                if cands:
                    dest = iter_path / f"{eid}{cands[0].suffix}"
                    shutil.copy2(cands[0], dest)
                    selected_refs[eid] = dest.name
                    search_stats["selected"] += 1
        else:
            session_id = session_file.read_text().strip()

            prompt = _build_selection_prompt(iter_path, queries_needing_selection, candidates_map)

            # Save prompt for debugging
            (iter_path / "prompt.txt").write_text(prompt)

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

                if result.returncode != 0:
                    error_detail = result.stderr or result.stdout
                    logger.error(f"Claude returned code {result.returncode}: {error_detail}")
                else:
                    logger.info("Claude selection completed")

            except subprocess.TimeoutExpired:
                logger.error("Claude timed out during reference selection")
            except Exception as e:
                logger.error(f"Exception calling Claude: {e}")

            # Check what Claude selected — look for copied files in iter_path
            for q in queries_needing_selection:
                eid = q["entity_id"]
                if eid in selected_refs:
                    continue  # already handled

                # Check if Claude copied the file
                for ext in [".jpg", ".jpeg", ".png", ".webp"]:
                    dest = iter_path / f"{eid}{ext}"
                    if dest.exists() and dest.stat().st_size > 1000:
                        selected_refs[eid] = dest.name
                        search_stats["selected"] += 1
                        print(f"    {eid}: selected → {dest.name}")
                        break
                else:
                    # Fallback: use first candidate
                    cands = candidates_map.get(eid, [])
                    if cands:
                        dest = iter_path / f"{eid}{cands[0].suffix}"
                        if not dest.exists():
                            shutil.copy2(cands[0], dest)
                        selected_refs[eid] = dest.name
                        search_stats["selected"] += 1
                        logger.warning(f"Fallback to first candidate for {eid}")

    # Save references manifest
    manifest = {
        "entity_references": selected_refs,
        "search_queries": {q["entity_id"]: q["query"] for q in queries},
        "stats": search_stats,
    }
    (iter_path / "references.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2)
    )
    logger.info(f"Saved references.json: {search_stats}")

    # Save metadata
    metadata = {
        "stage": "05-references",
        "iteration": iteration,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "based_on": inputs,
        "refined_from": refined_from,
        "rating": None,
        "result": f"05-references/{iteration}/",
        "needs_references": True,
        "references": selected_refs,
        "search_stats": search_stats,
    }
    save_metadata(iter_path / "metadata.json", metadata)
    logger.info("Saved metadata")

    # Log event
    log_event(
        story_path, "stage_complete",
        {"action": "references_searched", "iteration": iteration,
         "selected": len(selected_refs), "total_needed": len(queries)},
        take=take_path.name, stage="05-references", iteration=iteration
    )

    # Build user message
    total_needed = len(queries)
    total_found = len(selected_refs)
    missing = [q["entity_id"] for q in queries if q["entity_id"] not in selected_refs]

    parts = [f"Reference search complete: {total_found}/{total_needed} found."]
    if search_stats["existing"] > 0:
        parts.append(f"{search_stats['existing']} from existing project references.")
    if search_stats["selected"] > 0:
        parts.append(f"{search_stats['selected']} found via Yandex search.")
    if missing:
        parts.append(f"Missing: {', '.join(missing)}. Add manually to 05-references/{iteration}/.")

    msg = " ".join(parts)
    (iter_path / "user_message.txt").write_text(msg)

    print(f"\n{'=' * 60}")
    print(f"References: {total_found}/{total_needed} found")
    print(f"  Existing: {search_stats['existing']}")
    print(f"  Searched: {search_stats['searched']}, Downloaded: {search_stats['downloaded']}")
    print(f"  Selected: {search_stats['selected']}")
    if missing:
        print(f"  Missing: {', '.join(missing)}")
    print(f"{'=' * 60}")

    return {
        "status": "success",
        "iteration": iteration,
        "result": f"05-references/{iteration}/",
        "user_message": msg,
    }
