"""Gemini-powered style extraction helpers."""

from __future__ import annotations

import json
import mimetypes
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Any

import httpx
from google import genai
from google.genai import types

from nanoslides.core.provider_errors import is_service_unavailable_error
from nanoslides.core.style import ProjectStyleConfig

_STYLE_ANALYSIS_PRIMARY_MODEL = "gemini-3-pro-preview"
_STYLE_ANALYSIS_FALLBACK_MODEL = "gemini-2.5-pro"


class StyleStealSourceKind(str, Enum):
    """Supported style-steal source kinds."""

    IMAGE = "image"


@dataclass(frozen=True)
class StyleStealSource:
    """Normalized style-steal source payload."""

    kind: StyleStealSourceKind
    path: Path
    mime_type: str
    bytes_data: bytes


@dataclass(frozen=True)
class StyleStealSuggestion:
    """Structured style suggestion generated by Gemini 3 Pro."""

    base_prompt: str
    negative_prompt: str
    reference_comments: list[str]
    use_as_base_reference: bool
    base_reference_reason: str


@dataclass(frozen=True)
class InferredProjectStyle:
    """In-memory style inference output that callers may choose to persist."""

    project_style: ProjectStyleConfig
    suggestion: StyleStealSuggestion
    reference_sources: list[StyleStealSource]


class GeminiStyleStealAnalyzer:
    """Analyze a source asset and infer reusable slide style instructions."""

    def __init__(self, *, api_key: str, timeout_seconds: float = 120.0) -> None:
        if timeout_seconds <= 0:
            raise ValueError("timeout_seconds must be greater than zero.")
        timeout_milliseconds = max(1, int(timeout_seconds * 1000))
        self._client = genai.Client(
            api_key=api_key,
            http_options=types.HttpOptions(
                api_version="v1alpha",
                timeout=timeout_milliseconds,
                retry_options=types.HttpRetryOptions(attempts=2),
            ),
        )
        self.last_model_used = _STYLE_ANALYSIS_PRIMARY_MODEL

    def analyze(self, source: StyleStealSource) -> StyleStealSuggestion:
        if source.kind != StyleStealSourceKind.IMAGE:
            raise ValueError(f"Unsupported style-steal source kind: {source.kind}")

        return self._analyze_contents(
            contents=[
                _STYLE_STEAL_PROMPT,
                types.Part.from_bytes(data=source.bytes_data, mime_type=source.mime_type),
            ],
        )

    def analyze_instruction(
        self,
        *,
        instruction: str,
        reference_sources: list[StyleStealSource] | None = None,
    ) -> StyleStealSuggestion:
        cleaned_instruction = instruction.strip()
        if not cleaned_instruction:
            raise ValueError("Instruction cannot be empty.")

        contents: list[Any] = [
            _STYLE_GENERATE_PROMPT,
            f"Style instruction:\n{cleaned_instruction}",
        ]
        for source in reference_sources or []:
            if source.kind != StyleStealSourceKind.IMAGE:
                raise ValueError(f"Unsupported style-steal source kind: {source.kind}")
            contents.append(
                types.Part.from_bytes(data=source.bytes_data, mime_type=source.mime_type)
            )

        return self._analyze_contents(contents=contents)

    def _analyze_contents(self, *, contents: list[Any]) -> StyleStealSuggestion:
        models_to_try = [_STYLE_ANALYSIS_PRIMARY_MODEL, _STYLE_ANALYSIS_FALLBACK_MODEL]
        response = None
        last_error: BaseException | None = None
        for index, model in enumerate(models_to_try):
            try:
                response = self._client.models.generate_content(
                    model=model,
                    contents=contents,
                    config=types.GenerateContentConfig(
                        temperature=1.0,
                        response_mime_type="application/json",
                    ),
                )
                self.last_model_used = model
                break
            except httpx.TimeoutException as exc:
                raise RuntimeError(
                    "Style analysis timed out. Retry with a higher --timeout-seconds value."
                ) from exc
            except Exception as exc:
                last_error = exc
                should_retry = (
                    index < len(models_to_try) - 1 and is_service_unavailable_error(exc)
                )
                if not should_retry:
                    raise
        if response is None:
            if last_error is not None:
                raise last_error
            raise RuntimeError("Style analysis failed before receiving a response.")
        payload = _parse_json_response(response)

        base_prompt = str(payload.get("base_prompt", "")).strip()
        negative_prompt = str(payload.get("negative_prompt", "")).strip()
        use_as_base_reference = bool(payload.get("use_as_base_reference", False))
        base_reference_reason = str(payload.get("base_reference_reason", "")).strip()
        reference_comments = _normalize_comments(payload.get("reference_comments", []))

        if not base_prompt:
            raise RuntimeError("Style analysis returned an empty base_prompt.")
        if not base_reference_reason:
            raise RuntimeError(
                "Style analysis returned no explanation for base-reference decision."
            )

        return StyleStealSuggestion(
            base_prompt=base_prompt,
            negative_prompt=negative_prompt,
            reference_comments=reference_comments,
            use_as_base_reference=use_as_base_reference,
            base_reference_reason=base_reference_reason,
        )


