"""OpenAI-compatible client (OpenAI, OpenRouter, Groq, etc.)."""

from __future__ import annotations

from collections.abc import AsyncIterator
from typing import Any

import httpx

from llm_rotator._types import LLMResponse, StreamChunk, ToolCall, Usage
from llm_rotator.clients import AbstractLLMClient
from llm_rotator.exceptions import KeyDeadError, ModelRateLimitError, ServerError

_DEFAULT_BASE_URL = "https://api.openai.com/v1"


class OpenAIClient(AbstractLLMClient):
    """Client for OpenAI-compatible APIs."""

    def __init__(self, timeout: float = 60.0) -> None:
        self._timeout = timeout

    async def generate(
        self,
        messages: list[dict],
        model: str,
        api_key: str,
        base_url: str | None = None,
        **kwargs: Any,
    ) -> LLMResponse:
        url = f"{base_url or _DEFAULT_BASE_URL}/chat/completions"
        headers = {
            "authorization": f"Bearer {api_key}",
            "content-type": "application/json",
        }
        body: dict[str, Any] = {"model": model, "messages": messages}

        # Forward tools, tool_choice, response_format
        if kwargs.get("tools"):
            body["tools"] = kwargs["tools"]
        if kwargs.get("tool_choice") is not None:
            body["tool_choice"] = kwargs["tool_choice"]
        if kwargs.get("response_format") is not None:
            body["response_format"] = _build_response_format(kwargs["response_format"])

        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.post(url, json=body, headers=headers)

        _check_status(resp, model=model, api_key=api_key)

        data = resp.json()
        choice = data["choices"][0]
        message = choice["message"]
        usage_data = data.get("usage", {})

        tool_calls = _parse_tool_calls(message.get("tool_calls"))

        return LLMResponse(
            content=message.get("content"),
            usage=Usage(
                prompt_tokens=usage_data.get("prompt_tokens", 0),
                completion_tokens=usage_data.get("completion_tokens", 0),
                total_tokens=usage_data.get("total_tokens", 0),
            ),
            model=data.get("model", model),
            provider="openai",
            key_alias=api_key[:8] + "...",
            tool_calls=tool_calls,
            raw=data,
        )

    async def stream(
        self,
        messages: list[dict],
        model: str,
        api_key: str,
        base_url: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[StreamChunk]:
        url = f"{base_url or _DEFAULT_BASE_URL}/chat/completions"
        headers = {
            "authorization": f"Bearer {api_key}",
            "content-type": "application/json",
        }
        body: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": True,
            "stream_options": {"include_usage": True},
        }

        if kwargs.get("tools"):
            body["tools"] = kwargs["tools"]
        if kwargs.get("tool_choice") is not None:
            body["tool_choice"] = kwargs["tool_choice"]
        if kwargs.get("response_format") is not None:
            body["response_format"] = _build_response_format(kwargs["response_format"])

        # Accumulate tool calls across stream chunks
        pending_tool_calls: dict[int, dict] = {}  # index → {id, name, arguments}

        async with (
            httpx.AsyncClient(timeout=self._timeout) as client,
            client.stream("POST", url, json=body, headers=headers) as resp,
        ):
                if resp.status_code != 200:
                    await resp.aread()
                    _check_status(resp, model=model, api_key=api_key)

                async for line in resp.aiter_lines():
                    if not line.startswith("data: "):
                        continue
                    payload = line[6:].strip()
                    if payload == "[DONE]":
                        # Emit accumulated tool calls in done chunk
                        final_tool_calls = _finalize_tool_calls(pending_tool_calls)
                        yield StreamChunk(
                            delta="", done=True, tool_calls=final_tool_calls
                        )
                        return

                    import json

                    data = json.loads(payload)
                    delta = data.get("choices", [{}])[0].get("delta", {})
                    content = delta.get("content", "")

                    # Accumulate streamed tool_calls
                    for tc in delta.get("tool_calls", []):
                        idx = tc["index"]
                        if idx not in pending_tool_calls:
                            pending_tool_calls[idx] = {
                                "id": "",
                                "name": "",
                                "arguments": "",
                            }
                        if tc.get("id"):
                            pending_tool_calls[idx]["id"] = tc["id"]
                        fn = tc.get("function", {})
                        if fn.get("name"):
                            pending_tool_calls[idx]["name"] = fn["name"]
                        if fn.get("arguments"):
                            pending_tool_calls[idx]["arguments"] += fn["arguments"]

                    usage_data = data.get("usage")
                    usage = None
                    if usage_data:
                        usage = Usage(
                            prompt_tokens=usage_data.get("prompt_tokens", 0),
                            completion_tokens=usage_data.get("completion_tokens", 0),
                            total_tokens=usage_data.get("total_tokens", 0),
                        )

                    if content:
                        yield StreamChunk(delta=content, usage=usage)


def _check_status(
    resp: httpx.Response, *, model: str, api_key: str
) -> None:
    """Raise typed exceptions based on HTTP status code."""
    if resp.status_code < 400:
        return

    alias = api_key[:8] + "..."
    status = resp.status_code

    if status in (401, 403, 402):
        raise KeyDeadError(key_alias=alias, status_code=status)
    if status == 429:
        retry_after_raw = resp.headers.get("retry-after")
        retry_after = int(retry_after_raw) if retry_after_raw else None
        raise ModelRateLimitError(
            key_alias=alias, model=model, retry_after=retry_after
        )
    if status >= 500:
        raise ServerError(provider="openai", status_code=status)

    raise ServerError(
        provider="openai",
        status_code=status,
        message=f"Unexpected HTTP {status}",
    )


def _parse_tool_calls(raw: list[dict] | None) -> list[ToolCall] | None:
    """Parse OpenAI tool_calls from response message."""
    if not raw:
        return None
    return [
        ToolCall(
            id=tc["id"],
            name=tc["function"]["name"],
            arguments=tc["function"]["arguments"],
        )
        for tc in raw
    ]


def _finalize_tool_calls(
    pending: dict[int, dict],
) -> list[ToolCall] | None:
    """Convert accumulated stream tool call fragments into ToolCall list."""
    if not pending:
        return None
    return [
        ToolCall(
            id=pending[idx]["id"],
            name=pending[idx]["name"],
            arguments=pending[idx]["arguments"],
        )
        for idx in sorted(pending)
    ]


def _build_response_format(fmt: Any) -> dict:
    """Convert response_format to OpenAI API format."""
    if isinstance(fmt, dict):
        return fmt
    # Pydantic BaseModel → JSON Schema
    if hasattr(fmt, "model_json_schema"):
        return {
            "type": "json_schema",
            "json_schema": {
                "name": fmt.__name__,
                "schema": fmt.model_json_schema(),
            },
        }
    return fmt
