"""Google (Gemini) adapter: conversion helpers and async API wrapper."""

import base64
import json
import threading
from enum import Enum
from typing import Any

from google import genai
from pydantic import BaseModel

from dotpromptz.image import resolve_image_ref
from dotpromptz.logging import get_logger
from dotpromptz.typing import (
    DataPart,
    MediaPart,
    Message,
    Part,
    RenderedPrompt,
    Role,
    TextPart,
)

from ._base import (
    Adapter,
    Credential,
    GeneratedImage,
    GenerateResponse,
    Usage,
    resolve_config_params,
    resolve_rendered_config_and_model,
)

logger = get_logger(__name__)

# ---------------------------------------------------------------------------
# Conversion helpers
# ---------------------------------------------------------------------------

_ROLE_MAP: dict[Role, str] = {
    Role.USER: 'user',
    Role.MODEL: 'model',
    Role.SYSTEM: 'user',  # system handled separately
}
_GOOGLE_IMAGE_SIZES = frozenset({'1k', '2k', '4k'})
_GOOGLE_IMAGE_ASPECT_RATIOS = frozenset({'1:1', '2:3', '3:2', '3:4', '4:3', '9:16', '16:9', '21:9'})


def _to_google_thinking_config(level: str) -> dict[str, Any]:
    """Map a canonical thinking level to Google ``thinking_config``.

    Args:
        level: Canonical thinking level (low/medium/high).

    Returns:
        A ``thinking_config`` dict.
    """
    budget_map = {
        'low': 1024,
        'medium': 8192,
        'high': -1,  # automatic / highest
    }
    return {'thinking_budget': budget_map[level]}


def _part_to_genai(part: Part) -> dict[str, Any] | None:
    """Convert a dotprompt ``Part`` to a Google GenAI part dict.

    Args:
        part: A dotprompt content part.

    Returns:
        A dict for the Gemini ``parts`` list, or ``None``.
    """
    if isinstance(part, TextPart):
        return {'text': part.text}
    if isinstance(part, MediaPart):
        resolved = resolve_image_ref(
            part.media.url,
            content_type=part.media.content_type,
            read_bytes=True,
        )
        if resolved.kind == 'local_file':
            if resolved.data is None:
                raise ValueError('Local image bytes are missing during Google request conversion.')
            encoded = base64.b64encode(resolved.data).decode('ascii')
            return {
                'inline_data': {
                    'mime_type': resolved.mime_type,
                    'data': encoded,
                },
            }
        return {
            'file_data': {
                'mime_type': resolved.mime_type,
                'file_uri': resolved.url or part.media.url,
            },
        }
    if isinstance(part, DataPart):
        return {'text': json.dumps(part.data, ensure_ascii=False)}
    return None


def to_genai_contents(messages: list[Message]) -> tuple[str | None, list[dict[str, Any]]]:
    """Convert dotprompt messages to Gemini ``contents`` format.

    System messages are extracted separately for the ``system_instruction`` parameter.

    Args:
        messages: List of dotprompt ``Message`` objects.

    Returns:
        A tuple of ``(system_instruction_text, contents_list)``.
    """
    system_parts: list[str] = []
    contents: list[dict[str, Any]] = []

    for msg in messages:
        if msg.role == Role.SYSTEM:
            for part in msg.content:
                if isinstance(part, TextPart):
                    system_parts.append(part.text)
            continue

        role = _ROLE_MAP.get(msg.role, 'user')
        parts: list[dict[str, Any]] = []
        for part in msg.content:
            converted = _part_to_genai(part)
            if converted is not None:
                parts.append(converted)

        if parts:
            contents.append({'role': role, 'parts': parts})

    system_text = '\n'.join(system_parts) if system_parts else None
    return system_text, contents


def to_genai_request(rendered: RenderedPrompt) -> dict[str, Any]:
    """Convert a ``RenderedPrompt`` to a Google GenAI request dict.

    Args:
        rendered: The rendered prompt.

    Returns:
        A dict suitable for ``client.models.generate_content(**d)``.

    Raises:
        ValueError: If no model is specified.
    """
    if rendered.output is not None and rendered.output.format == 'image':
        return to_genai_image_request(rendered)

    system_text, contents = to_genai_contents(rendered.messages)
    config, resolved_model = resolve_rendered_config_and_model(rendered)

    request: dict[str, Any] = {
        'model': resolved_model,
        'contents': contents,
    }

    if system_text:
        request['config'] = request.get('config', {})
        request['config']['system_instruction'] = system_text

    # Resolve generation parameters via centralised helper.
    params = resolve_config_params(config)
    gen_config: dict[str, Any] = {}
    if params.get('temperature') is not None:
        gen_config['temperature'] = params['temperature']
    if params.get('top_p') is not None:
        gen_config['top_p'] = params['top_p']
    if params.get('max_tokens') is not None:
        gen_config['max_output_tokens'] = params['max_tokens']
    gen_config['thinking_config'] = _to_google_thinking_config(params.get('thinking', 'high'))
    stop = params.get('stop')
    if stop:
        gen_config['stop_sequences'] = stop if isinstance(stop, list) else [stop]

    if rendered.output and rendered.output.format in ('json', 'yaml'):
        gen_config['response_mime_type'] = 'application/json'
        if rendered.output.schema_value is not None:
            gen_config['response_json_schema'] = rendered.output.schema_value

    if gen_config:
        request['config'] = {**request.get('config', {}), **gen_config}

    return request


