"""
Reference Photo Analysis & Indexing

Analyzes reference photos using Claude CLI (multimodal) and generates
a rich index (references/index.json) with descriptions, types, and
usefulness ratings. Downstream stages use this index to pick the best
photos automatically.

Usage:
    from lib.reference_analysis import analyze_references, get_best_photos

    index = analyze_references(project_path)
    faces = get_best_photos(index, "people", person="main", count=3)
"""

import json
import logging
import subprocess
import time
from pathlib import Path
from typing import Optional

from . import claude_env

logger = logging.getLogger(__name__)

PHOTO_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp"}
NOTE_EXTENSIONS = {".md", ".txt"}
INDEX_FILENAME = "index.json"


def analyze_references(project_path: Path, force: bool = False) -> Optional[dict]:
    """Analyze all reference photos and generate/update index.json.

    Args:
        project_path: Path to project root (contains references/)
        force: If True, re-analyze even if index is up-to-date

    Returns:
        Index dict, or None if no references directory exists
    """
    refs_path = project_path / "references"
    if not refs_path.exists():
        logger.warning(f"No references directory: {refs_path}")
        return None

    index_path = refs_path / INDEX_FILENAME

    # Collect current photos and notes
    photos = _collect_photos(refs_path)
    user_notes = _collect_notes(refs_path)

    if not photos:
        logger.warning("No photos found in references/")
        return None

    # Check if re-analysis needed
    existing_index = None
    if index_path.exists():
        try:
            existing_index = json.loads(index_path.read_text())
        except (json.JSONDecodeError, OSError) as e:
            logger.warning(f"Failed to read existing index: {e}")

    if not force and existing_index and not _needs_reanalysis(photos, existing_index):
        logger.info(f"Index is up-to-date ({len(photos)} photos)")
        return existing_index

    # Run analysis
    logger.info(f"Analyzing {len(photos)} reference photos...")
    print(f"Analyzing {len(photos)} reference photos...")

    index = _run_claude_analysis(refs_path, photos, user_notes)

    if index:
        # Save index
        index_path.write_text(json.dumps(index, indent=2, ensure_ascii=False))
        logger.info(f"Saved index: {index_path}")
        print(f"Index saved: {index_path}")
        print(f"  Photos analyzed: {len(index.get('photos', []))}")
    else:
        print("Analysis failed — see logs for details")

    return index


def _collect_photos(refs_path: Path) -> list[str]:
    """Collect all photo paths relative to references/."""
    photos = []
    for f in sorted(refs_path.rglob("*")):
        if f.is_file() and f.suffix.lower() in PHOTO_EXTENSIONS:
            photos.append(str(f.relative_to(refs_path)))
    return photos


def _collect_notes(refs_path: Path) -> dict[str, str]:
    """Collect all text/markdown notes, keyed by relative path."""
    notes = {}
    for f in sorted(refs_path.rglob("*")):
        if f.is_file() and f.suffix.lower() in NOTE_EXTENSIONS:
            rel = str(f.relative_to(refs_path))
            try:
                notes[rel] = f.read_text()
            except OSError:
                pass
    return notes


def _needs_reanalysis(current_photos: list[str], index_data: dict) -> bool:
    """Check if index is stale by comparing photo sets."""
    indexed_photos = {p["path"] for p in index_data.get("photos", [])}
    current_set = set(current_photos)

    if indexed_photos != current_set:
        added = current_set - indexed_photos
        removed = indexed_photos - current_set
        if added:
            logger.info(f"New photos: {added}")
        if removed:
            logger.info(f"Removed photos: {removed}")
        return True

    return False


