import logging
import os
from collections.abc import Callable
from typing import Any, TypeVar

import requests
from pydantic import ValidationError as PydanticValidationError
from tenacity import (
    retry,
    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__)

F = TypeVar("F", bound=Callable[..., Any])


class ModelPricingClient:
    """Synchronous client for the ModelPricing.ai API.

    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.

    @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 ``requests.Session``.
    @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: requests.Session | 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 or requests.Session()
        self.max_retries = max_retries

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

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

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if self._owns_session:
            self._session.close()

    # -- 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",
        }

    def _retry_on_server_error(self, func: F) -> F:
        """Wrap *func* with tenacity retry on server/connection errors."""
        return retry(
            retry=retry_if_exception_type(
                (ServerError, requests.exceptions.RequestException)
            ),
            stop=stop_after_attempt(self.max_retries + 1),
            wait=wait_exponential_jitter(initial=0.1, max=2, jitter=0.5),
            reraise=True,
        )(func)

    def _handle_response(self, response: requests.Response) -> EstimateResponse:
        """Parse an HTTP response into an :class:`EstimateResponse`.

        @param response - The raw ``requests.Response``.
        @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.
        """
        if response.status_code == 200:
            data = response.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 response.status_code == 401:
            raise Unauthorized("Unauthorized", status_code=401)
        if response.status_code == 404:
            raise NotFound("Not found", status_code=404)
        if response.status_code == 422:
            try:
                data = response.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 = response.json()
            message = data.get("error") or response.text or "HTTP error"
        except (ValueError, KeyError):
            message = response.text or "HTTP error"

        if 500 <= response.status_code <= 599:
            raise ServerError(message, status_code=response.status_code)
        raise ModelPricingError(message, status_code=response.status_code)

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

    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")

        def _make_request() -> EstimateResponse:
            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)
            response = self._session.post(
                url, json=payload, headers=self._headers(), timeout=self.timeout
            )
            return self._handle_response(response)

        return self._retry_on_server_error(_make_request)()