def _normalize_google_image_size(config: dict[str, Any]) -> str | None:
    """Normalize ``config.resolution`` for Google image generation."""
    raw_resolution = config.get('resolution')
    if raw_resolution is None:
        return None
    if not isinstance(raw_resolution, str):
        raise ValueError('config.resolution must be a string for Google image generation.')
    normalized = raw_resolution.strip().lower()
    if normalized not in _GOOGLE_IMAGE_SIZES:
        raise ValueError("config.resolution for Google image generation must be one of ['1K', '2K', '4K'].")
    return normalized.upper()


def _normalize_google_aspect_ratio(config: dict[str, Any]) -> str | None:
    """Normalize ``config.aspect_ratio`` for Google image generation."""
    raw_aspect_ratio = config.get('aspect_ratio')
    if raw_aspect_ratio is None:
        return None
    if not isinstance(raw_aspect_ratio, str):
        raise ValueError('config.aspect_ratio must be a string for Google image generation.')
    normalized = raw_aspect_ratio.replace(' ', '')
    if normalized not in _GOOGLE_IMAGE_ASPECT_RATIOS:
        raise ValueError(
            'config.aspect_ratio for Google image generation must be one of '
            "['1:1', '2:3', '3:2', '3:4', '4:3', '9:16', '16:9', '21:9']."
        )
    return normalized


def to_genai_image_request(rendered: RenderedPrompt) -> dict[str, Any]:
    """Convert a ``RenderedPrompt`` to Gemini image-generation request dict."""
    system_text, contents = to_genai_contents(rendered.messages)
    config, resolved_model = resolve_rendered_config_and_model(rendered)
    params = resolve_config_params(config)

    request: dict[str, Any] = {
        'model': resolved_model,
        'contents': contents,
        'config': {
            'response_modalities': ['IMAGE'],
        },
    }
    if system_text:
        request['config']['system_instruction'] = system_text

    image_config: dict[str, Any] = {}
    image_size = _normalize_google_image_size(config)
    if image_size is not None:
        image_config['image_size'] = image_size
    aspect_ratio = _normalize_google_aspect_ratio(config)
    if aspect_ratio is not None:
        image_config['aspect_ratio'] = aspect_ratio
    if image_config:
        request['config']['image_config'] = image_config

    if config.get('transparent') is not None:
        if not isinstance(config.get('transparent'), bool):
            raise ValueError('config.transparent must be a boolean for Google image generation.')
        logger.warning('config.transparent is not supported by Google Gemini image API; ignoring value.')

    if params.get('temperature') is not None:
        request['config']['temperature'] = params['temperature']
    if params.get('top_p') is not None:
        request['config']['top_p'] = params['top_p']

    return request


def _extract_google_response_text(response: Any) -> str | None:
    """Extract text from a Gemini response."""
    direct_text = getattr(response, 'text', None)
    if isinstance(direct_text, str):
        return direct_text

    text: str | None = None
    if response.candidates and response.candidates[0].content:
        for part in response.candidates[0].content.parts:
            if isinstance(getattr(part, 'thought', None), bool) and part.thought:
                continue
            if hasattr(part, 'text') and part.text:
                text = (text or '') + part.text
    return text


def _serialize_google_parsed(parsed: Any) -> str | None:
    """Serialize Gemini parsed payload to JSON text."""
    if parsed is None:
        return None

    value: Any = parsed
    if isinstance(parsed, BaseModel):
        value = parsed.model_dump(mode='json')
    elif isinstance(parsed, Enum):
        value = parsed.value

    try:
        return json.dumps(value, ensure_ascii=False)
    except TypeError:
        return None


