"""LLM token cost tracking."""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from bilgeai.core.constants import (
    ESTIMATED_TOKENS_CONFIG_OPTIMIZE,
    ESTIMATED_TOKENS_DATASET_EVAL,
    ESTIMATED_TOKENS_SCRIPT_REVIEW,
    MODEL_PRICING,
)

logger = logging.getLogger(__name__)


@dataclass
class TokenUsage:
    """Single LLM call usage record."""

    model: str
    tokens_used: int
    timestamp: str = ""
    estimated_cost_usd: float = 0.0
    purpose: str = ""

    def __post_init__(self):
        if not self.timestamp:
            self.timestamp = datetime.now(tz=timezone.utc).isoformat()


@dataclass
class CostTracker:
    """Tracks LLM token usage and estimated costs across a session."""

    usages: list[TokenUsage] = field(default_factory=list)

    @property
    def total_tokens(self) -> int:
        return sum(u.tokens_used for u in self.usages)

    @property
    def total_cost_usd(self) -> float:
        return sum(u.estimated_cost_usd for u in self.usages)

    def record(
        self,
        model: str,
        tokens_used: int,
        purpose: str = "",
    ) -> TokenUsage:
        """Record a token usage event."""
        cost = estimate_cost(model, tokens_used)
        usage = TokenUsage(
            model=model,
            tokens_used=tokens_used,
            estimated_cost_usd=cost,
            purpose=purpose,
        )
        self.usages.append(usage)
        logger.debug(
            "Token usage: model=%s tokens=%d cost=$%.6f purpose=%s",
            model,
            tokens_used,
            cost,
            purpose,
        )
        return usage

    def summary(self) -> dict:
        """Return a summary of all usage."""
        by_model: dict[str, dict] = {}
        for u in self.usages:
            if u.model not in by_model:
                by_model[u.model] = {"tokens": 0, "cost_usd": 0.0, "calls": 0}
            by_model[u.model]["tokens"] += u.tokens_used
            by_model[u.model]["cost_usd"] += u.estimated_cost_usd
            by_model[u.model]["calls"] += 1

        return {
            "total_tokens": self.total_tokens,
            "total_cost_usd": round(self.total_cost_usd, 6),
            "total_calls": len(self.usages),
            "by_model": by_model,
        }

    def save(self, path: Path) -> None:
        """Save usage log to JSON file."""
        data = {
            "usages": [
                {
                    "model": u.model,
                    "tokens_used": u.tokens_used,
                    "estimated_cost_usd": u.estimated_cost_usd,
                    "timestamp": u.timestamp,
                    "purpose": u.purpose,
                }
                for u in self.usages
            ],
            "summary": self.summary(),
        }
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(data, indent=2), encoding="utf-8")

    @classmethod
    def load(cls, path: Path) -> CostTracker:
        """Load usage log from JSON file."""
        tracker = cls()
        if path.exists():
            data = json.loads(path.read_text(encoding="utf-8"))
            for u in data.get("usages", []):
                tracker.usages.append(
                    TokenUsage(
                        model=u["model"],
                        tokens_used=u["tokens_used"],
                        estimated_cost_usd=u.get("estimated_cost_usd", 0.0),
                        timestamp=u.get("timestamp", ""),
                        purpose=u.get("purpose", ""),
                    )
                )
        return tracker


def estimate_cost(model: str, tokens_used: int) -> float:
    """Estimate cost in USD for a given model and token count.

    Uses approximate average of input/output pricing.
    """
    model_lower = model.lower()
    for prefix, (input_cost, output_cost) in MODEL_PRICING.items():
        if prefix in model_lower:
            avg_cost_per_token = (input_cost + output_cost) / 2 / 1_000_000
            return tokens_used * avg_cost_per_token

    # Unknown model: conservative estimate
    return tokens_used * 5.0 / 1_000_000


def estimate_run_cost(provider: str, model: str) -> float:
    """Estimate total LLM cost for a full bilge run session.

    A bilge run makes 2-3 LLM calls:
    1. Config optimization (~3K tokens)
    2. Script review (~5K tokens)
    3. Optional dataset evaluation (~1.5K tokens)
    """
    if provider == "local":
        return 0.0

    total_tokens = (
        ESTIMATED_TOKENS_CONFIG_OPTIMIZE
        + ESTIMATED_TOKENS_SCRIPT_REVIEW
        + ESTIMATED_TOKENS_DATASET_EVAL
    )
    return estimate_cost(model, total_tokens)
