"""
frame_editor.py — Mask-free image editing via Gemini generate_content.

Replaces the generate-from-scratch extraction pipeline. Instead of building
refs + prompt and generating a new image, we pass the hero frame + an edit
instruction to Gemini. The model modifies the existing image rather than
generating from scratch, preserving background pixels.

Cost: same as NBP generation (~$0.134/frame for Pro, ~$0.039 for Flash).
"""

import io
import logging
import os
from pathlib import Path
from typing import Optional

import numpy as np
from PIL import Image

logger = logging.getLogger(__name__)

# Cost estimate for a single edit call (same as NBP generation)
EDIT_COST = 0.134


def edit_hero_pose(
    hero_image_path: Path,
    edit_instruction: str,
    model: str | None = None,
) -> dict:
    """Edit a hero frame to change the character's pose via Gemini mask-free edit.

    Uses the same generate_content endpoint as NBP generation, but with a
    different payload structure: [image, edit_instruction] instead of
    [ref_images..., scene_prompt]. This triggers image-to-image translation
    rather than from-scratch generation.

    Args:
        hero_image_path: Path to the hero/anchor frame to edit.
        edit_instruction: Natural language description of the pose change.
        model: Override model ID. Defaults to gemini-3-pro-image-preview.

    Returns:
        {"success": True, "image_data": bytes, "cost": float}
        or {"success": False, "error": str}
    """
    try:
        from google import genai
        from google.genai import types as genai_types

        api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
        if not api_key:
            return {"success": False, "error": "GEMINI_API_KEY not set"}

        if not hero_image_path.is_file():
            return {"success": False, "error": f"Hero image not found: {hero_image_path}"}

        client = genai.Client(api_key=api_key)

        # Load hero image as PIL Image
        hero_image = Image.open(hero_image_path)

        # Edit payload: [image, instruction] triggers edit mode
        config = genai_types.GenerateContentConfig(
            temperature=0.3,
            responseModalities=["IMAGE", "TEXT"],
            imageConfig=genai_types.ImageConfig(
                aspectRatio="9:16",
            ),
        )

        from recoil.core.model_profiles import get_model
        edit_model = model or get_model("production", "image")

        response = client.models.generate_content(
            model=edit_model,
            contents=[hero_image, edit_instruction],
            config=config,
        )

        # Extract image from response
        image_data = None
        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:
                            image_data = part.inline_data.data

        if image_data:
            return {"success": True, "image_data": image_data, "cost": EDIT_COST}

        return {"success": False, "error": "No image in edit response"}

    except Exception as e:
        return {"success": False, "error": str(e)}


def validate_companion(
    hero_bytes: bytes,
    companion_bytes: bytes,
    threshold: float = 0.70,
) -> dict:
    """Validate that a companion frame preserved the hero's background.

    Compares color histograms between hero and companion. High correlation
    means similar color distribution = similar background/lighting.

    Args:
        hero_bytes: Raw bytes of the hero/anchor frame.
        companion_bytes: Raw bytes of the generated companion frame.
        threshold: Minimum correlation score to pass (0.0-1.0).

    Returns:
        {"passed": bool, "correlation": float}
    """
    try:
        hero_img = Image.open(io.BytesIO(hero_bytes)).convert("RGB")
        comp_img = Image.open(io.BytesIO(companion_bytes)).convert("RGB")

        hero_hist = np.array(hero_img.histogram(), dtype=np.float64)
        comp_hist = np.array(comp_img.histogram(), dtype=np.float64)

        hero_hist = hero_hist / (hero_hist.sum() + 1e-10)
        comp_hist = comp_hist / (comp_hist.sum() + 1e-10)

        hero_mean = hero_hist.mean()
        comp_mean = comp_hist.mean()
        numerator = np.sum((hero_hist - hero_mean) * (comp_hist - comp_mean))
        denominator = np.sqrt(
            np.sum((hero_hist - hero_mean) ** 2) * np.sum((comp_hist - comp_mean) ** 2)
        )
        correlation = float(numerator / (denominator + 1e-10))

        return {
            "passed": correlation >= threshold,
            "correlation": round(correlation, 4),
        }

    except Exception as e:
        logger.warning("Companion validation failed: %s", e)
        return {"passed": False, "correlation": 0.0}


def build_edit_instruction(
    target_type: str,
    pose_description: str,
) -> str:
    """Build the edit instruction string for Gemini mask-free editing.

    Args:
        target_type: "anticipation" or "aftermath"
        pose_description: Description of the target pose from Flash text call.

    Returns:
        The full edit instruction string.
    """
    if target_type == "anticipation":
        return (
            f"Edit this image: change the character's body position and pose to show "
            f"the moment BEFORE the current action — {pose_description}. "
            f"Keep the background, lighting, environment, camera angle, walls, floor, "
            f"and all non-character elements completely identical. "
            f"Keep the character's clothing, hair style, and appearance the same. "
            f"Only change the character's body position, pose, and expression."
        )
    else:
        return (
            f"Edit this image: change the character's body position and pose to show "
            f"the moment AFTER the current action — {pose_description}. "
            f"Keep the background, lighting, environment, camera angle, walls, floor, "
            f"and all non-character elements completely identical. "
            f"Keep the character's clothing, hair style, and appearance the same. "
            f"Only change the character's body position, pose, and expression."
        )