"""Token usage tracking and cost estimation."""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from datetime import datetime, timezone


# Pricing table: model -> {input_cost_per_1k, output_cost_per_1k} in USD
COST_PER_TOKEN: dict[str, dict[str, float]] = {
    "claude-sonnet-4-20250514": {"input_cost_per_1k": 0.003, "output_cost_per_1k": 0.015},
    "claude-opus-4-20250514": {"input_cost_per_1k": 0.015, "output_cost_per_1k": 0.075},
    "claude-3.5-sonnet": {"input_cost_per_1k": 0.003, "output_cost_per_1k": 0.015},
    "claude-3-opus": {"input_cost_per_1k": 0.015, "output_cost_per_1k": 0.075},
    "claude-3-haiku": {"input_cost_per_1k": 0.00025, "output_cost_per_1k": 0.00125},
    "gpt-4o": {"input_cost_per_1k": 0.005, "output_cost_per_1k": 0.015},
    "gpt-4": {"input_cost_per_1k": 0.03, "output_cost_per_1k": 0.06},
    "gpt-3.5-turbo": {"input_cost_per_1k": 0.0005, "output_cost_per_1k": 0.0015},
}


@dataclass
class TokenUsage:
    """A single token usage record from a tool invocation."""

    tool: str
    input_tokens: int
    output_tokens: int
    total_tokens: int
    estimated_cost_usd: float
    model: str
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


class CostTracker:
    """Accumulates TokenUsage records for a session."""

    def __init__(self) -> None:
        self._records: list[TokenUsage] = []

    def record(self, usage: TokenUsage) -> None:
        """Add a token usage record."""
        self._records.append(usage)

    @property
    def records(self) -> list[TokenUsage]:
        """Return all recorded usage entries."""
        return list(self._records)

    def session_total(self) -> dict:
        """Aggregate all usage for the session.

        Returns a dict with:
        - total_input_tokens
        - total_output_tokens
        - total_tokens
        - total_cost_usd
        - by_tool: {tool_name: {input_tokens, output_tokens, total_tokens, cost_usd}}
        """
        total_input = 0
        total_output = 0
        total_tokens = 0
        total_cost = 0.0
        by_tool: dict[str, dict] = {}

        for r in self._records:
            total_input += r.input_tokens
            total_output += r.output_tokens
            total_tokens += r.total_tokens
            total_cost += r.estimated_cost_usd

            if r.tool not in by_tool:
                by_tool[r.tool] = {
                    "input_tokens": 0,
                    "output_tokens": 0,
                    "total_tokens": 0,
                    "cost_usd": 0.0,
                }
            by_tool[r.tool]["input_tokens"] += r.input_tokens
            by_tool[r.tool]["output_tokens"] += r.output_tokens
            by_tool[r.tool]["total_tokens"] += r.total_tokens
            by_tool[r.tool]["cost_usd"] += r.estimated_cost_usd

        return {
            "total_input_tokens": total_input,
            "total_output_tokens": total_output,
            "total_tokens": total_tokens,
            "total_cost_usd": total_cost,
            "by_tool": by_tool,
        }


def _parse_int(value: str) -> int:
    """Parse an integer from a string, stripping commas and whitespace."""
    return int(value.replace(",", "").strip())


def _parse_float(value: str) -> float:
    """Parse a float from a string, stripping dollar signs, commas, and whitespace."""
    return float(value.replace("$", "").replace(",", "").strip())


# Flexible regex patterns for Claude Code output parsing
_TOTAL_TOKENS_RE = re.compile(r"[Tt]otal\s+tokens?\s*[:=]\s*([\d,]+)", re.IGNORECASE)
_INPUT_TOKENS_RE = re.compile(r"[Ii]nput\s+tokens?\s*[:=]\s*([\d,]+)", re.IGNORECASE)
_OUTPUT_TOKENS_RE = re.compile(r"[Oo]utput\s+tokens?\s*[:=]\s*([\d,]+)", re.IGNORECASE)
_COST_RE = re.compile(r"[Cc]ost\s*[:=]\s*\$?([\d,]+\.?\d*)", re.IGNORECASE)
_MODEL_RE = re.compile(r"[Mm]odel\s*[:=]\s*(\S+)", re.IGNORECASE)

# Alternative patterns: "Total cost: $X.XX" or "Session cost: $X.XX"
_ALT_COST_RE = re.compile(
    r"(?:total|session|estimated)\s+cost\s*[:=]\s*\$?([\d,]+\.?\d*)", re.IGNORECASE
)

# Token pattern: "Tokens used: X (input: Y, output: Z)"
_TOKENS_USED_RE = re.compile(
    r"[Tt]okens?\s+used\s*[:=]\s*([\d,]+)\s*"
    r"(?:\(\s*input\s*[:=]\s*([\d,]+)\s*,\s*output\s*[:=]\s*([\d,]+)\s*\))?",
    re.IGNORECASE,
)


def parse_claude_code_output(output: str) -> TokenUsage | None:
    """Parse Claude Code's session summary for token/cost data.

    Looks for patterns like:
    - "Total tokens: 12345"
    - "Input tokens: 5000"
    - "Output tokens: 7345"
    - "Cost: $0.45"
    - "Model: claude-3.5-sonnet"
    - "Tokens used: 12345 (input: 5000, output: 7345)"

    Returns a TokenUsage or None if no token data is found.
    """
    if not output:
        return None

    input_tokens = 0
    output_tokens = 0
    total_tokens = 0
    cost = 0.0
    model = "unknown"

    # Try the combined "Tokens used" pattern first
    m = _TOKENS_USED_RE.search(output)
    if m:
        total_tokens = _parse_int(m.group(1))
        if m.group(2) and m.group(3):
            input_tokens = _parse_int(m.group(2))
            output_tokens = _parse_int(m.group(3))

    # Try individual token patterns
    m = _INPUT_TOKENS_RE.search(output)
    if m:
        input_tokens = _parse_int(m.group(1))

    m = _OUTPUT_TOKENS_RE.search(output)
    if m:
        output_tokens = _parse_int(m.group(1))

    m = _TOTAL_TOKENS_RE.search(output)
    if m:
        total_tokens = _parse_int(m.group(1))

    # If we have input + output but no total, calculate it
    if total_tokens == 0 and (input_tokens > 0 or output_tokens > 0):
        total_tokens = input_tokens + output_tokens

    # If no token data at all, bail
    if total_tokens == 0:
        return None

    # Parse cost
    m = _ALT_COST_RE.search(output)
    if m:
        cost = _parse_float(m.group(1))
    else:
        m = _COST_RE.search(output)
        if m:
            cost = _parse_float(m.group(1))

    # Parse model
    m = _MODEL_RE.search(output)
    if m:
        model = m.group(1)

    return TokenUsage(
        tool="claude-code",
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        total_tokens=total_tokens,
        estimated_cost_usd=cost,
        model=model,
    )
