import httpx

from latch.models import (
    DecisionRequest,
    DecisionResponse,
    ExecuteResponse,
    LatchConfig,
    Message,
    ReportRequest,
    ReportResponse,
)
from latch.exceptions import (
    AuthenticationError,
    BudgetExceededError,
    LatchError,
    RateLimitError,
    RequestBlockedError,
)


class Latch:
    """Synchronous Latch client"""

    def __init__(
        self,
        *,
        api_key: str,
        base_url: str = "https://api.latch.dev",
        timeout: float = 30.0,
        mode: str = "observe",
    ):
        self.config = LatchConfig(
            api_key=api_key,
            base_url=base_url,
            timeout=timeout,
            mode=mode,
        )
        self._client = httpx.Client(
            base_url=base_url,
            timeout=timeout,
            headers={"Authorization": f"Bearer {api_key}"},
        )

    def decide(
        self,
        *,
        messages: list[dict],
        task: str,
        budget: float | None = None,
        max_latency: int | None = None,
        metadata: dict | None = None,
    ) -> DecisionResponse:
        """Get a routing decision without executing the LLM call"""
        request = DecisionRequest(
            messages=[Message(**m) for m in messages],
            task=task,
            budget=budget,
            max_latency=max_latency,
            metadata=metadata,
        )

        response = self._client.post(
            "/v1/decision",
            json=request.model_dump(exclude_none=True, by_alias=True),
        )

        self._handle_response_errors(response=response)

        data = response.json()["data"]
        return DecisionResponse(**data)

    def run(
        self,
        *,
        messages: list[dict],
        task: str,
        budget: float | None = None,
        max_latency: int | None = None,
        metadata: dict | None = None,
    ) -> ExecuteResponse:
        """Get a routing decision and execute the LLM call"""
        request = DecisionRequest(
            messages=[Message(**m) for m in messages],
            task=task,
            budget=budget,
            max_latency=max_latency,
            metadata=metadata,
        )

        response = self._client.post(
            "/v1/execute",
            json=request.model_dump(exclude_none=True, by_alias=True),
        )

        self._handle_response_errors(response=response)

        data = response.json()["data"]
        return ExecuteResponse(**data)

    def report(
        self,
        *,
        decision_id: str,
        input_tokens: int,
        output_tokens: int,
        llm_latency_ms: int,
        model: str | None = None,
        success: bool = True,
        error_message: str | None = None,
    ) -> ReportResponse:
        """Report usage after making your own LLM call"""
        request = ReportRequest(
            decision_id=decision_id,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            llm_latency_ms=llm_latency_ms,
            model=model,
            success=success,
            error_message=error_message,
        )

        response = self._client.post(
            "/v1/report",
            json=request.model_dump(exclude_none=True, by_alias=True),
        )

        self._handle_response_errors(response=response)

        data = response.json()["data"]
        return ReportResponse(**data)

    def _handle_response_errors(self, *, response: httpx.Response) -> None:
        if response.status_code == 200 or response.status_code == 201:
            return

        try:
            error_data = response.json().get("error", {})
            code = error_data.get("code", "")
            message = error_data.get("message", "Unknown error")
        except Exception:
            code = ""
            message = f"HTTP {response.status_code}"

        if response.status_code == 401:
            raise AuthenticationError(message)

        if response.status_code == 402:
            if code == "BUDGET_EXCEEDED":
                raise BudgetExceededError(message)
            raise RequestBlockedError(message)

        if response.status_code == 429:
            raise RateLimitError(message)

        raise LatchError(message, code=code)

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

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class AsyncLatch:
    """Asynchronous Latch client"""

    def __init__(
        self,
        *,
        api_key: str,
        base_url: str = "https://api.latch.dev",
        timeout: float = 30.0,
        mode: str = "observe",
    ):
        self.config = LatchConfig(
            api_key=api_key,
            base_url=base_url,
            timeout=timeout,
            mode=mode,
        )
        self._client = httpx.AsyncClient(
            base_url=base_url,
            timeout=timeout,
            headers={"Authorization": f"Bearer {api_key}"},
        )

    async def decide(
        self,
        *,
        messages: list[dict],
        task: str,
        budget: float | None = None,
        max_latency: int | None = None,
        metadata: dict | None = None,
    ) -> DecisionResponse:
        """Get a routing decision without executing the LLM call"""
        request = DecisionRequest(
            messages=[Message(**m) for m in messages],
            task=task,
            budget=budget,
            max_latency=max_latency,
            metadata=metadata,
        )

        response = await self._client.post(
            "/v1/decision",
            json=request.model_dump(exclude_none=True, by_alias=True),
        )

        self._handle_response_errors(response=response)

        data = response.json()["data"]
        return DecisionResponse(**data)

    async def run(
        self,
        *,
        messages: list[dict],
        task: str,
        budget: float | None = None,
        max_latency: int | None = None,
        metadata: dict | None = None,
    ) -> ExecuteResponse:
        """Get a routing decision and execute the LLM call"""
        request = DecisionRequest(
            messages=[Message(**m) for m in messages],
            task=task,
            budget=budget,
            max_latency=max_latency,
            metadata=metadata,
        )

        response = await self._client.post(
            "/v1/execute",
            json=request.model_dump(exclude_none=True, by_alias=True),
        )

        self._handle_response_errors(response=response)

        data = response.json()["data"]
        return ExecuteResponse(**data)

    async def report(
        self,
        *,
        decision_id: str,
        input_tokens: int,
        output_tokens: int,
        llm_latency_ms: int,
        model: str | None = None,
        success: bool = True,
        error_message: str | None = None,
    ) -> ReportResponse:
        """Report usage after making your own LLM call"""
        request = ReportRequest(
            decision_id=decision_id,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            llm_latency_ms=llm_latency_ms,
            model=model,
            success=success,
            error_message=error_message,
        )

        response = await self._client.post(
            "/v1/report",
            json=request.model_dump(exclude_none=True, by_alias=True),
        )

        self._handle_response_errors(response=response)

        data = response.json()["data"]
        return ReportResponse(**data)

    def _handle_response_errors(self, *, response: httpx.Response) -> None:
        if response.status_code == 200 or response.status_code == 201:
            return

        try:
            error_data = response.json().get("error", {})
            code = error_data.get("code", "")
            message = error_data.get("message", "Unknown error")
        except Exception:
            code = ""
            message = f"HTTP {response.status_code}"

        if response.status_code == 401:
            raise AuthenticationError(message)

        if response.status_code == 402:
            if code == "BUDGET_EXCEEDED":
                raise BudgetExceededError(message)
            raise RequestBlockedError(message)

        if response.status_code == 429:
            raise RateLimitError(message)

        raise LatchError(message, code=code)

    async def close(self) -> None:
        await self._client.aclose()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()
