"""Usage tracking for monitoring costs and limits."""

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

from pydantic_ai.usage import Usage


@dataclass
class UsageRecord:
    """Record of usage for a single request."""

    timestamp: datetime
    user_id: str
    agent_name: str
    model: str
    request_tokens: int
    response_tokens: int
    total_tokens: int
    requests: int
    cost_estimate: float = 0.0


@dataclass
class UsageSummary:
    """Summary of usage over a time period."""

    total_requests: int = 0
    total_tokens: int = 0
    request_tokens: int = 0
    response_tokens: int = 0
    estimated_cost: float = 0.0
    by_model: dict[str, int] = field(default_factory=dict)
    by_agent: dict[str, int] = field(default_factory=dict)


class UsageTracker:
    """Track and report on LLM usage.

    Helps monitor costs, enforce limits, and analyze usage patterns.

    Example:
        ```python
        tracker = UsageTracker(database_pool)

        # Track usage after agent run
        await tracker.record_usage(
            user_id="user-123",
            agent_name="research_agent",
            model="gpt-4o",
            usage=result.usage(),
        )

        # Get usage summary
        summary = await tracker.get_user_summary("user-123")
        print(f"Total tokens: {summary.total_tokens}")
        ```
    """

    # Approximate costs per 1K tokens (update as needed)
    COST_PER_1K_TOKENS = {
        "gpt-4o": {"input": 0.005, "output": 0.015},
        "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
        "claude-sonnet-4-5": {"input": 0.003, "output": 0.015},
        "claude-3-5-haiku": {"input": 0.0008, "output": 0.004},
    }

    def __init__(self, db: Any) -> None:
        """Initialize the usage tracker.

        Args:
            db: Database connection pool.
        """
        self.db = db

    def estimate_cost(
        self,
        model: str,
        request_tokens: int,
        response_tokens: int,
    ) -> float:
        """Estimate cost for a request.

        Args:
            model: Model name.
            request_tokens: Number of input tokens.
            response_tokens: Number of output tokens.

        Returns:
            Estimated cost in USD.
        """
        # Find matching cost config
        costs = None
        for model_key in self.COST_PER_1K_TOKENS:
            if model_key in model:
                costs = self.COST_PER_1K_TOKENS[model_key]
                break

        if not costs:
            return 0.0

        input_cost = (request_tokens / 1000) * costs["input"]
        output_cost = (response_tokens / 1000) * costs["output"]

        return input_cost + output_cost

    async def record_usage(
        self,
        user_id: str,
        agent_name: str,
        model: str,
        usage: Usage,
    ) -> UsageRecord:
        """Record usage from an agent run.

        Args:
            user_id: The user's ID.
            agent_name: Name of the agent that was run.
            model: Model that was used.
            usage: Usage data from the agent run.

        Returns:
            The created usage record.
        """
        request_tokens = usage.request_tokens or 0
        response_tokens = usage.response_tokens or 0
        total_tokens = usage.total_tokens or (request_tokens + response_tokens)
        requests = usage.requests or 1

        cost = self.estimate_cost(model, request_tokens, response_tokens)

        record = UsageRecord(
            timestamp=datetime.utcnow(),
            user_id=user_id,
            agent_name=agent_name,
            model=model,
            request_tokens=request_tokens,
            response_tokens=response_tokens,
            total_tokens=total_tokens,
            requests=requests,
            cost_estimate=cost,
        )

        await self.db.execute(
            """
            INSERT INTO usage_records (
                user_id, agent_name, model,
                request_tokens, response_tokens, total_tokens,
                requests, cost_estimate, created_at
            ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, NOW())
            """,
            record.user_id,
            record.agent_name,
            record.model,
            record.request_tokens,
            record.response_tokens,
            record.total_tokens,
            record.requests,
            record.cost_estimate,
        )

        return record

    async def get_user_summary(
        self,
        user_id: str,
        days: int = 30,
    ) -> UsageSummary:
        """Get usage summary for a user.

        Args:
            user_id: The user's ID.
            days: Number of days to summarize.

        Returns:
            UsageSummary with aggregated data.
        """
        result = await self.db.fetch_one(
            f"""
            SELECT
                COUNT(*) as total_requests,
                SUM(total_tokens) as total_tokens,
                SUM(request_tokens) as request_tokens,
                SUM(response_tokens) as response_tokens,
                SUM(cost_estimate) as estimated_cost
            FROM usage_records
            WHERE user_id = $1
              AND created_at > NOW() - INTERVAL '{days} days'
            """,
            user_id,
        )

        summary = UsageSummary(
            total_requests=result["total_requests"] or 0,
            total_tokens=result["total_tokens"] or 0,
            request_tokens=result["request_tokens"] or 0,
            response_tokens=result["response_tokens"] or 0,
            estimated_cost=result["estimated_cost"] or 0.0,
        )

        # Get breakdown by model
        model_rows = await self.db.fetch_all(
            f"""
            SELECT model, SUM(total_tokens) as tokens
            FROM usage_records
            WHERE user_id = $1
              AND created_at > NOW() - INTERVAL '{days} days'
            GROUP BY model
            """,
            user_id,
        )
        summary.by_model = {row["model"]: row["tokens"] for row in model_rows}

        # Get breakdown by agent
        agent_rows = await self.db.fetch_all(
            f"""
            SELECT agent_name, SUM(total_tokens) as tokens
            FROM usage_records
            WHERE user_id = $1
              AND created_at > NOW() - INTERVAL '{days} days'
            GROUP BY agent_name
            """,
            user_id,
        )
        summary.by_agent = {row["agent_name"]: row["tokens"] for row in agent_rows}

        return summary

    async def check_limit(
        self,
        user_id: str,
        token_limit: int,
        period_days: int = 30,
    ) -> tuple[bool, int]:
        """Check if user is within usage limits.

        Args:
            user_id: The user's ID.
            token_limit: Maximum allowed tokens.
            period_days: Period to check.

        Returns:
            Tuple of (is_within_limit, tokens_used).
        """
        result = await self.db.fetch_one(
            f"""
            SELECT COALESCE(SUM(total_tokens), 0) as tokens_used
            FROM usage_records
            WHERE user_id = $1
              AND created_at > NOW() - INTERVAL '{period_days} days'
            """,
            user_id,
        )

        tokens_used = result["tokens_used"]
        return tokens_used < token_limit, tokens_used