def infer_project_style_from_source(
    *,
    analyzer: GeminiStyleStealAnalyzer,
    source: StyleStealSource,
    set_base_reference: bool = False,
) -> InferredProjectStyle:
    """Infer an in-memory project style from a single source asset."""
    suggestion = analyzer.analyze(source)
    should_set_reference = set_base_reference or suggestion.use_as_base_reference
    project_style = ProjectStyleConfig(
        style_id=None,
        base_prompt=suggestion.base_prompt,
        negative_prompt=suggestion.negative_prompt,
        reference_images=[str(source.path)] if should_set_reference else [],
        reference_comments=suggestion.reference_comments,
    )
    return InferredProjectStyle(
        project_style=project_style,
        suggestion=suggestion,
        reference_sources=[source],
    )


def infer_project_style_from_instruction(
    *,
    analyzer: GeminiStyleStealAnalyzer,
    instruction: str,
    reference_sources: list[StyleStealSource] | None = None,
    set_base_reference: bool = False,
) -> InferredProjectStyle:
    """Infer an in-memory project style from instruction + optional references."""
    resolved_sources = list(reference_sources or [])
    suggestion = analyzer.analyze_instruction(
        instruction=instruction,
        reference_sources=resolved_sources,
    )
    should_set_reference = bool(resolved_sources) and (
        set_base_reference or suggestion.use_as_base_reference
    )
    project_style = ProjectStyleConfig(
        style_id=None,
        base_prompt=suggestion.base_prompt,
        negative_prompt=suggestion.negative_prompt,
        reference_images=(
            [str(source.path) for source in resolved_sources]
            if should_set_reference
            else []
        ),
        reference_comments=suggestion.reference_comments,
    )
    return InferredProjectStyle(
        project_style=project_style,
        suggestion=suggestion,
        reference_sources=resolved_sources,
    )


def load_style_steal_source(path: Path) -> StyleStealSource:
    """Load and validate a style-steal source file."""
    source_path = path.expanduser().resolve()
    if not source_path.exists() or not source_path.is_file():
        raise ValueError(f"Style source file not found: {source_path}")

    mime_type = _guess_mime_type(source_path)
    kind = _detect_source_kind(mime_type)
    if kind is None:
        raise ValueError(
            "Unsupported style source type. Use an image file for now "
            "(future PDF/PPTX support planned)."
        )

    return StyleStealSource(
        kind=kind,
        path=source_path,
        mime_type=mime_type,
        bytes_data=source_path.read_bytes(),
    )


def _guess_mime_type(path: Path) -> str:
    mime_type, _ = mimetypes.guess_type(str(path))
    if not mime_type:
        raise ValueError(f"Unable to infer MIME type for style source: {path}")
    return mime_type


def _detect_source_kind(mime_type: str) -> StyleStealSourceKind | None:
    if mime_type.startswith("image/"):
        return StyleStealSourceKind.IMAGE
    return None


def _normalize_comments(raw_comments: Any) -> list[str]:
    if not isinstance(raw_comments, list):
        return []
    normalized: list[str] = []
    for raw_comment in raw_comments:
        comment = str(raw_comment).strip()
        if comment:
            normalized.append(comment)
    return normalized


def _parse_json_response(response: Any) -> dict[str, Any]:
    raw_text = getattr(response, "text", None)
    if not raw_text:
        raw_text = "\n".join(_response_text_parts(response))
    if not raw_text:
        raise RuntimeError("Style analysis returned no text output.")

    cleaned = raw_text.strip()
    if cleaned.startswith("```"):
        cleaned = cleaned.strip("`")
        if cleaned.startswith("json"):
            cleaned = cleaned[len("json") :].strip()
    try:
        payload = json.loads(cleaned)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"Style analysis returned invalid JSON: {exc}") from exc
    if not isinstance(payload, dict):
        raise RuntimeError("Style analysis JSON payload must be an object.")
    return payload


def _response_text_parts(response: Any) -> list[str]:
    text_parts: list[str] = []
    candidates = getattr(response, "candidates", None) or []
    for candidate in candidates:
        content = getattr(candidate, "content", None)
        parts = getattr(content, "parts", None) or []
        for part in parts:
            text = getattr(part, "text", None)
            if text:
                text_parts.append(str(text))
    return text_parts


_STYLE_STEAL_PROMPT = """
You are extracting a reusable visual style for AI-generated slides.

Analyze the source and return ONLY strict JSON with this schema:
{
  "base_prompt": string,
  "negative_prompt": string,
  "reference_comments": string[],
  "use_as_base_reference": boolean,
  "base_reference_reason": string
}

Rules:
- base_prompt must describe reusable visual style (color palette, typography feel, composition tendencies, texture, tone) and avoid content-specific details.
- negative_prompt should list what to avoid so future slides stay clean and consistent.
- reference_comments should be short actionable notes about the style.
- Decide use_as_base_reference:
  - true when the source is a broadly reusable style anchor that improves consistency.
  - false when the source has a very specific layout/composition that would overconstrain generations.
- base_reference_reason must explain the decision in one concise sentence.
- Never include markdown or extra keys.
""".strip()

_STYLE_GENERATE_PROMPT = """
You are generating a reusable visual style for AI-generated slides from user direction.

Analyze the style instruction and any attached reference images, then return ONLY strict JSON
with this schema:
{
  "base_prompt": string,
  "negative_prompt": string,
  "reference_comments": string[],
  "use_as_base_reference": boolean,
  "base_reference_reason": string
}

Rules:
- base_prompt must capture reusable style traits, not slide-specific content.
- negative_prompt should list concrete things to avoid.
- reference_comments should be short actionable style notes.
- use_as_base_reference should be true only if attached references are broad reusable style anchors.
- base_reference_reason must explain the decision in one concise sentence.
- Never include markdown or extra keys.
""".strip()