def _build_analysis_prompt(refs_path: Path, photos: list[str], user_notes: dict) -> str:
    """Build the Claude prompt for photo analysis."""

    # Build photo list with categories inferred from path
    photo_entries = []
    for p in photos:
        parts = Path(p).parts
        category = parts[0] if len(parts) > 1 else "uncategorized"
        person = parts[1] if category == "people" and len(parts) > 2 else None
        photo_entries.append({
            "path": p,
            "full_path": str(refs_path / p),
            "category": category,
            "person": person,
        })

    # Group by category for clearer prompt
    by_category = {}
    for entry in photo_entries:
        cat = entry["category"]
        if cat not in by_category:
            by_category[cat] = []
        by_category[cat].append(entry)

    prompt = """You are analyzing reference photos for a video production project.

Read each photo listed below using the Read tool, then write a JSON index file.

## Photos to analyze

"""

    for cat, entries in by_category.items():
        prompt += f"### {cat}/\n"
        for e in entries:
            person_note = f" (person: {e['person']})" if e.get("person") else ""
            prompt += f"- `{e['full_path']}`{person_note}\n"
        prompt += "\n"

    if user_notes:
        prompt += "## User-provided descriptions\n\n"
        for note_path, content in user_notes.items():
            prompt += f"### {note_path}\n```\n{content}\n```\n\n"

    prompt += """## Instructions

For EACH photo, determine:

1. **type** — one of: `face_portrait` | `full_body` | `lifestyle` | `mood_style` | `landscape` | `location` | `object`
   - `face_portrait`: face clearly visible and is the main subject
   - `full_body`: full body visible, face may or may not be clear
   - `lifestyle`: person in context/activity, face may be hidden
   - `mood_style`: abstract mood/aesthetic reference
   - `landscape`: outdoor/nature scene without people
   - `location`: indoor/outdoor setting (could be used as background)
   - `object`: product, prop, or other object

2. **face_visible**: true/false — is a human face visible at all?

3. **face_clarity**: `high` | `medium` | `low` | `none`
   - high: clear, well-lit, front-facing or 3/4 view
   - medium: visible but partially obscured, side view, or small in frame
   - low: barely visible, very small, or heavily obscured
   - none: no face visible

4. **body_visible**: `full` | `upper` | `lower` | `full_back` | `partial` | `none`

5. **usefulness** (1-10): How useful is this photo for AI video generation?
   - 9-10: Clear face portrait with good lighting — best for face reference
   - 7-8: Good body/style reference with some face visibility
   - 5-6: Useful for style/mood but limited face reference value
   - 3-4: Minimal reference value (back view, tiny figure, etc.)
   - 1-2: Almost no reference value (blurry, irrelevant, etc.)

6. **description**: One-sentence description of what's in the photo. Include: pose, lighting, clothing, setting, distinguishing features.

## Output format

Write the result as a JSON file to: `""" + str(refs_path / INDEX_FILENAME) + """`

The JSON must follow this exact structure:
```json
{
  "analyzed_at": "<ISO timestamp>",
  "photos": [
    {
      "path": "<relative path from references/>",
      "category": "<inferred from directory>",
      "person": "<person name if under people/, else null>",
      "type": "<photo type>",
      "face_visible": true/false,
      "face_clarity": "high/medium/low/none",
      "body_visible": "full/upper/lower/full_back/partial/none",
      "usefulness": 1-10,
      "description": "<one sentence>"
    }
  ],
  "user_notes": {<notes dict>}
}
```

Important:
- Read EVERY photo listed above — do not skip any
- Use the RELATIVE path (from references/) in the "path" field, not the full path
- Sort photos by usefulness (highest first) within each category
- Be accurate about face visibility — this directly affects downstream video quality
"""

    return prompt


