"""OpenAI adapter: data models, conversion helpers, and async API wrapper."""

import json
import threading
from enum import StrEnum
from typing import Any, Literal

from openai import AsyncOpenAI
from pydantic import BaseModel, ConfigDict, Field

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

from ._base import Adapter, Credential, GenerateResponse, ToolCallResult, Usage, resolve_config_params


class DetailKind(StrEnum):
    """The kind of Image URL detail."""

    AUTO = 'auto'
    LOW = 'low'
    HIGH = 'high'


# Define TypedDicts to represent the interfaces
class ImageURLDetail(BaseModel):
    """Image URL Detail."""

    url: str
    detail: DetailKind | None = None


class ContentItemType(StrEnum):
    """Enumveration variants for content item type."""

    TEXT = 'text'
    IMAGE_URL = 'image_url'


class ContentItem(BaseModel):
    """Content Item: can be text or image."""

    type: ContentItemType
    text: str | None = None
    image_url: ImageURLDetail | None = None


class ToolFunction(BaseModel):
    """Tool function."""

    name: str
    arguments: str


class ToolCallType(StrEnum):
    """Enumeration variants for tool call type."""

    FUNCTION = 'function'


class ToolCall(BaseModel):
    """Tool Call."""

    id: str
    type: ToolCallType
    function: ToolFunction


class OpenAIMessage(BaseModel):
    """Open AI Message."""

    role: str
    content: str | list[ContentItem] | None = None
    name: str | None = None
    tool_calls: list[ToolCall] | None = None
    tool_call_id: str | None = None


class OpenAIToolFunction(BaseModel):
    """OpenAI Tool Function."""

    name: str
    description: str | None = None
    parameters: dict[str, Any] | None = None


class OpenAIToolDefinition(BaseModel):
    """OpenAI Tool Definition."""

    type: ToolCallType
    function: OpenAIToolFunction


class ToolChoiceFunction(BaseModel):
    """Tool Choice Function."""

    name: str


class ToolChoice(BaseModel):
    """Tool Choice."""

    type: ToolCallType
    function: ToolChoiceFunction


class ResponseFormatType(StrEnum):
    """Enum variants for response format type."""

    TEXT = 'text'
    JSON_OBJECT = 'json_object'
    JSON_SCHEMA = 'json_schema'


class JsonSchemaSpec(BaseModel):
    """JSON Schema specification within a structured output response format."""

    name: str
    strict: bool = True
    schema_: dict[str, Any] = Field(..., alias='schema')
    model_config = ConfigDict(populate_by_name=True, serialize_by_alias=True)


class ResponseFormat(BaseModel):
    """Expected Response Format."""

    type: ResponseFormatType
    json_schema: JsonSchemaSpec | None = None


ToolChoiceOptions = Literal['none', 'auto']


class OpenAIRequest(BaseModel):
    """Open AI Request."""

    messages: list[OpenAIMessage]
    model: str
    frequency_penalty: float | None = None
    logit_bias: dict[str, int] | None = None
    max_tokens: int | None = None
    n: int | None = None
    presence_penalty: float | None = None
    response_format: ResponseFormat | None = None
    seed: int | None = None
    stop: str | list[str] | None = None
    temperature: float | None = None
    tool_choice: ToolChoiceOptions | ToolChoice | None = None
    tools: list[OpenAIToolDefinition] | None = None
    top_p: float | None = None
    user: str | None = None


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

_ROLE_MAP: dict[Role, str] = {
    Role.USER: 'user',
    Role.MODEL: 'assistant',
    Role.SYSTEM: 'system',
    Role.TOOL: 'tool',
}


def _role_to_openai(role: Role) -> str:
    """Map a dotprompt ``Role`` to the OpenAI role string.

    Args:
        role: The dotprompt role enum value.

    Returns:
        The corresponding OpenAI role string.
    """
    return _ROLE_MAP.get(role, role.value)


