"""Gemini LLM provider — google-genai SDK, Gemini 3.1 Flash (analysis + image generation)."""

import json
import logging
from typing import Any

from google import genai
from google.genai import types
from pydantic import BaseModel

from app.config import settings
from app.services.llm.base import (
    BaseLLMProvider,
    GarmentAnalysis,
    StyleAnalysis,
    StyleAnalysisTemp,
)
from app.services.llm.prompts import (
    GARMENT_ANALYSIS_PROMPT,
    GHOST_MANNEQUIN_EDIT_PROMPT,
    STYLE_ANALYSIS_PROMPT,
    STYLE_ANALYSIS_TEMP,
    VTON_PROMPT,
)

logger = logging.getLogger(__name__)

# Analysis (JSON structured output) — gemini-2.5-flash supports vision + JSON mode
GEMINI_ANALYSIS_MODEL = "gemini-2.5-flash"
# Image generation — gemini-2.5-flash-image is optimized for image gen/editing
GEMINI_IMAGE_MODEL = "gemini-2.5-flash-image"


class _GarmentAnalysisList(BaseModel):
    """Wrapper schema to enforce structured list output from Gemini."""

    garments: list[GarmentAnalysis]


class GeminiProvider(BaseLLMProvider):
    """Gemini API provider using the official google-genai SDK."""

    def __init__(self) -> None:
        self.api_key = settings.GEMINI_API_KEY
        if not self.api_key:
            logger.warning("GEMINI_API_KEY is not set — Gemini calls will fail")
        self.client = genai.Client(api_key=self.api_key)

    async def _analyze(
        self, prompt: str, image: bytes, schema: type[BaseModel] | None = None
    ) -> Any:
        """Send a vision analysis request and return parsed JSON."""
        config = types.GenerateContentConfig(
            max_output_tokens=4096,
            response_mime_type="application/json",
        )
        if schema:
            config.response_json_schema = schema.model_json_schema()

        response = await self.client.aio.models.generate_content(
            model=GEMINI_ANALYSIS_MODEL,
            contents=[
                prompt,
                types.Part.from_bytes(data=image, mime_type="image/jpeg"),
            ],
            config=config,
        )
        return json.loads(response.text)

    async def analyze_garment(self, image: bytes) -> list[GarmentAnalysis]:
        """Analyze garments in image using Gemini Vision."""
        parsed = await self._analyze(
            GARMENT_ANALYSIS_PROMPT, image, schema=_GarmentAnalysisList
        )

        # Schema enforces {"garments": [...]}, extract the list
        if isinstance(parsed, dict) and "garments" in parsed:
            return [GarmentAnalysis(**item) for item in parsed["garments"]]

        # Fallback for plain list
        if isinstance(parsed, list):
            return [GarmentAnalysis(**item) for item in parsed]

        logger.error("Unexpected Gemini response format: %s", type(parsed))
        return []

    async def generate_ghost_mannequin(
        self,
        image: bytes,
        item: GarmentAnalysis,
        mode: str,
        image_model: str | None = None,
    ) -> bytes | None:
        """Generate ghost mannequin using Gemini image generation.

        ``image_model`` is accepted for signature compatibility with the
        OpenAI provider but ignored — Gemini always uses ``GEMINI_IMAGE_MODEL``.
        """
        try:
            prompt = GHOST_MANNEQUIN_EDIT_PROMPT.format(
                category_main=item.category_main,
                category_sub=item.category_sub or item.category_main,
            )

            response = await self.client.aio.models.generate_content(
                model=GEMINI_IMAGE_MODEL,
                contents=[
                    prompt,
                    types.Part.from_bytes(data=image, mime_type="image/jpeg"),
                ],
                config=types.GenerateContentConfig(
                    response_modalities=["IMAGE", "TEXT"],
                ),
            )

            # Extract image from response parts
            if not response.candidates:
                logger.warning("Gemini image edit returned no candidates")
                return None

            for part in response.candidates[0].content.parts:
                if part.inline_data and part.inline_data.data:
                    return part.inline_data.data

            logger.warning("Gemini image edit response contained no image data")
            return None
        except Exception:
            logger.exception("Ghost mannequin generation failed (gemini)")
            return None

    async def generate_vton(
        self,
        person_image: bytes,
        ghost_images: list[bytes],
        categories: list[str],
        image_model: str | None = None,
    ) -> tuple[bytes | None, str | None]:
        """Composite ghost garments onto a person via Gemini multi-image generation.

        ``image_model`` is accepted for signature compatibility with the
        OpenAI provider but ignored — Gemini always uses ``GEMINI_IMAGE_MODEL``.
        """
        catalog = ", ".join(
            f"#{i + 1}={cat}" for i, cat in enumerate(categories)
        )
        prompt = VTON_PROMPT.format(
            garment_count=len(ghost_images),
            catalog=catalog,
        )

        contents: list = [prompt, types.Part.from_bytes(data=person_image, mime_type="image/jpeg")]
        for g in ghost_images:
            contents.append(types.Part.from_bytes(data=g, mime_type="image/png"))

        try:
            response = await self.client.aio.models.generate_content(
                model=GEMINI_IMAGE_MODEL,
                contents=contents,
                config=types.GenerateContentConfig(
                    response_modalities=["IMAGE", "TEXT"],
                ),
            )
        except Exception:
            logger.exception("VTON generation failed (gemini)")
            return None, None

        if not response.candidates:
            logger.warning("Gemini VTON returned no candidates")
            return None, None

        image_bytes: bytes | None = None
        note_chunks: list[str] = []
        for part in response.candidates[0].content.parts:
            if part.inline_data and part.inline_data.data:
                image_bytes = part.inline_data.data
            elif getattr(part, "text", None):
                note_chunks.append(part.text)

        note = " ".join(chunk.strip() for chunk in note_chunks if chunk.strip()) or None
        if image_bytes is None:
            logger.warning("Gemini VTON response contained no image data")
        return image_bytes, note

    async def analyze_style(self, image: bytes) -> StyleAnalysis:
        """Analyze overall outfit style using Gemini Vision."""
        parsed = await self._analyze(
            STYLE_ANALYSIS_PROMPT, image, schema=StyleAnalysis
        )
        return StyleAnalysis(**parsed)

    async def analyze_style_temp(self, image: bytes) -> StyleAnalysisTemp:
        """Context-aware outfit style analysis using Gemini Vision."""
        parsed = await self._analyze(
            STYLE_ANALYSIS_TEMP, image, schema=StyleAnalysisTemp
        )
        return StyleAnalysisTemp(**parsed)