def _extract_google_generated_images(response: Any) -> list[GeneratedImage]:
    """Extract inline generated image bytes from Gemini response payload."""
    images: list[GeneratedImage] = []
    candidates = getattr(response, 'candidates', None)
    if not isinstance(candidates, list):
        raise ValueError('Google image response contained no candidates.')

    for candidate in candidates:
        content = getattr(candidate, 'content', None)
        parts = getattr(content, 'parts', None)
        if not isinstance(parts, list):
            continue
        for part in parts:
            inline_data = getattr(part, 'inline_data', None)
            if inline_data is None:
                continue
            mime_type = getattr(inline_data, 'mime_type', None)
            data = getattr(inline_data, 'data', None)
            if not isinstance(mime_type, str) or not mime_type.startswith('image/'):
                continue
            if not isinstance(data, bytes):
                continue
            if mime_type.lower() != 'image/png':
                raise ValueError(f'Google image output requires MIME image/png; got {mime_type!r}.')
            images.append(GeneratedImage(data=data, mime_type=mime_type))

    if not images:
        raise ValueError('Google image response contained no inline image data.')
    return images


# ---------------------------------------------------------------------------
# Google Adapter
# ---------------------------------------------------------------------------


class GoogleAdapter(Adapter):
    """Adapter for the Google (Gemini) API.

    Uses the ``google-genai`` SDK (``google.genai``).

    Args:
        credential: A ``Credential`` with ``api_key`` and optional ``base_url``.
        default_model: Fallback model name.
    """

    def __init__(
        self,
        *,
        credential: Credential,
        default_model: str = 'gemini-3.1-flash-image-preview',
    ) -> None:
        """Initialise the Google adapter."""
        self._credential = credential
        self._default_model = default_model
        self._client: Any = None
        self._client_lock = threading.Lock()

    def _get_client(self) -> Any:
        """Return (and lazily create) the ``google.genai.Client``."""

        def _factory() -> genai.Client:
            kwargs: dict[str, Any] = {'api_key': self._credential.api_key}
            if self._credential.base_url:
                kwargs['http_options'] = {'base_url': self._credential.base_url}
            return genai.Client(**kwargs)

        return self._get_or_create_client(_factory)

    async def aclose(self) -> None:
        """Close the underlying ``google.genai.Client``.

        Closes both the sync and async HTTP transports.
        """
        if self._client is not None:
            await self._client.aio.aclose()
            self._client.close()
            self._client = None

    def convert(self, rendered: RenderedPrompt) -> dict[str, Any]:
        """Convert a ``RenderedPrompt`` to a Google GenAI request dict.

        Args:
            rendered: The rendered prompt.

        Returns:
            A dict suitable for ``client.models.generate_content(**d)``.
        """
        return self._convert_with_model(rendered, to_genai_request)

    async def generate(
        self,
        rendered: RenderedPrompt,
        **kwargs: Any,
    ) -> GenerateResponse:
        """Call the Google GenAI API.

        Args:
            rendered: The rendered prompt.
            **kwargs: Extra keyword args forwarded to the SDK.

        Returns:
            A ``GenerateResponse``.
        """
        client = self._get_client()
        request_dict = self.convert(rendered)
        request_dict.update(kwargs)

        response = await client.aio.models.generate_content(**request_dict)

        if rendered.output and rendered.output.format == 'image':
            usage: Usage | None = None
            if response.usage_metadata:
                um = response.usage_metadata
                prompt_tokens = getattr(um, 'prompt_token_count', 0) or 0
                completion_tokens = getattr(um, 'candidates_token_count', 0) or 0
                usage = Usage(
                    prompt_tokens=prompt_tokens,
                    completion_tokens=completion_tokens,
                    total_tokens=prompt_tokens + completion_tokens,
                )

            finish_reason = None
            if response.candidates:
                fr = getattr(response.candidates[0], 'finish_reason', None)
                if fr is not None:
                    finish_reason = str(fr)

            return GenerateResponse(
                text=None,
                usage=usage,
                raw=response,
                model=request_dict.get('model'),
                finish_reason=finish_reason,
                images=_extract_google_generated_images(response),
            )

        is_structured_output = rendered.output and rendered.output.format in ('json', 'yaml')
        text: str | None = None
        if is_structured_output:
            text = _serialize_google_parsed(getattr(response, 'parsed', None))
        if text is None:
            text = _extract_google_response_text(response)

        usage: Usage | None = None
        if response.usage_metadata:
            um = response.usage_metadata
            prompt_tokens = getattr(um, 'prompt_token_count', 0) or 0
            completion_tokens = getattr(um, 'candidates_token_count', 0) or 0
            usage = Usage(
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=prompt_tokens + completion_tokens,
            )

        finish_reason = None
        if response.candidates:
            fr = getattr(response.candidates[0], 'finish_reason', None)
            if fr is not None:
                finish_reason = str(fr)

        return GenerateResponse(
            text=text,
            usage=usage,
            raw=response,
            model=request_dict.get('model'),
            finish_reason=finish_reason,
        )