def _run_claude_analysis(refs_path: Path, photos: list[str], user_notes: dict) -> Optional[dict]:
    """Call Claude CLI to analyze photos and generate index."""

    prompt = _build_analysis_prompt(refs_path, photos, user_notes)
    index_path = refs_path / INDEX_FILENAME

    # Remove old index so Claude writes fresh
    if index_path.exists():
        index_path.unlink()

    start = time.time()

    try:
        result = subprocess.run(
            [
                "claude",
                "--allowedTools", "Read,Write",
                "-p", prompt,
            ],
            capture_output=True,
            text=True,
            cwd=str(refs_path.parent),  # project root
            timeout=600,
            env=claude_env(),
        )
    except subprocess.TimeoutExpired:
        logger.error("Claude analysis timed out (600s)")
        return None
    except FileNotFoundError:
        logger.error("Claude CLI not found — is it installed?")
        return None

    elapsed = time.time() - start
    logger.info(f"Claude analysis completed in {elapsed:.1f}s (exit code {result.returncode})")

    if result.returncode != 0:
        error_detail = result.stderr or result.stdout
        logger.error(f"Claude error output: {error_detail[:500] if error_detail else 'none'}")
        # Still try to read index — Claude may have written it before erroring

    if result.stdout:
        logger.debug(f"Claude stdout (last 300 chars): {result.stdout[-300:]}")

    # Read the index file Claude should have written
    if index_path.exists():
        try:
            index = json.loads(index_path.read_text())
            # Validate structure
            if "photos" not in index:
                logger.error("Index missing 'photos' key")
                return None
            if not isinstance(index["photos"], list):
                logger.error("Index 'photos' is not a list")
                return None
            logger.info(f"Index loaded: {len(index['photos'])} photos analyzed")
            return index
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse index.json: {e}")
            return None
    else:
        # Try to extract JSON from stdout as fallback
        logger.warning("Claude did not write index.json — trying stdout extraction")
        return _extract_json_from_output(result.stdout, user_notes)


def _extract_json_from_output(stdout: str, user_notes: dict) -> Optional[dict]:
    """Try to extract JSON index from Claude's stdout as fallback."""
    if not stdout:
        return None

    # Look for JSON block in output
    import re
    # Try to find a JSON object with "photos" key
    json_match = re.search(r'\{[^{}]*"photos"\s*:\s*\[.*?\]\s*(?:,\s*"user_notes"\s*:\s*\{.*?\})?\s*\}',
                          stdout, re.DOTALL)
    if json_match:
        try:
            return json.loads(json_match.group())
        except json.JSONDecodeError:
            pass

    # Try the whole stdout
    try:
        data = json.loads(stdout)
        if "photos" in data:
            return data
    except json.JSONDecodeError:
        pass

    return None


def get_best_photos(index: dict, category: str, person: str = None,
                    count: int = 3, min_usefulness: int = 0,
                    require_face: bool = False) -> list[dict]:
    """Get top N photos from index, sorted by usefulness.

    Args:
        index: The index dict (from analyze_references or index.json)
        category: Category to filter by ("people", "locations", "style", etc.)
        person: Person name to filter by (only for category="people")
        count: Max number of photos to return
        min_usefulness: Minimum usefulness rating to include
        require_face: If True, only return photos where face_visible=True

    Returns:
        List of photo entry dicts, sorted by usefulness (highest first)
    """
    if not index or "photos" not in index:
        return []

    filtered = []
    for photo in index["photos"]:
        if photo.get("category") != category:
            continue
        if person is not None and photo.get("person") != person:
            continue
        if photo.get("usefulness", 0) < min_usefulness:
            continue
        if require_face and not photo.get("face_visible", False):
            continue
        filtered.append(photo)

    # Sort by usefulness descending
    filtered.sort(key=lambda p: p.get("usefulness", 0), reverse=True)

    return filtered[:count]


def _load_char_description(image_path: Path) -> Optional[str]:
    """Load char_description.txt from the same directory as a characteristic shot.

    Returns the description text (whitespace-normalized to single line) or None.
    """
    desc_file = image_path.parent / "char_description.txt"
    if desc_file.exists():
        text = desc_file.read_text().strip()
        # Normalize to compact form for prefix
        return " ".join(text.split())
    return None


