"""
cortex/llm/usage_tracker.py — LLM Usage Tracking (Issue #419).
"""
from __future__ import annotations

import math
import threading
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional


@dataclass
class _CallRecord:
    tier: str
    operation: str
    input_tokens: int
    output_tokens: int
    latency_ms: float
    error: bool
    timestamp: float = field(default_factory=time.time)


@dataclass
class TierStats:
    call_count: int = 0
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    error_count: int = 0
    _latencies: List[float] = field(default_factory=list)

    def avg_latency_ms(self) -> float:
        if not self._latencies:
            return 0.0
        return sum(self._latencies) / len(self._latencies)

    def p50_latency_ms(self) -> float:
        return _percentile(sorted(self._latencies), 0.50)

    def p95_latency_ms(self) -> float:
        return _percentile(sorted(self._latencies), 0.95)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "call_count": self.call_count,
            "total_input_tokens": self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
            "avg_latency_ms": round(self.avg_latency_ms(), 3),
            "p50_latency_ms": round(self.p50_latency_ms(), 3),
            "p95_latency_ms": round(self.p95_latency_ms(), 3),
            "error_count": self.error_count,
        }


class LLMUsageTracker:
    """Thread-safe, singleton-per-owner LLM call tracker."""
    _instances: Dict[str, "LLMUsageTracker"] = {}
    _instances_lock: threading.Lock = threading.Lock()
    TIERS = ("fast", "main", "reasoning")

    @classmethod
    def get(cls, owner_id: str = "default") -> "LLMUsageTracker":
        with cls._instances_lock:
            if owner_id not in cls._instances:
                cls._instances[owner_id] = cls(owner_id=owner_id)
            return cls._instances[owner_id]

    @classmethod
    def reset_instances(cls) -> None:
        with cls._instances_lock:
            cls._instances.clear()

    def __init__(self, owner_id: str = "default") -> None:
        self.owner_id = owner_id
        self._lock = threading.Lock()
        self._records: List[_CallRecord] = []
        self._tier_stats: Dict[str, TierStats] = {t: TierStats() for t in self.TIERS}
        self._operation_stats: Dict[str, Dict[str, TierStats]] = {}
        self._created_at = time.time()

    def record_call(self, tier: str, operation: str, input_tokens: int = 0,
                    output_tokens: int = 0, latency_ms: float = 0.0, error: bool = False) -> None:
        record = _CallRecord(tier=tier, operation=operation, input_tokens=input_tokens,
                             output_tokens=output_tokens, latency_ms=latency_ms, error=error)
        with self._lock:
            self._records.append(record)
            if tier not in self._tier_stats:
                self._tier_stats[tier] = TierStats()
            ts = self._tier_stats[tier]
            ts.call_count += 1; ts.total_input_tokens += input_tokens
            ts.total_output_tokens += output_tokens; ts._latencies.append(latency_ms)
            if error: ts.error_count += 1
            if operation not in self._operation_stats:
                self._operation_stats[operation] = {}
            if tier not in self._operation_stats[operation]:
                self._operation_stats[operation][tier] = TierStats()
            os_ = self._operation_stats[operation][tier]
            os_.call_count += 1; os_.total_input_tokens += input_tokens
            os_.total_output_tokens += output_tokens; os_._latencies.append(latency_ms)
            if error: os_.error_count += 1

    def get_summary(self, since: Optional[datetime] = None, tier_filter: Optional[str] = None) -> Dict[str, Any]:
        since_ts: Optional[float] = None
        if since is not None:
            since_ts = since.timestamp()
        if since_ts is None and tier_filter is None:
            with self._lock:
                tier_data = {t: s.to_dict() for t, s in self._tier_stats.items()}
                op_data: Dict[str, Any] = {op: {t: s.to_dict() for t, s in tm.items()}
                                            for op, tm in self._operation_stats.items()}
            return self._build_response(tier_data, op_data)
        with self._lock:
            records = list(self._records)
        if since_ts is not None:
            records = [r for r in records if r.timestamp >= since_ts]
        if tier_filter is not None:
            records = [r for r in records if r.tier == tier_filter]
        return self._build_response(self._agg_tiers(records), self._agg_ops(records))

    def _agg_tiers(self, records: List[_CallRecord]) -> Dict[str, Any]:
        temp: Dict[str, TierStats] = {}
        for r in records:
            if r.tier not in temp: temp[r.tier] = TierStats()
            ts = temp[r.tier]; ts.call_count += 1
            ts.total_input_tokens += r.input_tokens; ts.total_output_tokens += r.output_tokens
            ts._latencies.append(r.latency_ms)
            if r.error: ts.error_count += 1
        result = {t: temp.get(t, TierStats()).to_dict() for t in self.TIERS}
        for t, s in temp.items():
            if t not in result: result[t] = s.to_dict()
        return result

    def _agg_ops(self, records: List[_CallRecord]) -> Dict[str, Any]:
        temp: Dict[str, Dict[str, TierStats]] = {}
        for r in records:
            if r.operation not in temp: temp[r.operation] = {}
            if r.tier not in temp[r.operation]: temp[r.operation][r.tier] = TierStats()
            os_ = temp[r.operation][r.tier]; os_.call_count += 1
            os_.total_input_tokens += r.input_tokens; os_.total_output_tokens += r.output_tokens
            os_._latencies.append(r.latency_ms)
            if r.error: os_.error_count += 1
        return {op: {t: s.to_dict() for t, s in tm.items()} for op, tm in temp.items()}

    def _build_response(self, tier_data: Dict[str, Any], op_data: Dict[str, Any]) -> Dict[str, Any]:
        return {
            "tiers": tier_data,
            "by_operation": op_data,
            "total": {
                "call_count": sum(t.get("call_count", 0) for t in tier_data.values()),
                "total_input_tokens": sum(t.get("total_input_tokens", 0) for t in tier_data.values()),
                "total_output_tokens": sum(t.get("total_output_tokens", 0) for t in tier_data.values()),
                "error_count": sum(t.get("error_count", 0) for t in tier_data.values()),
            },
            "generated_at": datetime.now(timezone.utc).isoformat(),
        }

    def reset(self) -> None:
        with self._lock:
            self._records.clear()
            self._tier_stats = {t: TierStats() for t in self.TIERS}
            self._operation_stats.clear()
            self._created_at = time.time()


def _percentile(sorted_values: List[float], pct: float) -> float:
    if not sorted_values: return 0.0
    n = len(sorted_values)
    idx = max(0, min(n - 1, int(math.ceil(pct * n)) - 1))
    return round(sorted_values[idx], 3)
