"""OpenAI backend – supports chat, vision (VLM), function calling, streaming, and embeddings."""

from __future__ import annotations

import copy
import json
from collections.abc import AsyncIterator
from typing import Any, Self

from thryve.core.backends.base import BaseBackend
from thryve.core.backends.openai.param import OpenAIBackendParam
from thryve.core.context.models import (
    BackendInfo,
    ChatResponse,
    FinishReason,
    ImagePart,
    Message,
    TextPart,
    ToolCall,
    UsageInfo,
)
from thryve.utils import get_logger

logger = get_logger("backends.openai")


class OpenAIBackend(BaseBackend):
    """Concrete backend using the ``openai`` Python SDK."""

    def __init__(self, param: OpenAIBackendParam) -> None:
        try:
            import openai  # noqa: F811
        except ImportError as exc:
            raise ImportError("Install the 'openai' package: pip install openai") from exc

        self._param = param
        client_kwargs: dict[str, Any] = {}
        if param.api_key:
            client_kwargs["api_key"] = param.api_key
        if param.base_url:
            client_kwargs["base_url"] = param.base_url

        self._client = openai.OpenAI(**client_kwargs)
        self._async_client = openai.AsyncOpenAI(**client_kwargs)
        self._bound_tools: list[Any] | None = None  # Tool objects

    # ------------------------------------------------------------------
    # Internal: message format conversion
    # ------------------------------------------------------------------

    def _to_openai_messages(self, messages: list[Message]) -> list[dict[str, Any]]:
        """Convert internal ``Message`` list to the OpenAI API dict format."""
        result: list[dict[str, Any]] = []
        for msg in messages:
            d: dict[str, Any] = {"role": msg.role}

            # Content handling
            if isinstance(msg.content, str):
                d["content"] = msg.content
            else:
                parts: list[dict[str, Any]] = []
                for part in msg.content:
                    if isinstance(part, TextPart):
                        parts.append({"type": "text", "text": part.text})
                    elif isinstance(part, ImagePart):
                        if isinstance(part.data, str) and (
                            part.data.startswith("http://") or part.data.startswith("https://")
                        ):
                            image_url = {"url": part.data}
                        else:
                            # base64 data
                            data_str = part.data if isinstance(part.data, str) else __import__("base64").b64encode(part.data).decode()
                            image_url = {"url": f"data:{part.mime_type};base64,{data_str}"}
                        parts.append({"type": "image_url", "image_url": image_url})
                d["content"] = parts

            # Tool calls from assistant
            if msg.tool_calls:
                d["tool_calls"] = [
                    {
                        "id": tc.id,
                        "type": "function",
                        "function": {
                            "name": tc.name,
                            "arguments": json.dumps(tc.arguments),
                        },
                    }
                    for tc in msg.tool_calls
                ]

            # Tool result message
            if msg.role == "tool" and msg.tool_call_id:
                d["tool_call_id"] = msg.tool_call_id

            result.append(d)
        return result

    def _build_api_kwargs(self, **overrides: Any) -> dict[str, Any]:
        """Build keyword args for the OpenAI API call."""
        kwargs: dict[str, Any] = {
            "model": self._param.model,
            "temperature": self._param.temperature,
        }
        if self._param.max_tokens is not None:
            kwargs["max_tokens"] = self._param.max_tokens
        if self._bound_tools:
            kwargs["tools"] = [self._format_tool(t) for t in self._bound_tools]
        kwargs.update(self._param.extra)
        kwargs.update(overrides)
        return kwargs

    @staticmethod
    def _format_tool(tool: Any) -> dict[str, Any]:
        """Convert an internal ``Tool`` object to the OpenAI tools format."""
        return tool.to_openai_format()

    def _parse_response(self, response: Any) -> ChatResponse:
        """Parse an OpenAI chat completion into our unified ``ChatResponse``."""
        choice = response.choices[0]
        message = choice.message

        # Parse tool calls
        tool_calls: list[ToolCall] | None = None
        if message.tool_calls:
            tool_calls = [
                ToolCall(
                    id=tc.id,
                    name=tc.function.name,
                    arguments=json.loads(tc.function.arguments),
                )
                for tc in message.tool_calls
            ]

        # Finish reason
        fr_map = {
            "stop": FinishReason.STOP,
            "tool_calls": FinishReason.TOOL_CALLS,
            "length": FinishReason.LENGTH,
        }
        finish = fr_map.get(choice.finish_reason, FinishReason.STOP)

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

        return ChatResponse(
            content=message.content or "",
            tool_calls=tool_calls,
            backend=self.get_info(),
            finish_reason=finish,
            usage=usage,
        )

    # ------------------------------------------------------------------
    # BaseBackend implementation
    # ------------------------------------------------------------------

    def generate(self, messages: list[Message], **kwargs: Any) -> ChatResponse:
        api_messages = self._to_openai_messages(messages)
        api_kwargs = self._build_api_kwargs(**kwargs)
        response = self._client.chat.completions.create(
            messages=api_messages,  # type: ignore[arg-type]
            **api_kwargs,
        )
        return self._parse_response(response)

    async def generate_async(self, messages: list[Message], **kwargs: Any) -> ChatResponse:
        api_messages = self._to_openai_messages(messages)
        api_kwargs = self._build_api_kwargs(**kwargs)
        response = await self._async_client.chat.completions.create(
            messages=api_messages,  # type: ignore[arg-type]
            **api_kwargs,
        )
        return self._parse_response(response)

    async def generate_stream(self, messages: list[Message], **kwargs: Any) -> AsyncIterator[str]:
        # When tools are bound, fall back to non-streaming to capture tool calls.
        if self._bound_tools:
            resp = await self.generate_async(messages, **kwargs)
            if resp.content:
                yield resp.content
            return

        api_messages = self._to_openai_messages(messages)
        api_kwargs = self._build_api_kwargs(stream=True, **kwargs)
        stream = await self._async_client.chat.completions.create(
            messages=api_messages,  # type: ignore[arg-type]
            **api_kwargs,
        )
        async for chunk in stream:
            delta = chunk.choices[0].delta if chunk.choices else None
            if delta and delta.content:
                yield delta.content

    def embeddings(self, texts: list[str], **kwargs: Any) -> list[list[float]]:
        response = self._client.embeddings.create(
            model=kwargs.get("model", self._param.model),
            input=texts,
            **{k: v for k, v in kwargs.items() if k != "model"},
        )
        return [item.embedding for item in response.data]

    def bind_tools(self, tools: list[Any]) -> Self:
        new = copy.copy(self)
        new._bound_tools = list(tools)
        return new  # type: ignore[return-value]

    def get_info(self) -> BackendInfo:
        return BackendInfo(
            provider="openai",
            model=self._param.model,
            supports_vision=self._param.vision_enabled,
            supports_tools=True,
        )
