import os
from typing import Any, Dict, Optional, Union

import aiohttp

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


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

    def __init__(
        self,
        api_key: str,
        *,
        base_url: Optional[str] = None,
        timeout: float = 30.0,
        session: Optional[aiohttp.ClientSession] = None,
    ) -> None:
        if not api_key or not isinstance(api_key, str):
            raise ValueError("api_key is required")
        self.api_key = api_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._session = session

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

    async def _ensure_session(self) -> 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:
        if self._session is not None and not self._session.closed:
            await self._session.close()

    async def estimate(
        self,
        *,
        model: str,
        tokens_in: int,
        tokens_out: int,
        trace_id: Optional[Dict[str, Any]] = None,
    ) -> Union[Dict[str, Any], EstimateResponse]:
        if not model:
            raise ValueError("model is required")
        if tokens_in is None or tokens_out is None:
            raise ValueError("tokens_in and tokens_out are required")

        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

        session = await self._ensure_session()
        async with session.post(url, json=payload, headers=self._headers()) as resp:
            status = resp.status
            if status == 200:
                data = await resp.json()
                try:
                    return EstimateResponse.model_validate(data)
                except Exception:
                    return data
            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 Exception:
                    message = "Unprocessable Entity"
                raise ValidationError(message, status_code=422)

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


