"""HTTP client for the Traces API."""

from __future__ import annotations

from datetime import datetime
from typing import Any

import httpx

from mcp_traces.schemas import (
    AnalyticsSummary,
    ModelMetrics,
    SpanDetail,
    TraceDetail,
    TracesResponse,
    SpanStatus,
)


class TracesClient:
    """Client for interacting with the Traces API.

    Authenticates using an API key passed via the X-API-Key header.
    """

    def __init__(
        self,
        base_url: str,
        api_key: str,
        timeout: float = 30.0,
    ):
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.timeout = timeout
        self._client: httpx.AsyncClient | None = None

    @property
    def client(self) -> httpx.AsyncClient:
        if self._client is None:
            self._client = httpx.AsyncClient(
                base_url=self.base_url,
                headers={
                    "X-API-Key": self.api_key,
                    "Content-Type": "application/json",
                },
                timeout=self.timeout,
            )
        return self._client

    async def close(self) -> None:
        """Close the HTTP client."""
        if self._client:
            await self._client.aclose()
            self._client = None

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

    async def __aexit__(self, *args: Any) -> None:
        await self.close()

    async def list_traces(
        self,
        *,
        status: SpanStatus | None = None,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
        search: str | None = None,
        session_id: str | None = None,
        user_id: str | None = None,
        project_id: str | None = None,
        tags: list[str] | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> TracesResponse:
        """List traces with optional filters.

        Args:
            status: Filter by trace status (SUCCESS, ERROR, RUNNING)
            start_time: Filter traces starting after this time
            end_time: Filter traces starting before this time
            search: Full-text search query
            session_id: Filter by session ID
            user_id: Filter by user ID
            project_id: Filter by project ID
            tags: Filter by tags (traces must have ALL specified tags)
            limit: Maximum number of traces to return (1-100)
            offset: Offset for pagination

        Returns:
            TracesResponse with list of traces and pagination info
        """
        params: dict[str, Any] = {"limit": limit, "offset": offset}

        if status:
            params["status"] = status
        if start_time:
            params["start_time"] = start_time.isoformat()
        if end_time:
            params["end_time"] = end_time.isoformat()
        if search:
            params["search"] = search
        if session_id:
            params["session_id"] = session_id
        if user_id:
            params["user_id"] = user_id
        if project_id:
            params["project_id"] = project_id
        if tags:
            params["tags"] = ",".join(tags)

        response = await self.client.get("/traces", params=params)
        response.raise_for_status()
        return TracesResponse.model_validate(response.json())

    async def get_trace(self, trace_id: str) -> TraceDetail:
        """Get full trace details including all spans.

        Args:
            trace_id: The trace ID to fetch

        Returns:
            TraceDetail with full span data including input/output
        """
        response = await self.client.get(f"/traces/{trace_id}")
        response.raise_for_status()
        return TraceDetail.model_validate(response.json())

    async def get_trace_summary(self, trace_id: str) -> TraceDetail:
        """Get trace with lightweight span summaries (faster).

        Args:
            trace_id: The trace ID to fetch

        Returns:
            TraceDetail with spans (no input/output for faster loading)
        """
        response = await self.client.get(f"/traces/{trace_id}/summary")
        response.raise_for_status()
        return TraceDetail.model_validate(response.json())

    async def get_span(self, span_id: str) -> SpanDetail:
        """Get full span details including input/output.

        Args:
            span_id: The span ID to fetch

        Returns:
            SpanDetail with full input/output data
        """
        response = await self.client.get(f"/traces/spans/{span_id}")
        response.raise_for_status()
        data = response.json()
        return SpanDetail.model_validate(data.get("span", data))

    async def get_analytics(
        self,
        *,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
        granularity: str = "day",
        project_id: str | None = None,
    ) -> AnalyticsSummary:
        """Get analytics summary with time series data.

        Args:
            start_time: Start of time range
            end_time: End of time range
            granularity: Time series granularity (hour, day, week)
            project_id: Filter by project ID

        Returns:
            AnalyticsSummary with aggregated metrics and time series
        """
        params: dict[str, Any] = {"granularity": granularity}

        if start_time:
            params["start_time"] = start_time.isoformat()
        if end_time:
            params["end_time"] = end_time.isoformat()
        if project_id:
            params["project_id"] = project_id

        response = await self.client.get("/traces/analytics/summary", params=params)
        response.raise_for_status()
        return AnalyticsSummary.model_validate(response.json())

    async def get_model_analytics(
        self,
        project_id: str | None = None,
    ) -> list[ModelMetrics]:
        """Get per-model performance metrics.

        Args:
            project_id: Filter by project ID

        Returns:
            List of ModelMetrics with per-model stats
        """
        params: dict[str, Any] = {}
        if project_id:
            params["project_id"] = project_id

        response = await self.client.get("/traces/analytics/models", params=params)
        response.raise_for_status()
        data = response.json()
        models = data.get("models", data) if isinstance(data, dict) else data
        return [ModelMetrics.model_validate(m) for m in models]
