#!/usr/bin/env python3
"""
screen_test_gen.py — Image generation for the Screen Test feature.

Generates character images for each wardrobe phase using NBP
(Gemini 3 Pro Image) with the cast hero as reference.

Functions:
  - build_phase_prompt: Builds generation prompt from character + phase data
  - enrich_director_note: Translates natural-language note into image prompt via Flash
  - generate_phase_image: Generates a single phase image via NBP with references
"""

import logging
import sys
from pathlib import Path
from typing import Optional

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

# Optional imports — allow tests to run without google.genai installed
try:
    from google.genai import types as genai_types
    _HAS_GENAI = True
except ImportError:
    genai_types = None
    _HAS_GENAI = False

try:
    from google import genai
    _HAS_GENAI_CLIENT = True
except ImportError:
    genai = None
    _HAS_GENAI_CLIENT = False

from recoil.core.model_profiles import get_model
from recoil.core.prompt_config import get_constant

logger = logging.getLogger(__name__)


# ── Private client helpers ───────────────────────────────────────────


def _get_text_client():
    """Return a genai.Client() for text-only calls (Flash enrichment).

    Uses default credentials / GEMINI_API_KEY from environment.
    """
    if not _HAS_GENAI_CLIENT:
        raise RuntimeError(
            "google-genai SDK not installed. Install with: pip install google-genai"
        )
    return genai.Client()


def _get_image_client():
    """Return the raw genai client for NBP image generation.

    Uses lib.api_client.get_client() to get the GoogleGenaiClient,
    then unwraps via _get_client() for direct API access.
    """
    from recoil.execution.api_client import get_client

    client = get_client(get_model("production", "image"))
    return client._get_client()


# ── build_phase_prompt ───────────────────────────────────────────────


DEFAULT_TONE = (
    "Favor muted, naturalistic tones over saturated primaries. "
    "Let texture and fit tell the story, not color alone."
)


def build_phase_prompt(
    char: dict,
    phase: dict,
    enriched_note: Optional[str] = None,
    props: Optional[list] = None,
    aesthetic_directives: Optional[dict] = None,
) -> str:
    """Build a generation prompt from character visual description + phase details.

    Args:
        char: Character dict with at least 'visual_description'.
        phase: Phase dict with wardrobe, hair, makeup, marks fields.
        enriched_note: Optional enriched director's note to include.
        props: Optional list of prop dicts with 'description' and 'state_notes'.
        aesthetic_directives: Optional dict from global bible with show-level
            aesthetic preferences (e.g. 'costume_palette_direction').

    Returns:
        Formatted photographic prompt string for NBP generation.
    """
    visual_desc = char.get("visual_description", "")
    wardrobe = phase.get("wardrobe", "")
    hair = phase.get("hair", "")
    makeup = phase.get("makeup", "")
    marks = phase.get("marks", "")

    # Build the subject description block
    sections = [
        f"SUBJECT: {visual_desc}",
    ]

    if wardrobe:
        sections.append(f"WARDROBE: {wardrobe}")
    if hair:
        sections.append(f"HAIR: {hair}")
    if makeup:
        sections.append(f"MAKEUP: {makeup}")
    if marks:
        sections.append(f"MARKS/DISTINGUISHING FEATURES: {marks}")

    # Props block — wardrobe props visible on or carried by the character
    if props:
        prop_descs = []
        for p in props:
            desc = p.get("description", "")
            state = p.get("state_notes", "")
            if desc:
                prop_descs.append(f"{desc} ({state})" if state else desc)
        if prop_descs:
            sections.append(f"PROPS/ACCESSORIES: {'; '.join(prop_descs)}")

    subject_block = "\n".join(sections)

    # Director's note block
    note_block = ""
    if enriched_note:
        note_block = f"\n\nDIRECTOR'S NOTE: {enriched_note}"

    # Show-level tone from bible — global creative direction for the whole show
    directives = aesthetic_directives or {}
    tone = directives.get("tone", DEFAULT_TONE)

    prompt = (
        f"Generate a full-body 9:16 portrait photograph of this character.\n"
        f"Match the reference images EXACTLY — same face, same identity, same person.\n\n"
        f"{subject_block}"
        f"{note_block}\n\n"
        f"COSTUME DESIGN PRINCIPLES:\n"
        f"- Apply the 60-30-10 color rule: 60% dominant neutral (base garment), "
        f"30% secondary complementary tone (jacket, vest, or outer layer), "
        f"10% accent color (accessories, trim, belt, insignia).\n"
        f"- All wardrobe pieces, props, and accessories must feel like they belong to "
        f"the same world — unified material palette, consistent wear/aging, coherent era.\n"
        f"- TONE: {tone}\n\n"
        f"PHOTOGRAPHIC ANCHORS:\n"
        f"- Medium: 35mm motion picture film still, shot on {get_constant('casting', 'casting_camera')}.\n"
        f"- Lighting: Soft, directional studio lighting. {get_constant('casting', 'casting_background')}.\n"
        f"- Framing: Full-body 9:16 portrait, centered subject, head-to-toe.\n"
        f"- Texture: {get_constant('casting', 'casting_texture_human')}\n"
        f"- Negative constraints: {get_constant('casting', 'casting_anti_airbrush')}"
    )

    return prompt


