"""OpenAI-compatible provider: direct httpx calls to any OpenAI-format API."""

from __future__ import annotations

import json
import time
from typing import Any, AsyncIterator

import httpx

from .base import LLMResponse, StreamChunk
from . import model_registry


class OpenAICompatibleProvider:
    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        provider_type: str | None = None,
    ):
        self.api_key = api_key
        self.base_url = (base_url or "https://api.openai.com/v1").rstrip("/")
        self.provider_type = provider_type

    def _headers(self) -> dict[str, str]:
        h = {"Content-Type": "application/json"}
        if self.api_key:
            h["Authorization"] = f"Bearer {self.api_key}"
        return h

    def _build_payload(self, model: str, messages: list[dict], **extra: Any) -> dict:
        payload: dict[str, Any] = {"model": model, "messages": messages}
        for key in ("temperature", "max_tokens", "top_p", "stop",
                    "response_format", "tools", "tool_choice"):
            if key in extra and extra[key] is not None:
                payload[key] = extra[key]
        if "extra_params" in extra and isinstance(extra["extra_params"], dict):
            payload.update(extra["extra_params"])
        if extra.get("stream"):
            payload["stream"] = True
            payload["stream_options"] = {"include_usage": True}
        return payload

    async def complete(
        self, messages: list[dict], model: str, **kwargs: Any
    ) -> LLMResponse:
        payload = self._build_payload(model, messages, **kwargs)
        start = time.perf_counter()

        async with httpx.AsyncClient() as client:
            resp = await client.post(
                f"{self.base_url}/chat/completions",
                json=payload,
                headers=self._headers(),
                timeout=300.0,
            )
            resp.raise_for_status()
            data = resp.json()

        duration_ms = (time.perf_counter() - start) * 1000

        content = data["choices"][0]["message"].get("content", "") or ""
        usage = data.get("usage", {})
        tokens_in = usage.get("prompt_tokens", 0)
        tokens_out = usage.get("completion_tokens", 0)
        cost = model_registry.compute_cost(model, tokens_in, tokens_out)
        tps = (tokens_out / (duration_ms / 1000)) if duration_ms > 0 else 0.0

        return LLMResponse(
            content=content,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            cost=cost,
            duration_ms=round(duration_ms, 1),
            tokens_per_sec=round(tps, 1),
            model=model,
            raw_response=data,
        )

    async def stream(
        self, messages: list[dict], model: str, **kwargs: Any
    ) -> AsyncIterator[StreamChunk]:
        payload = self._build_payload(model, messages, stream=True, **kwargs)
        start = time.perf_counter()
        chunks_text: list[str] = []
        tokens_in = 0
        tokens_out = 0

        async with httpx.AsyncClient() as client:
            async with client.stream(
                "POST",
                f"{self.base_url}/chat/completions",
                json=payload,
                headers=self._headers(),
                timeout=httpx.Timeout(300.0, connect=10.0),
            ) as resp:
                resp.raise_for_status()
                async for line in resp.aiter_lines():
                    if not line.startswith("data: "):
                        continue
                    if line.strip() == "data: [DONE]":
                        break

                    chunk = json.loads(line[6:])
                    delta = ""
                    choices = chunk.get("choices", [])
                    if choices and choices[0].get("delta", {}).get("content"):
                        delta = choices[0]["delta"]["content"]
                        chunks_text.append(delta)

                    usage = chunk.get("usage")
                    if usage:
                        tokens_in = usage.get("prompt_tokens", 0) or 0
                        tokens_out = usage.get("completion_tokens", 0) or 0

                    finish = choices[0].get("finish_reason") if choices else None
                    if finish:
                        duration_ms = (time.perf_counter() - start) * 1000
                        full_content = "".join(chunks_text)
                        if tokens_out == 0:
                            tokens_out = len(full_content) // 4
                        tps = (tokens_out / (duration_ms / 1000)) if duration_ms > 0 else 0.0
                        cost = model_registry.compute_cost(model, tokens_in, tokens_out)
                        yield StreamChunk(
                            delta=delta,
                            done=True,
                            response=LLMResponse(
                                content=full_content,
                                tokens_in=tokens_in,
                                tokens_out=tokens_out,
                                cost=cost,
                                duration_ms=round(duration_ms, 1),
                                tokens_per_sec=round(tps, 1),
                                model=model,
                            ),
                        )
                    elif delta:
                        yield StreamChunk(delta=delta)
