"""
Cost Tracker — Real token accounting
======================================
Intercepts every API response, reads actual usage fields from the response body,
accumulates totals, and converts to USD using current published pricing.

Groq pricing (as of mid-2025, per 1M tokens):
  llama-3.3-70b-versatile   input $0.59  output $0.79
  llama-3.1-70b-versatile   input $0.59  output $0.79
  mixtral-8x7b-32768        input $0.24  output $0.24
  llama-3.1-8b-instant      input $0.05  output $0.08
  gemma2-9b-it              input $0.20  output $0.20

NVIDIA NIM pricing (per 1M tokens):
  meta/codellama-70b-instruct           input $0.99  output $0.99
  mistralai/codestral-22b-v0-1          input $0.60  output $0.60
  nvidia/llama-3.1-nemotron-70b-instruct input $4.20 output $4.20
  meta/llama-3.1-70b-instruct           input $0.99  output $0.99
"""
from __future__ import annotations

import json
import threading
from dataclasses import dataclass, field
from typing import Optional


# ── Pricing tables ────────────────────────────────────────────────────────────
# All costs in USD per 1,000,000 tokens

GROQ_PRICING: dict[str, tuple[float, float]] = {
    "llama-3.3-70b-versatile":   (0.59, 0.79),
    "llama-3.1-70b-versatile":   (0.59, 0.79),
    "mixtral-8x7b-32768":        (0.24, 0.24),
    "llama-3.1-8b-instant":      (0.05, 0.08),
    "gemma2-9b-it":              (0.20, 0.20),
}

NVIDIA_PRICING: dict[str, tuple[float, float]] = {
    "meta/codellama-70b-instruct":             (0.99, 0.99),
    "mistralai/codestral-22b-v0-1":            (0.60, 0.60),
    "nvidia/llama-3.1-nemotron-70b-instruct":  (4.20, 4.20),
    "meta/llama-3.1-70b-instruct":             (0.99, 0.99),
}

DEFAULT_GROQ_PRICE  = (0.59, 0.79)   # fallback for unknown groq models
DEFAULT_NVIDIA_PRICE = (0.99, 0.99)  # fallback for unknown nvidia models


@dataclass
class ModelUsage:
    provider: str          # 'groq' | 'nvidia' | 'ollama'
    model: str
    calls: int = 0
    prompt_tokens: int = 0
    completion_tokens: int = 0
    cost_usd: float = 0.0

    @property
    def total_tokens(self) -> int:
        return self.prompt_tokens + self.completion_tokens

    def to_dict(self) -> dict:
        return {
            "provider": self.provider,
            "model": self.model,
            "calls": self.calls,
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
            "total_tokens": self.total_tokens,
            "cost_usd": round(self.cost_usd, 6),
        }


@dataclass
class CostReport:
    models: dict[str, ModelUsage] = field(default_factory=dict)
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)

    def record(
        self,
        provider: str,
        model: str,
        prompt_tokens: int,
        completion_tokens: int,
    ) -> float:
        """
        Record one API call. Returns the incremental cost for this call in USD.
        Thread-safe.
        """
        incremental = _calc_cost(provider, model, prompt_tokens, completion_tokens)

        with self._lock:
            key = f"{provider}/{model}"
            if key not in self.models:
                self.models[key] = ModelUsage(provider=provider, model=model)
            m = self.models[key]
            m.calls += 1
            m.prompt_tokens += prompt_tokens
            m.completion_tokens += completion_tokens
            m.cost_usd += incremental

        return incremental

    @property
    def total_prompt_tokens(self) -> int:
        return sum(m.prompt_tokens for m in self.models.values())

    @property
    def total_completion_tokens(self) -> int:
        return sum(m.completion_tokens for m in self.models.values())

    @property
    def total_tokens(self) -> int:
        return self.total_prompt_tokens + self.total_completion_tokens

    @property
    def total_cost_usd(self) -> float:
        return sum(m.cost_usd for m in self.models.values())

    @property
    def total_calls(self) -> int:
        return sum(m.calls for m in self.models.values())

    def summary_lines(self) -> list[str]:
        lines = []
        for m in sorted(self.models.values(), key=lambda x: -x.cost_usd):
            lines.append(
                f"  {m.provider.upper():<8} {m.model:<45} "
                f"{m.calls:>3} calls  "
                f"{m.prompt_tokens:>7,}p + {m.completion_tokens:>7,}c tokens  "
                f"${m.cost_usd:.4f}"
            )
        lines.append(
            f"\n  {'TOTAL':<54} "
            f"{self.total_calls:>3} calls  "
            f"{self.total_prompt_tokens:>7,}p + {self.total_completion_tokens:>7,}c tokens  "
            f"${self.total_cost_usd:.4f}"
        )
        return lines

    def to_dict(self) -> dict:
        return {
            "total_calls": self.total_calls,
            "total_prompt_tokens": self.total_prompt_tokens,
            "total_completion_tokens": self.total_completion_tokens,
            "total_tokens": self.total_tokens,
            "total_cost_usd": round(self.total_cost_usd, 6),
            "models": {k: v.to_dict() for k, v in self.models.items()},
        }

    def save(self, path) -> None:
        import pathlib
        pathlib.Path(path).write_text(json.dumps(self.to_dict(), indent=2))