def build_analyze_prefix(index: Optional[dict], reference_image_paths: list[str],
                        refs_path: Optional[Path] = None,
                        character_labels: Optional[list] = None) -> str:
    """Build [Analyze...] prefix block for NanoBanana prompts.

    Instructs NanoBanana to analyze attached reference images for face/appearance
    matching while generating a completely fresh scene composition. Also enforces
    sharp, in-focus output.

    Args:
        index: The reference index dict (may be None if no index)
        reference_image_paths: List of absolute paths to reference images being sent
        refs_path: Path to references/ directory (for resolving relative paths)
        character_labels: Optional list parallel to reference_image_paths. Each entry
            is either a dict {"character_id": str, "label": str} or None (for
            non-character images like characteristic shots). When provided, builds
            a multi-character analyze prefix.

    Returns:
        Prefix string to prepend to NanoBanana prompt, or empty string if no refs
    """
    if not reference_image_paths:
        return ""

    # Multi-character mode
    if character_labels:
        return _build_multi_character_prefix(
            index, reference_image_paths, refs_path, character_labels
        )

    # Single-person mode (original behavior)
    # Build description of what each reference image contains
    image_descriptions = []
    for i, img_path in enumerate(reference_image_paths, 1):
        p = Path(img_path)

        # Check for characteristic shot with description file
        if "characteristic_shot" in p.stem:
            char_desc = _load_char_description(p)
            if char_desc:
                image_descriptions.append(
                    f"Image {i}: characteristic shot (visual style and character reference). "
                    f"SCENE CONTENTS: {char_desc}"
                )
                continue

        analysis = None
        if index and refs_path:
            # Try to find this image in the index
            try:
                rel = str(Path(img_path).relative_to(refs_path))
            except ValueError:
                rel = img_path
            analysis = get_photo_analysis(index, rel)

        if analysis:
            desc = analysis.get("description", "")
            img_type = analysis.get("type", "reference")
            face_note = ""
            if analysis.get("face_visible"):
                face_note = f", face clarity: {analysis.get('face_clarity', 'unknown')}"
            image_descriptions.append(f"Image {i}: {img_type}{face_note}. {desc}")
        else:
            # No index data — generic description based on filename
            name = Path(img_path).name
            image_descriptions.append(f"Image {i}: {name}")

    descriptions_block = " ".join(image_descriptions)

    prefix = (
        f"[Analyze image inputs for the person's face, facial features, hair, "
        f"body type, and overall appearance. {descriptions_block} "
        f"Use these as strict reference for character likeness — the generated "
        f"person must match the reference photos. "
        f"RESET SCENE GEOMETRY completely. Do not edit or add to the previous "
        f"image's composition. Create a completely fresh framing and camera "
        f"perspective to depict the specific scene described below. "
        f"Output must be tack-sharp with crisp focus on the subject — no soft "
        f"focus, no blur, no dreamy haze, no gaussian glow. Razor-sharp facial "
        f"details, sharp eyes, sharp hair strands, sharp fabric texture.] "
    )

    return prefix