def _part_to_content_item(part: Part) -> ContentItem | str | None:
    """Convert a single dotprompt ``Part`` to an OpenAI content item.

    Args:
        part: A dotprompt content part.

    Returns:
        A ``ContentItem``, a plain string, or ``None`` if unsupported.
    """
    if isinstance(part, TextPart):
        return ContentItem(type=ContentItemType.TEXT, text=part.text)
    if isinstance(part, MediaPart):
        return ContentItem(
            type=ContentItemType.IMAGE_URL,
            image_url=ImageURLDetail(url=part.media.url),
        )
    if isinstance(part, DataPart):
        return ContentItem(type=ContentItemType.TEXT, text=json.dumps(part.data, ensure_ascii=False))
    return None


def _parts_to_content(parts: list[Part]) -> str | list[ContentItem]:
    """Convert a list of dotprompt ``Part`` objects to OpenAI content.

    If all parts are text-only, returns a single concatenated string.
    Otherwise returns a list of ``ContentItem``.

    Args:
        parts: List of dotprompt content parts.

    Returns:
        Either a plain string or a list of ``ContentItem``.
    """
    if all(isinstance(p, TextPart) for p in parts):
        return ''.join(p.text for p in parts if isinstance(p, TextPart))

    items: list[ContentItem] = []
    for part in parts:
        item = _part_to_content_item(part)
        if isinstance(item, ContentItem):
            items.append(item)
    return items


def _extract_tool_calls(parts: list[Part]) -> list[ToolCall] | None:
    """Extract ``ToolCall`` objects from tool-request parts.

    Args:
        parts: List of dotprompt content parts.

    Returns:
        A list of ``ToolCall`` or ``None`` if no tool requests exist.
    """
    calls: list[ToolCall] = []
    for part in parts:
        if isinstance(part, ToolRequestPart):
            req = part.tool_request
            calls.append(
                ToolCall(
                    id=req.ref or req.name,
                    type=ToolCallType.FUNCTION,
                    function=ToolFunction(
                        name=req.name,
                        arguments=json.dumps(req.input or {}, ensure_ascii=False),
                    ),
                )
            )
    return calls or None


def _message_to_openai(msg: Message) -> OpenAIMessage:
    """Convert a single dotprompt ``Message`` to an ``OpenAIMessage``.

    Args:
        msg: A dotprompt message.

    Returns:
        An ``OpenAIMessage``.
    """
    role = _role_to_openai(msg.role)
    tool_calls = _extract_tool_calls(msg.content)

    # Tool response messages require a tool_call_id.
    tool_call_id: str | None = None
    if msg.role == Role.TOOL:
        for part in msg.content:
            if isinstance(part, ToolResponsePart):
                tool_call_id = part.tool_response.ref or part.tool_response.name
                break

    content = _parts_to_content([
        p for p in msg.content if not isinstance(p, (ToolRequestPart, ToolResponsePart, PendingPart))
    ])
    # For tool response messages, content must be a string.
    if msg.role == Role.TOOL and isinstance(content, list):
        content = json.dumps([item.model_dump(exclude_none=True) for item in content], ensure_ascii=False)

    return OpenAIMessage(
        role=role,
        content=content if content else None,
        tool_calls=tool_calls,
        tool_call_id=tool_call_id,
    )


def to_openai_messages(messages: list[Message]) -> list[OpenAIMessage]:
    """Convert dotprompt messages to OpenAI message format.

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

    Returns:
        List of ``OpenAIMessage`` objects.
    """
    return [_message_to_openai(m) for m in messages]


def to_openai_tools(tool_defs: list[ToolDefinition] | None) -> list[OpenAIToolDefinition] | None:
    """Convert dotprompt ``ToolDefinition`` list to OpenAI tool format.

    Args:
        tool_defs: List of dotprompt tool definitions.

    Returns:
        List of ``OpenAIToolDefinition`` or ``None``.
    """
    if not tool_defs:
        return None
    return [
        OpenAIToolDefinition(
            type=ToolCallType.FUNCTION,
            function=OpenAIToolFunction(
                name=td.name,
                description=td.description,
                parameters=td.input_schema,
            ),
        )
        for td in tool_defs
    ]