def _calc_cost(provider: str, model: str, prompt: int, completion: int) -> float:
    if provider == "groq":
        in_price, out_price = GROQ_PRICING.get(model, DEFAULT_GROQ_PRICE)
    elif provider == "nvidia":
        in_price, out_price = NVIDIA_PRICING.get(model, DEFAULT_NVIDIA_PRICE)
    else:
        return 0.0  # Ollama / local = free
    return (prompt * in_price + completion * out_price) / 1_000_000


# ── Instrumented API wrappers ────────────────────────────────────────────────
# These wrap the real GroqClient / NvidiaClient methods to intercept
# the `usage` field that every OpenAI-compatible API returns.

def _patch_groq_client(client, report: CostReport) -> None:
    """Monkey-patch GroqClient.complete to record token usage."""
    import httpx

    original_post = client._timeout  # just to verify client object

    _orig_complete = client.complete.__func__  # unbound

    async def _tracked_complete(self_inner, messages, model=None, temperature=0.3, max_tokens=8000):
        model = model or "llama-3.3-70b-versatile"
        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        async with httpx.AsyncClient(timeout=self_inner._timeout) as http:
            resp = await http.post(
                self_inner.BASE_URL, json=payload, headers=self_inner._headers
            )
            if not resp.is_success:
                self_inner._handle_error(resp)
            data = resp.json()
            usage = data.get("usage", {})
            report.record(
                provider="groq",
                model=model,
                prompt_tokens=usage.get("prompt_tokens", 0),
                completion_tokens=usage.get("completion_tokens", 0),
            )
            return data["choices"][0]["message"]["content"]

    import types
    client.complete = types.MethodType(_tracked_complete, client)


def _patch_nvidia_client(client, report: CostReport) -> None:
    """Monkey-patch NvidiaClient.generate_code to record token usage."""
    import httpx

    async def _tracked_generate(self_inner, system_prompt, user_prompt,
                                model="meta/codellama-70b-instruct", max_tokens=4096):
        from forge.core.api import strip_fences
        payload = {
            "model": model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            "temperature": 0.12,
            "max_tokens": max_tokens,
            "stream": False,
        }
        async with httpx.AsyncClient(timeout=self_inner._timeout) as http:
            resp = await http.post(
                self_inner.BASE_URL, json=payload, headers=self_inner._headers
            )
            if not resp.is_success:
                self_inner._handle_error(resp)
            data = resp.json()
            usage = data.get("usage", {})
            report.record(
                provider="nvidia",
                model=model,
                prompt_tokens=usage.get("prompt_tokens", 0),
                completion_tokens=usage.get("completion_tokens", 0),
            )
            raw = data["choices"][0]["message"]["content"]
            return strip_fences(raw)

    import types
    client.generate_code = types.MethodType(_tracked_generate, client)
