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

import base64
import json
import threading
from typing import Any

from google import genai

from dotpromptz.typing import (
    DataPart,
    MediaPart,
    Message,
    Part,
    RenderedPrompt,
    Role,
    TextPart,
    ToolDefinition,
    ToolRequestPart,
    ToolResponsePart,
)

from ._base import (
    Adapter,
    Credential,
    GenerateResponse,
    ImageResult,
    ToolCallResult,
    Usage,
    parse_data_uri,
    resolve_config_params,
)

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

_ROLE_MAP: dict[Role, str] = {
    Role.USER: 'user',
    Role.MODEL: 'model',
    Role.TOOL: 'user',  # tool results sent as user
    Role.SYSTEM: 'user',  # system handled separately
}


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):
        url = part.media.url
        mime_type = part.media.content_type or 'image/png'
        parsed = parse_data_uri(url)
        if parsed:
            inferred, data = parsed
            return {
                'inline_data': {
                    'mime_type': inferred or mime_type,
                    'data': data,
                },
            }
        return {
            'file_data': {
                'mime_type': mime_type,
                'file_uri': url,
            },
        }
    if isinstance(part, DataPart):
        return {'text': json.dumps(part.data, ensure_ascii=False)}
    if isinstance(part, ToolRequestPart):
        req = part.tool_request
        return {
            'function_call': {
                'name': req.name,
                'args': req.input or {},
            },
        }
    if isinstance(part, ToolResponsePart):
        resp = part.tool_response
        return {
            'function_response': {
                'name': resp.name,
                'response': resp.output if resp.output is not None else {},
            },
        }
    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_tools(tool_defs: list[ToolDefinition] | None) -> list[dict[str, Any]] | None:
    """Convert dotprompt tool definitions to Gemini tool format.

    Args:
        tool_defs: List of dotprompt tool definitions.

    Returns:
        A list containing a single tool dict with ``function_declarations``, or ``None``.
    """
    if not tool_defs:
        return None
    declarations = []
    for td in tool_defs:
        decl: dict[str, Any] = {
            'name': td.name,
            'description': td.description or '',
        }
        if td.input_schema:
            decl['parameters'] = td.input_schema
        declarations.append(decl)
    return [{'function_declarations': declarations}]


def to_genai_request(rendered: RenderedPrompt[Any]) -> 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.
    """
    system_text, contents = to_genai_contents(rendered.messages)
    config: dict[str, Any] = rendered.config if isinstance(rendered.config, dict) else {}

    resolved_model = config.get('model')
    if not resolved_model:
        raise ValueError('No model specified in rendered prompt.')

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

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

    tools = to_genai_tools(rendered.tool_defs)
    if tools:
        request['config'] = request.get('config', {})
        request['config']['tools'] = tools

    # 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']
    stop = params.get('stop')
    if stop:
        gen_config['stop_sequences'] = stop if isinstance(stop, list) else [stop]

    if rendered.output and rendered.output.format == 'json':
        gen_config['response_mime_type'] = 'application/json'
        if rendered.output.schema_ is not None:
            gen_config['response_json_schema'] = rendered.output.schema_
    elif rendered.output and rendered.output.format == 'image':
        gen_config['response_modalities'] = ['IMAGE']

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

    return request


# ---------------------------------------------------------------------------
# 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-2.5-flash',
    ) -> 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``."""
        if self._client is None:
            with self._client_lock:
                if self._client is None:
                    kwargs: dict[str, Any] = {'api_key': self._credential.api_key}
                    if self._credential.base_url:
                        kwargs['http_options'] = {'base_url': self._credential.base_url}
                    self._client = genai.Client(**kwargs)
        return self._client

    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[Any]) -> 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)``.
        """
        rendered = self._ensure_model(rendered)
        return to_genai_request(rendered)

    async def generate(
        self,
        rendered: RenderedPrompt[Any],
        **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)

        text: str | None = None
        tool_calls: list[ToolCallResult] = []
        image_result: ImageResult | None = None

        if response.candidates and response.candidates[0].content:
            for part in response.candidates[0].content.parts:
                if hasattr(part, 'text') and part.text:
                    text = (text or '') + part.text
                elif hasattr(part, 'inline_data') and part.inline_data:
                    raw_data = part.inline_data.data
                    if isinstance(raw_data, str):
                        raw_data = base64.b64decode(raw_data)
                    mime_type = part.inline_data.mime_type or 'image/png'

                    image_result = ImageResult(
                        data=raw_data,
                        mime_type=mime_type,
                    )
                    break  # Only take the first image
                elif hasattr(part, 'function_call') and part.function_call:
                    fc = part.function_call
                    tool_calls.append(
                        ToolCallResult(
                            id=fc.name,  # Gemini doesn't have separate IDs
                            name=fc.name,
                            arguments=dict(fc.args) if fc.args else {},
                        )
                    )

        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,
            image=image_result,
            tool_calls=tool_calls,
            usage=usage,
            raw=response,
            model=request_dict.get('model'),
            finish_reason=finish_reason,
        )