def to_openai_request(rendered: RenderedPrompt[Any]) -> OpenAIRequest:
    """Convert a ``RenderedPrompt`` to a full ``OpenAIRequest``.

    Args:
        rendered: The rendered prompt from ``Dotprompt.render()``.

    Returns:
        An ``OpenAIRequest`` ready for the OpenAI API.

    Raises:
        ValueError: If no model is specified.
    """
    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.')

    response_format: ResponseFormat | None = None
    if rendered.output and rendered.output.format == 'json':
        if rendered.output.schema_ is not None:
            response_format = ResponseFormat(
                type=ResponseFormatType.JSON_SCHEMA,
                json_schema=JsonSchemaSpec(
                    name='response',
                    strict=True,
                    schema=rendered.output.schema_,
                ),
            )
        else:
            response_format = ResponseFormat(type=ResponseFormatType.JSON_OBJECT)

    params = resolve_config_params(config)

    return OpenAIRequest(
        messages=to_openai_messages(rendered.messages),
        model=resolved_model,
        tools=to_openai_tools(rendered.tool_defs),
        temperature=params.get('temperature'),
        top_p=params.get('top_p'),
        max_tokens=params.get('max_tokens'),
        stop=params.get('stop'),
        frequency_penalty=params.get('frequency_penalty'),
        presence_penalty=params.get('presence_penalty'),
        seed=params.get('seed'),
        response_format=response_format,
    )


# ---------------------------------------------------------------------------
# OpenAI Adapter
# ---------------------------------------------------------------------------


class OpenAIAdapter(Adapter):
    """Adapter for the OpenAI Chat Completions API.

    Args:
        credential: A ``Credential`` with ``api_key`` and optional ``base_url``.
        default_model: Fallback model when the prompt has none.
        organization: Optional OpenAI organization ID.
    """

    def __init__(
        self,
        *,
        credential: Credential,
        default_model: str = 'gpt-4o',
        organization: str | None = None,
    ) -> None:
        """Initialise the OpenAI adapter."""
        self._credential = credential
        self._default_model = default_model
        self._organization = organization
        self._client: Any = None
        self._client_lock = threading.Lock()

    def _get_client(self) -> Any:
        """Return (and lazily create) the ``AsyncOpenAI`` 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['base_url'] = self._credential.base_url
                    if self._organization:
                        kwargs['organization'] = self._organization
                    self._client = AsyncOpenAI(**kwargs)
        return self._client

    async def aclose(self) -> None:
        """Close the underlying ``AsyncOpenAI`` client."""
        if self._client is not None:
            await self._client.close()
            self._client = None

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

        Args:
            rendered: The rendered prompt.

        Returns:
            A dict suitable for ``client.chat.completions.create(**d)``.
        """
        rendered = self._ensure_model(rendered)
        req = to_openai_request(rendered)
        return req.model_dump(exclude_none=True)

    async def generate(
        self,
        rendered: RenderedPrompt[Any],
        **kwargs: Any,
    ) -> GenerateResponse:
        """Call the OpenAI Chat Completions API.

        Args:
            rendered: The rendered prompt from ``Dotprompt.render()``.
            **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.chat.completions.create(**request_dict)
        choice = response.choices[0]
        msg = choice.message

        tool_calls: list[ToolCallResult] = []
        if msg.tool_calls:
            for tc in msg.tool_calls:
                tool_calls.append(
                    ToolCallResult(
                        id=tc.id,
                        name=tc.function.name,
                        arguments=json.loads(tc.function.arguments) if tc.function.arguments else {},
                    )
                )

        usage: Usage | None = None
        if response.usage:
            usage = Usage(
                prompt_tokens=response.usage.prompt_tokens,
                completion_tokens=response.usage.completion_tokens,
                total_tokens=response.usage.total_tokens,
            )

        return GenerateResponse(
            text=msg.content,
            tool_calls=tool_calls,
            usage=usage,
            raw=response,
            model=response.model,
            finish_reason=choice.finish_reason,
        )