# ── enrich_director_note ─────────────────────────────────────────────


def enrich_director_note(base_description: str, director_note: str) -> str:
    """Translate a natural-language director's note into optimized image prompt.

    Uses Gemini Flash (cheap, fast) to convert narrative direction
    into concrete visual descriptors suitable for image generation.

    Args:
        base_description: Character's base visual description for context.
        director_note: Natural-language direction from the director.

    Returns:
        Enriched description string with concrete visual descriptors.
    """
    client = _get_text_client()

    prompt = (
        f"You are a visual prompt engineer for AI image generation.\n\n"
        f"CHARACTER DESCRIPTION:\n{base_description}\n\n"
        f"DIRECTOR'S NOTE:\n{director_note}\n\n"
        f"Translate the director's note into concrete, visual image-generation descriptors. "
        f"Focus on what the camera would SEE: body language, facial expression, "
        f"physical state (sweat, dirt, blood, tension), posture, gesture. "
        f"Return ONLY the visual descriptors as a single line, no preamble."
    )

    response = client.models.generate_content(
        model=get_model("flash", "text"),
        contents=prompt,
    )

    return response.text.strip()


# ── generate_phase_image ─────────────────────────────────────────────


def generate_phase_image(
    hero_path: Path,
    three_quarter_path: Optional[Path],
    prompt: str,
    output_path: Path,
    anchor_path: Optional[Path] = None,
) -> bool:
    """Generate a single phase image via NBP with reference images.

    Args:
        hero_path: Path to the hero (front) reference image.
        three_quarter_path: Optional path to three-quarter reference image.
        prompt: The full generation prompt (from build_phase_prompt).
        output_path: Where to save the generated image.
        anchor_path: Optional path to anchor phase image for cross-phase consistency.

    Returns:
        True on success, False on failure.
    """
    try:
        api_client = _get_image_client()

        # Assemble reference images as parts
        contents = []

        hero_bytes = hero_path.read_bytes()
        hero_part = genai_types.Part.from_bytes(data=hero_bytes, mime_type="image/png")
        contents.append(hero_part)

        if three_quarter_path is not None:
            tq_bytes = three_quarter_path.read_bytes()
            tq_part = genai_types.Part.from_bytes(data=tq_bytes, mime_type="image/png")
            contents.append(tq_part)

        if anchor_path is not None:
            anchor_bytes = anchor_path.read_bytes()
            anchor_part = genai_types.Part.from_bytes(data=anchor_bytes, mime_type="image/png")
            contents.append(anchor_part)

        # Text prompt goes last (recency bias ordering)
        contents.append(prompt)

        # API config: temperature=0.6, 9:16 portrait
        config = genai_types.GenerateContentConfig(
            temperature=0.6,
            response_modalities=["IMAGE", "TEXT"],
            image_config=genai_types.ImageConfig(aspect_ratio="9:16"),
        )

        response = api_client.models.generate_content(
            model=get_model("production", "image"),
            contents=contents,
            config=config,
        )

        # Extract image from response
        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:
                            # Create output directories if needed
                            output_path.parent.mkdir(parents=True, exist_ok=True)
                            output_path.write_bytes(part.inline_data.data)
                            logger.info("Phase image saved: %s", output_path)
                            return True

        logger.warning("No image data in NBP response for: %s", output_path)
        return False

    except Exception as e:
        logger.error("Phase image generation failed: %s", e)
        return False