"""OpenAI Responses API provider with multimodal input support."""

from __future__ import annotations

import os
from typing import Any, AsyncIterator

import httpx

from logos_adk.errors import AuthenticationError, ProviderError
from logos_adk.providers import Provider
from logos_adk.providers.multimodal import openai_responses_message, parse_openai_responses_output
from logos_adk.types import MediaChunk, Message, Response, TextChunk, Usage


class OpenAIResponsesProvider(Provider):
    """Provider for OpenAI Responses API with native multimodal inputs."""

    _http_client_pool: dict[str | None, httpx.AsyncClient] = {}

    @classmethod
    def _shared_http_client(cls, proxy: str | None) -> httpx.AsyncClient:
        if proxy not in cls._http_client_pool:
            kwargs: dict[str, Any] = {"timeout": 120.0}
            if proxy:
                kwargs["proxy"] = proxy
            cls._http_client_pool[proxy] = httpx.AsyncClient(**kwargs)
        return cls._http_client_pool[proxy]

    def __init__(
        self,
        model: str = "gpt-4.1",
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        proxy: str | None = None,
        client: httpx.AsyncClient | None = None,
    ) -> None:
        resolved_key = api_key or os.environ.get("OPENAI_API_KEY")
        if not resolved_key:
            raise AuthenticationError(
                "OpenAI API key not found. Set OPENAI_API_KEY env var or pass api_key= to the provider."
            )
        self._model = model
        self._api_key = resolved_key
        self._base_url = (base_url or "https://api.openai.com/v1").rstrip("/")
        self.proxy = proxy
        self._client = client

    async def _post(self, payload: dict[str, Any]) -> dict[str, Any]:
        owns_client = self._client is None
        client = self._client or self._shared_http_client(self.proxy)
        try:
            response = await client.post(
                f"{self._base_url}/responses",
                headers={
                    "Authorization": f"Bearer {self._api_key}",
                    "Content-Type": "application/json",
                },
                json=payload,
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPError as exc:
            raise ProviderError(f"OpenAI Responses API error: {exc}") from exc
        finally:
            if owns_client and self._client is not None:
                await client.aclose()

    def _payload(
        self,
        messages: list[Message],
        *,
        tools: list[dict] | None,
        output_schema: Any,
        **kwargs: Any,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "model": self._model,
            "input": [],
        }
        for message in messages:
            serialized = openai_responses_message(message)
            if isinstance(serialized, list):
                payload["input"].extend(serialized)
            else:
                payload["input"].append(serialized)
        if tools:
            payload["tools"] = [
                {
                    "type": "function",
                    "name": tool["name"],
                    "description": tool.get("description", ""),
                    "parameters": tool.get("parameters", {"type": "object", "properties": {}}),
                }
                for tool in tools
            ]
        if output_schema is not None:
            payload["text"] = {"format": {"type": "json_object"}}
        payload.update(kwargs)
        return payload

    def _parse(self, data: dict[str, Any]) -> Response:
        output_items = data.get("output") or []
        text, parts, tool_calls = parse_openai_responses_output(output_items)
        usage = data.get("usage") or {}
        return Response(
            content=text,
            tool_calls=tool_calls,
            usage=Usage(
                input_tokens=int(usage.get("input_tokens") or 0),
                output_tokens=int(usage.get("output_tokens") or 0),
            ),
            raw=data,
            content_parts=parts,
        )

    async def generate(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        output_schema: Any = None,
        **kwargs: Any,
    ) -> Response:
        payload = self._payload(messages, tools=tools, output_schema=output_schema, **kwargs)
        data = await self._post(payload)
        return self._parse(data)

    async def stream(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[Any]:
        response = await self.generate(messages, tools=tools, **kwargs)
        if response.content:
            yield TextChunk(text=response.content)
        for part in response.content_parts:
            if part.type != "text":
                yield MediaChunk(part=part)
        yield response
