# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

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

from __future__ import annotations

import json
import os
from enum import Enum
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, GenerateResponse, ToolCallResult, Usage


class DetailKind(str, Enum):
    """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(str, Enum):
    """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(str, Enum):
    """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(str, Enum):
    """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)

    return OpenAIRequest(
        messages=to_openai_messages(rendered.messages),
        model=resolved_model,
        tools=to_openai_tools(rendered.tool_defs),
        temperature=config.get('temperature'),
        top_p=config.get('topP') or config.get('top_p'),
        max_tokens=config.get('maxTokens') or config.get('max_tokens'),
        stop=config.get('stop') or config.get('stopSequences'),
        frequency_penalty=config.get('frequencyPenalty') or config.get('frequency_penalty'),
        presence_penalty=config.get('presencePenalty') or config.get('presence_penalty'),
        seed=config.get('seed'),
        response_format=response_format,
    )


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


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

    Args:
        api_key: OpenAI API key.  Falls back to ``OPENAI_API_KEY`` env var.
        base_url: Optional custom base URL (for Azure, proxies, etc.).
        default_model: Fallback model when the prompt has none.
        organization: Optional OpenAI organization ID.
    """

    def __init__(
        self,
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        default_model: str = 'gpt-4o',
        organization: str | None = None,
    ) -> None:
        """Initialise the OpenAI adapter."""
        self._api_key = api_key or os.environ.get('OPENAI_API_KEY', '')
        self._base_url = base_url or os.environ.get('OPENAI_BASE_URL')
        self._default_model = default_model
        self._organization = organization
        self._client: Any = None  # lazily created

    def _get_client(self) -> Any:
        """Return (and lazily create) the ``AsyncOpenAI`` client."""
        if self._client is None:
            kwargs: dict[str, Any] = {'api_key': self._api_key}
            if self._base_url:
                kwargs['base_url'] = self._base_url
            if self._organization:
                kwargs['organization'] = self._organization
            self._client = AsyncOpenAI(**kwargs)
        return self._client

    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,
        )
