import logging
import os
from typing import Any

import aiohttp
from pydantic import ValidationError as PydanticValidationError
from tenacity import (
    AsyncRetrying,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential_jitter,
)

from .errors import (
    ModelPricingError,
    NotFound,
    ServerError,
    Unauthorized,
    ValidationError,
)
from .models import EstimateResponse

logger = logging.getLogger(__name__)


class AsyncModelPricingClient:
    """Asynchronous client for the ModelPricing.ai API using aiohttp.

    Sends cost-estimation requests to the ``/v1/estimate`` endpoint and
    returns typed :class:`EstimateResponse` objects.  Retries automatically
    on transient server errors with exponential back-off.

    Supports ``async with`` for automatic session cleanup::

        async with AsyncModelPricingClient(api_key="sk-...") as client:
            estimate = await client.estimate(model="gpt-4o-mini", ...)

    @param api_key - API key (falls back to ``MODELPRICING_API_KEY`` env var).
    @param base_url - API base URL (falls back to ``MODELPRICING_BASE_URL``).
    @param timeout - Request timeout in seconds.
    @param session - Optional pre-configured ``aiohttp.ClientSession``.
    @param max_retries - Maximum retry attempts for transient errors.
    """

    def __init__(
        self,
        api_key: str | None = None,
        *,
        base_url: str | None = None,
        timeout: float = 30.0,
        session: aiohttp.ClientSession | None = None,
        max_retries: int = 3,
    ) -> None:
        resolved_key = api_key or os.environ.get("MODELPRICING_API_KEY")
        if not resolved_key or not isinstance(resolved_key, str):
            raise ValueError(
                "api_key is required — pass it as an argument or set the "
                "MODELPRICING_API_KEY environment variable"
            )
        self.api_key = resolved_key
        configured_base = base_url or os.environ.get(
            "MODELPRICING_BASE_URL", "https://api.modelpricing.ai"
        )
        self.base_url = configured_base.rstrip("/")
        self.timeout = timeout
        self._owns_session = session is None
        self._session = session
        self.max_retries = max_retries

    # -- context manager --------------------------------------------------

    async def __aenter__(self) -> "AsyncModelPricingClient":
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        await self.aclose()

    # -- internals --------------------------------------------------------

    def _headers(self) -> dict[str, str]:
        """Build the default request headers."""
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        }

    async def _ensure_session(self) -> aiohttp.ClientSession:
        """Lazily create the underlying ``aiohttp.ClientSession``."""
        if self._session is None:
            timeout = aiohttp.ClientTimeout(total=self.timeout)
            self._session = aiohttp.ClientSession(timeout=timeout)
        return self._session

    async def aclose(self) -> None:
        """Close the underlying HTTP session if it is owned by this client."""
        if self._owns_session and self._session is not None and not self._session.closed:
            await self._session.close()

    def _get_retry_handler(self) -> AsyncRetrying:
        """Build the tenacity async-retry configuration."""
        return AsyncRetrying(
            retry=retry_if_exception_type((ServerError, aiohttp.ClientError)),
            stop=stop_after_attempt(self.max_retries + 1),
            wait=wait_exponential_jitter(initial=0.1, max=2, jitter=0.5),
            reraise=True,
        )

    async def _handle_response(
        self, resp: aiohttp.ClientResponse
    ) -> EstimateResponse:
        """Parse an HTTP response into an :class:`EstimateResponse`.

        @param resp - The raw ``aiohttp.ClientResponse``.
        @returns Validated estimate response.
        @throws Unauthorized on 401.
        @throws NotFound on 404.
        @throws ValidationError on 422.
        @throws ServerError on 5xx.
        @throws ModelPricingError on other non-200 status codes.
        """
        status = resp.status
        if status == 200:
            data = await resp.json()
            try:
                return EstimateResponse.model_validate(data)
            except PydanticValidationError as exc:
                logger.warning("Failed to validate estimate response: %s", exc)
                raise ModelPricingError(
                    f"Unexpected response structure: {exc}",
                    status_code=200,
                ) from exc
        if status == 401:
            raise Unauthorized("Unauthorized", status_code=401)
        if status == 404:
            raise NotFound("Not found", status_code=404)
        if status == 422:
            try:
                data = await resp.json()
                message = data.get("error") or "Unprocessable Entity"
                details = data.get("details")
                if details:
                    message = f"{message}: {details}"
            except (ValueError, KeyError):
                message = "Unprocessable Entity"
            raise ValidationError(message, status_code=422)

        try:
            data = await resp.json()
            message = data.get("error") or await resp.text()
        except (ValueError, KeyError):
            message = await resp.text()
        if 500 <= status <= 599:
            raise ServerError(message, status_code=status)
        raise ModelPricingError(message, status_code=status)

    # -- public API -------------------------------------------------------

    async def estimate(
        self,
        *,
        model: str,
        tokens_in: int,
        tokens_out: int,
        trace_id: dict[str, Any] | None = None,
    ) -> EstimateResponse:
        """Estimate the cost of an LLM API call.

        Calls ``POST /v1/estimate`` and returns a validated response.
        Retries automatically on 5xx errors with exponential back-off.

        @param model - Canonical model identifier (e.g. ``"gpt-4o-mini"``).
        @param tokens_in - Number of input tokens (must be >= 0).
        @param tokens_out - Number of output tokens (must be >= 0).
        @param trace_id - Optional pass-through trace data returned in the response.
        @returns Validated :class:`EstimateResponse` with cost breakdown.
        @throws ValueError if model is empty or token counts are negative.
        @throws Unauthorized on invalid API key (401).
        @throws ValidationError on unknown model or bad metrics (422).
        @throws ServerError on 5xx after retries are exhausted.
        """
        if not isinstance(model, str) or not model.strip():
            raise ValueError("model must be a non-empty string")
        if not isinstance(tokens_in, (int, float)) or tokens_in < 0:
            raise ValueError("tokens_in must be a non-negative number")
        if not isinstance(tokens_out, (int, float)) or tokens_out < 0:
            raise ValueError("tokens_out must be a non-negative number")

        async for attempt in self._get_retry_handler():
            with attempt:
                url = f"{self.base_url}/v1/estimate"
                payload: dict[str, Any] = {
                    "model": model,
                    "metrics": {
                        "tokensIn": int(tokens_in),
                        "tokensOut": int(tokens_out),
                    },
                }
                if trace_id is not None:
                    payload["traceId"] = trace_id

                logger.debug("POST %s %s", url, payload)
                session = await self._ensure_session()
                async with session.post(
                    url, json=payload, headers=self._headers()
                ) as resp:
                    return await self._handle_response(resp)