def _build_multi_character_prefix(
    index: Optional[dict],
    reference_image_paths: list[str],
    refs_path: Optional[Path],
    character_labels: list,
) -> str:
    """Build [Analyze...] prefix for multiple named characters.

    Each image is labelled with its character name and role so NanoBanana
    can match the right face to the right character in the scene.
    """
    image_descriptions = []
    for i, (img_path, label) in enumerate(
        zip(reference_image_paths, character_labels), 1
    ):
        if label and label.get("type") == "corpus_ref":
            # Corpus reference image (composition guide)
            scene_num = label.get("scene", "?")
            desc = label.get("label", "")
            desc_note = f" — {desc[:120]}" if desc else ""
            image_descriptions.append(
                f"Image {i}: Frame {scene_num} COMPOSITION REFERENCE — "
                f"Adapt this reference's composition for Frame {scene_num}. "
                f"Preserve spatial arrangement and visual metaphor, "
                f"update characters{desc_note}."
            )
        elif label and label.get("type") == "frame_reference":
            # Episode frame reference (from frame index fallback)
            asset_id = label.get("asset_id", "unknown")
            asset_name = asset_id.replace("_", " ").title()
            desc = label.get("label", "")
            desc_note = f" — {desc[:120]}" if desc else ""
            image_descriptions.append(
                f"Image {i}: {asset_name} frame reference "
                f"(match this environment/object appearance){desc_note}."
            )
        elif label and label.get("character_id"):
            # Named character reference
            char_id = label["character_id"]
            char_name = char_id.replace("_", " ").title()
            role_note = f" ({label['label']})" if label.get("label") else ""
            img_type = Path(img_path).stem  # e.g. "sheet", "portrait"
            image_descriptions.append(
                f"Image {i}: {char_name} character {img_type} "
                f"(face/body/clothing reference) ({char_id}){role_note}."
            )
        else:
            # Non-character image — infer role from path
            p = Path(img_path)
            parent_name = p.parent.name
            grandparent_name = p.parent.parent.name if p.parent.parent != p.parent else ""

            if "characteristic_shot" in p.stem:
                char_desc = _load_char_description(p)
                if char_desc:
                    role_label = (
                        f"characteristic shot (visual style and character reference). "
                        f"SCENE CONTENTS: {char_desc}"
                    )
                else:
                    role_label = "characteristic shot (overall visual style and composition reference)"
            elif parent_name == "locations" or grandparent_name == "locations":
                asset_id = p.stem
                asset_name = asset_id.replace("_", " ").title()
                role_label = f"{asset_name} location reference (match this environment appearance) ({asset_id})"
            elif parent_name == "objects" or grandparent_name == "objects":
                asset_id = p.stem
                asset_name = asset_id.replace("_", " ").title()
                role_label = f"{asset_name} object reference (match this object appearance exactly) ({asset_id})"
            elif parent_name == "style" or grandparent_name == "style":
                role_label = f"style reference ({p.stem}) (match this visual style)"
            else:
                role_label = f"reference ({p.stem})"

            analysis = None
            if index and refs_path:
                try:
                    rel = str(Path(img_path).relative_to(refs_path))
                except ValueError:
                    rel = img_path
                analysis = get_photo_analysis(index, rel)

            if analysis:
                desc = analysis.get("description", "")
                image_descriptions.append(
                    f"Image {i}: {role_label}. {desc}"
                )
            else:
                image_descriptions.append(
                    f"Image {i}: {role_label}."
                )

    descriptions_block = " ".join(image_descriptions)

    prefix = (
        f"[REFERENCE IMAGES: {descriptions_block} "
        f"All reference images are STYLE REFERENCES — use them to match "
        f"character faces, clothing, objects, and locations but do NOT copy "
        f"their composition. Each character must match their specific reference "
        f"sheet (correct face, hair, body type, clothing). Each object and "
        f"location must visually match its reference image. Do NOT mix features "
        f"between different references. "
        f"RESET SCENE GEOMETRY completely. Do not edit or add to the previous "
        f"image's composition. Create a completely fresh framing and camera "
        f"perspective to depict the specific scene described below. "
        f"Output must be tack-sharp with crisp focus on the subject — no soft "
        f"focus, no blur, no dreamy haze, no gaussian glow. Razor-sharp facial "
        f"details, sharp eyes, sharp hair strands, sharp fabric texture.] "
    )

    return prefix


def get_photo_analysis(index: dict, photo_path: str) -> Optional[dict]:
    """Look up analysis for a specific photo by its path.

    Args:
        index: The index dict
        photo_path: Relative path from references/ (e.g., "people/main/1.jpg")
            or absolute path (will try to match the filename)

    Returns:
        Photo entry dict or None
    """
    if not index or "photos" not in index:
        return None

    # Try exact match first
    for photo in index["photos"]:
        if photo["path"] == photo_path:
            return photo

    # Try matching by filename (for absolute paths)
    filename = Path(photo_path).name
    for photo in index["photos"]:
        if Path(photo["path"]).name == filename:
            return photo

    return None
