#!/usr/bin/env python3
"""
scene_keyframe.py — Generate a single scene keyframe (NBP still) with one
or more reference images.

Composite still generator: takes a list of reference image paths (driver,
car, location, prop — anything) and a text prompt, and generates a single
NBP still that follows the prompt while locking identity to the refs.

Unlike prep_character_angles.py (character-specific angles), this tool is
subject-agnostic and is the canonical tool for building per-shot start
frames when you already have picked refs.

Usage:
    python -m tools.scene_keyframe \\
        --project driver-beware \\
        --slug v1_stay_in_car_start \\
        --refs refs/characters/new_driver_01/picks/frontal.png,refs/props/new_car_01/picks/frontal.png,refs/locations/dense_city/hero.png \\
        --aspect 16:9 \\
        --description "Wide establishing shot: the pictured driver sits inside the pictured car..."

Output: projects/{project}/sequences/ep_001/clean/{slug}.png
Cost: ~$0.134 per keyframe (single NBP call).
"""

import argparse
import json
import logging
import sys
from pathlib import Path

_PROJECT_ROOT = Path(__file__).parent.parent
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))

from recoil.core.paths import ProjectPaths
from recoil.core.model_profiles import get_model

logger = logging.getLogger("pipeline.scene_keyframe")

NBP_COST = 0.134


def generate_scene_keyframe(
    slug: str,
    description: str,
    ref_paths: list[Path],
    project: str,
    aspect_ratio: str = "16:9",
    episode: str = "ep_001",
    dry_run: bool = False,
) -> dict:
    """Generate a scene keyframe still via NBP with ref images.

    Args:
        slug: short identifier used in the output filename
        description: freeform prompt describing the shot
        ref_paths: list of Path objects pointing at ref images
        project: project name
        aspect_ratio: output aspect ratio (default 16:9)
        episode: episode slug for output dir (default ep_001)
        dry_run: print prompt + refs and exit

    Returns:
        dict with output_path, cost, and status.
    """
    # `episode` is a slug like "ep_001"; extract the integer for ProjectPaths.
    try:
        _ep_int = int(str(episode).replace("ep_", "").lstrip("0") or "0")
    except (TypeError, ValueError):
        _ep_int = 0
    out_dir = ProjectPaths.for_project(project).episode_prep_dir(_ep_int) / "clean"
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / f"{slug}.png"

    for p in ref_paths:
        if not p.is_file():
            return {
                "slug": slug,
                "output_path": None,
                "cost": 0.0,
                "error": f"Ref not found: {p}",
            }

    ref_list = "\n".join(f"  - {p.name}" for p in ref_paths)
    prompt = (
        f"REFERENCE IMAGES: {len(ref_paths)} attached — use them as the "
        f"source of truth for identity, color, shape, and style.\n"
        f"{ref_list}\n\n"
        f"SHOT BRIEF: {description}\n\n"
        f"CRITICAL: Maintain the EXACT visual style of the attached references "
        f"(2D cartoon, clean line art, flat colors, no photorealism). Do not "
        f"invent new characters or cars — use only what is shown in the refs. "
        f"Compose the scene exactly as described in the shot brief.\n\n"
        f"NEGATIVE: no photorealism, no extra characters, no text overlays, "
        f"no watermarks, no borders, no film frames."
    )

    if dry_run:
        print(f"[DRY RUN] Would generate scene keyframe: {out_path}")
        print(f"  Aspect: {aspect_ratio}")
        print(f"  Refs ({len(ref_paths)}):")
        for p in ref_paths:
            print(f"    {p}")
        print(f"  Prompt:\n{prompt}")
        return {"slug": slug, "output_path": None, "cost": 0.0, "dry_run": True}

    from recoil.execution.api_client import get_client
    from google.genai import types as genai_types

    model_id = get_model("production_hq", "image")
    client = get_client(model_id)

    # Build content parts: one Part per ref image, then the text prompt.
    content_parts = []
    for ref in ref_paths:
        data = ref.read_bytes()
        mime = "image/jpeg" if ref.suffix.lower() in {".jpg", ".jpeg"} else "image/png"
        content_parts.append(genai_types.Part.from_bytes(data=data, mime_type=mime))
    content_parts.append(prompt)

    config = genai_types.GenerateContentConfig(
        temperature=0.3,  # Low temp — want identity lock, not variety
        response_modalities=["IMAGE", "TEXT"],
        image_config=genai_types.ImageConfig(aspect_ratio=aspect_ratio),
    )

    try:
        api_client = client._get_client()
        response = api_client.models.generate_content(
            model=model_id,
            contents=content_parts,
            config=config,
        )
    except Exception as e:
        logger.error("Scene keyframe generation failed for %s: %s", slug, e)
        return {"slug": slug, "output_path": None, "cost": NBP_COST, "error": str(e)}

    if response and response.candidates:
        for candidate in response.candidates:
            if candidate.content and candidate.content.parts:
                for part in candidate.content.parts:
                    if hasattr(part, "inline_data") and part.inline_data:
                        out_path.write_bytes(part.inline_data.data)
                        logger.info("Scene keyframe saved: %s", out_path)
                        break

    if not out_path.exists():
        return {
            "slug": slug,
            "output_path": None,
            "cost": NBP_COST,
            "error": "NBP returned no image data",
        }

    return {"slug": slug, "output_path": str(out_path), "cost": NBP_COST}


def main():
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%H:%M:%S",
    )
    parser = argparse.ArgumentParser(description="NBP scene keyframe generator")
    parser.add_argument("--project", required=True)
    parser.add_argument("--slug", required=True, help="Output filename slug")
    parser.add_argument(
        "--refs",
        required=True,
        help="Comma-separated ref image paths (relative to project dir OR absolute)",
    )
    parser.add_argument("--description", required=True, help="Shot brief prompt")
    parser.add_argument("--aspect", default="16:9", help="Aspect ratio (default 16:9)")
    parser.add_argument("--episode", default="ep_001", help="Episode slug (default ep_001)")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    # Resolve ref paths: accept absolute or project-relative
    from recoil.core.paths import projects_root
    project_dir = projects_root() / args.project
    ref_paths = []
    for r in args.refs.split(","):
        r = r.strip()
        p = Path(r)
        if not p.is_absolute():
            p = project_dir / p
        ref_paths.append(p)

    result = generate_scene_keyframe(
        slug=args.slug,
        description=args.description,
        ref_paths=ref_paths,
        project=args.project,
        aspect_ratio=args.aspect,
        episode=args.episode,
        dry_run=args.dry_run,
    )
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()