"""
Latch Context Manager

Provides a clean interface for getting routing decisions and reporting analytics
without interfering with your existing LLM code.

Example:
    from latch import latch

    async def research_query(query: str):
        async with latch(task="legal-research", budget=0.10) as ctx:
            # ctx.model_name is the recommended model (e.g., "claude-sonnet-4-5-20250514")
            agent = Agent(model=get_model(ctx.model_name), ...)
            result = await agent.run(query)
            # Report usage for analytics
            ctx.record_usage(input_tokens=result.usage.input_tokens,
                           output_tokens=result.usage.output_tokens)
        return result
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Any

import httpx

from latch.models import DecisionResponse, ReportResponse
from latch.exceptions import LatchError, AuthenticationError


@dataclass
class LatchContext:
    """
    Context manager for Latch routing decisions.

    Handles the decide/report lifecycle automatically:
    - On enter: Calls /v1/decision to get routing recommendation
    - On exit: Calls /v1/report to send analytics

    Attributes:
        model_name: The recommended model name (e.g., "gpt-4o-mini", "claude-sonnet-4-5-20250514")
        decision_id: Unique ID for this decision (used for analytics correlation)
        should_call_llm: Whether Latch recommends making the LLM call
        estimated_cost: Estimated cost for this request
    """

    api_key: str
    base_url: str = "http://localhost:8081"
    task: str = "general"
    budget: float | None = None
    max_latency: int | None = None
    metadata: dict[str, Any] | None = None
    timeout: float = 30.0

    decision: DecisionResponse | None = field(default=None, init=False)
    _client: httpx.AsyncClient | None = field(default=None, init=False, repr=False)
    _start_time: float = field(default=0.0, init=False, repr=False)
    _usage: dict[str, Any] = field(default_factory=dict, init=False, repr=False)

    @property
    def model_name(self) -> str:
        """The recommended model name from Latch's routing decision"""
        if self.decision is None:
            raise RuntimeError("Context not entered - use 'async with latch(...) as ctx:'")
        return self.decision.model

    @property
    def decision_id(self) -> str:
        """Unique ID for this decision"""
        if self.decision is None:
            raise RuntimeError("Context not entered - use 'async with latch(...) as ctx:'")
        return self.decision.decision_id

    @property
    def should_call_llm(self) -> bool:
        """Whether Latch recommends making the LLM call"""
        if self.decision is None:
            raise RuntimeError("Context not entered - use 'async with latch(...) as ctx:'")
        return self.decision.should_call_llm

    @property
    def estimated_cost(self) -> float:
        """Estimated cost for this request"""
        if self.decision is None:
            raise RuntimeError("Context not entered - use 'async with latch(...) as ctx:'")
        return self.decision.estimated_cost

    def record_usage(
        self,
        *,
        input_tokens: int,
        output_tokens: int,
        model: str | None = None,
    ) -> None:
        """
        Record token usage for analytics reporting.

        Call this after your LLM call completes to track actual usage.
        If not called, the context will still report but without token counts.

        Args:
            input_tokens: Number of input tokens used
            output_tokens: Number of output tokens generated
            model: Optional override for the model name (if different from recommendation)
        """
        self._usage = {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "model": model,
        }

    async def __aenter__(self) -> LatchContext:
        self._start_time = time.perf_counter()

        self._client = httpx.AsyncClient(
            base_url=self.base_url,
            timeout=self.timeout,
            headers={"Authorization": f"Bearer {self.api_key}"},
        )

        request_body: dict[str, Any] = {
            "messages": [{"role": "user", "content": "placeholder"}],
            "task": self.task,
        }
        if self.budget is not None:
            request_body["budget"] = self.budget
        if self.max_latency is not None:
            request_body["maxLatency"] = self.max_latency
        if self.metadata is not None:
            request_body["metadata"] = self.metadata

        response = await self._client.post("/v1/decision", json=request_body)

        if response.status_code == 401:
            raise AuthenticationError("Invalid API key")
        if response.status_code != 200:
            raise LatchError(f"Decision request failed: {response.status_code}")

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

        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        if self._client is None or self.decision is None:
            return

        llm_latency_ms = int((time.perf_counter() - self._start_time) * 1000)

        try:
            report_body: dict[str, Any] = {
                "decisionId": self.decision.decision_id,
                "inputTokens": self._usage.get("input_tokens", 0),
                "outputTokens": self._usage.get("output_tokens", 0),
                "llmLatencyMs": llm_latency_ms,
                "success": exc_type is None,
            }

            if self._usage.get("model"):
                report_body["model"] = self._usage["model"]

            if exc_val is not None:
                report_body["errorMessage"] = str(exc_val)

            await self._client.post("/v1/report", json=report_body)
        except Exception:
            pass
        finally:
            await self._client.aclose()


def latch(
    *,
    api_key: str | None = None,
    base_url: str = "http://localhost:8081",
    task: str = "general",
    budget: float | None = None,
    max_latency: int | None = None,
    metadata: dict[str, Any] | None = None,
    timeout: float = 30.0,
) -> LatchContext:
    """
    Create a Latch context for routing decisions and analytics.

    Usage:
        async with latch(api_key="...", task="legal-research") as ctx:
            # Use ctx.model_name to get the recommended model
            model = get_pydantic_ai_model(ctx.model_name)
            result = await agent.run(query)
            ctx.record_usage(input_tokens=..., output_tokens=...)

    Args:
        api_key: Your Latch API key (or set LATCH_API_KEY env var)
        base_url: Latch API URL (default: http://localhost:8081)
        task: Task type for routing (e.g., "legal-research", "code-generation")
        budget: Optional budget limit for this request
        max_latency: Optional maximum latency in milliseconds
        metadata: Optional metadata to attach to this request
        timeout: HTTP timeout in seconds

    Returns:
        LatchContext: An async context manager
    """
    import os

    resolved_api_key = api_key or os.environ.get("LATCH_API_KEY")
    if not resolved_api_key:
        raise ValueError("api_key is required (or set LATCH_API_KEY env var)")

    return LatchContext(
        api_key=resolved_api_key,
        base_url=base_url,
        task=task,
        budget=budget,
        max_latency=max_latency,
        metadata=metadata,
        timeout=timeout,
    )
