"""KairosRoute client — drop-in replacement for OpenAI's client."""

from __future__ import annotations

import json
import time
from typing import Any, Generator, Iterator

import httpx

from .exceptions import (
    AuthenticationError,
    BudgetExceededError,
    KairosRouteError,
    ProviderError,
    RateLimitError,
)

DEFAULT_BASE_URL = "https://kairosroute.com/api/v1"
DEFAULT_TIMEOUT = 60.0
MAX_RETRIES = 3


class _ChatCompletions:
    """Chat completions API."""

    def __init__(self, client: "KairosRoute"):
        self._client = client

    def create(
        self,
        *,
        model: str = "kr-auto",
        messages: list[dict[str, str]],
        temperature: float | None = None,
        max_tokens: int | None = None,
        stream: bool = False,
        top_p: float | None = None,
        frequency_penalty: float | None = None,
        presence_penalty: float | None = None,
        stop: str | list[str] | None = None,
        n: int | None = None,
        **kwargs: Any,
    ) -> dict[str, Any] | Iterator[dict[str, Any]]:
        body: dict[str, Any] = {"model": model, "messages": messages}
        if temperature is not None: body["temperature"] = temperature
        if max_tokens is not None: body["max_tokens"] = max_tokens
        if stream: body["stream"] = True
        if top_p is not None: body["top_p"] = top_p
        if frequency_penalty is not None: body["frequency_penalty"] = frequency_penalty
        if presence_penalty is not None: body["presence_penalty"] = presence_penalty
        if stop is not None: body["stop"] = stop
        if n is not None: body["n"] = n
        body.update(kwargs)

        if stream:
            return self._stream(body)
        return self._client._post("/chat/completions", body)

    def _stream(self, body: dict[str, Any]) -> Generator[dict[str, Any], None, None]:
        with self._client._http.stream("POST", f"{self._client.base_url}/chat/completions", json=body, headers=self._client._headers()) as response:
            self._client._check_response(response)
            for line in response.iter_lines():
                if line.startswith("data: "):
                    data = line[6:]
                    if data == "[DONE]":
                        return
                    try:
                        yield json.loads(data)
                    except json.JSONDecodeError:
                        continue


class _Completions:
    """Legacy completions API."""

    def __init__(self, client: "KairosRoute"):
        self._client = client

    def create(
        self,
        *,
        model: str = "kr-auto",
        prompt: str | list[str],
        max_tokens: int | None = None,
        temperature: float | None = None,
        stream: bool = False,
        **kwargs: Any,
    ) -> dict[str, Any]:
        body: dict[str, Any] = {"model": model, "prompt": prompt}
        if max_tokens is not None: body["max_tokens"] = max_tokens
        if temperature is not None: body["temperature"] = temperature
        if stream: body["stream"] = True
        body.update(kwargs)
        return self._client._post("/completions", body)


class _Embeddings:
    """Embeddings API."""

    def __init__(self, client: "KairosRoute"):
        self._client = client

    def create(
        self,
        *,
        input: str | list[str],
        model: str = "kr-auto-embed",
        encoding_format: str | None = None,
        dimensions: int | None = None,
    ) -> dict[str, Any]:
        body: dict[str, Any] = {"input": input, "model": model}
        if encoding_format: body["encoding_format"] = encoding_format
        if dimensions: body["dimensions"] = dimensions
        return self._client._post("/embeddings", body)


class _Chat:
    """Chat namespace containing completions."""

    def __init__(self, client: "KairosRoute"):
        self.completions = _ChatCompletions(client)


class _Models:
    """Models API."""

    def __init__(self, client: "KairosRoute"):
        self._client = client

    def list(self) -> dict[str, Any]:
        return self._client._get("/models")


class KairosRoute:
    """KairosRoute API client.

    Drop-in replacement for the OpenAI Python client.

    Usage:
        from kairosroute import KairosRoute

        client = KairosRoute(api_key="kr_live_...")

        # Same API as OpenAI
        response = client.chat.completions.create(
            model="kr-auto",  # or any specific model
            messages=[{"role": "user", "content": "Hello!"}],
        )
        print(response["choices"][0]["message"]["content"])
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = DEFAULT_BASE_URL,
        timeout: float = DEFAULT_TIMEOUT,
        max_retries: int = MAX_RETRIES,
    ):
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.max_retries = max_retries
        self._http = httpx.Client(timeout=timeout)

        # Namespaces
        self.chat = _Chat(self)
        self.completions = _Completions(self)
        self.embeddings = _Embeddings(self)
        self.models = _Models(self)

    def _headers(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "User-Agent": "kairosroute-python/0.1.0",
        }

    def _check_response(self, response: httpx.Response) -> None:
        if response.status_code == 200:
            return
        if response.status_code == 401:
            raise AuthenticationError("Invalid API key", status_code=401)
        if response.status_code == 403:
            raise AuthenticationError("Unauthorized", status_code=403)
        if response.status_code == 429:
            try:
                data = response.json()
                code = data.get("error", {}).get("code", "")
            except Exception:
                code = ""
            retry_after = response.headers.get("Retry-After")
            if code == "BUDGET_EXCEEDED":
                raise BudgetExceededError(
                    "Daily budget exceeded. Upgrade to Pro for unlimited usage.",
                    code="BUDGET_EXCEEDED",
                    status_code=429,
                )
            raise RateLimitError(
                "Rate limit exceeded",
                retry_after=float(retry_after) if retry_after else None,
            )
        if response.status_code >= 500:
            raise ProviderError(
                f"Provider error: {response.status_code}",
                status_code=response.status_code,
            )
        try:
            data = response.json()
            msg = data.get("error", {}).get("message", response.text)
        except Exception:
            msg = response.text
        raise KairosRouteError(msg, status_code=response.status_code)

    def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]:
        url = f"{self.base_url}{path}"
        last_error: Exception | None = None

        for attempt in range(self.max_retries):
            try:
                response = self._http.post(url, json=body, headers=self._headers())
                self._check_response(response)
                return response.json()
            except RateLimitError as e:
                last_error = e
                wait = e.retry_after or (2 ** attempt)
                time.sleep(wait)
            except (httpx.ConnectError, httpx.TimeoutException) as e:
                last_error = e
                time.sleep(2 ** attempt)
            except KairosRouteError:
                raise

        raise last_error or KairosRouteError("Request failed after retries")

    def _get(self, path: str) -> dict[str, Any]:
        url = f"{self.base_url}{path}"
        response = self._http.get(url, headers=self._headers())
        self._check_response(response)
        return response.json()

    def close(self) -> None:
        self._http.close()

    def __enter__(self) -> "KairosRoute":
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()
