# retoor <retoor@molodetz.nl>

from __future__ import annotations

import contextvars
import logging
import os
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Optional

logger = logging.getLogger("devii.cost")

PRICE_MODEL = os.environ.get("DEVII_PRICE_MODEL", "deepseek-v4-flash")
PRICE_CACHE_HIT_PER_M = float(os.environ.get("DEVII_PRICE_CACHE_HIT", "0.0028"))
PRICE_CACHE_MISS_PER_M = float(os.environ.get("DEVII_PRICE_CACHE_MISS", "0.14"))
PRICE_OUTPUT_PER_M = float(os.environ.get("DEVII_PRICE_OUTPUT", "0.28"))
PER_MILLION = 1_000_000


@dataclass(frozen=True)
class Pricing:
    model: str = PRICE_MODEL
    cache_hit_per_m: float = PRICE_CACHE_HIT_PER_M
    cache_miss_per_m: float = PRICE_CACHE_MISS_PER_M
    output_per_m: float = PRICE_OUTPUT_PER_M


DEFAULT_PRICING = Pricing()

_active_tracker: contextvars.ContextVar = contextvars.ContextVar(
    "devii_cost_tracker", default=None
)


def _now() -> datetime:
    return datetime.now(timezone.utc)


@dataclass
class CostTracker:
    pricing: Pricing = field(default_factory=lambda: DEFAULT_PRICING)
    started_at: datetime = field(default_factory=_now)
    requests: int = 0
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    cache_hit_tokens: int = 0
    cache_miss_tokens: int = 0
    reasoning_tokens: int = 0
    native_cost_usd: float = 0.0
    has_native: bool = False

    def record_native(self, cost_usd: float) -> None:
        self.native_cost_usd += float(cost_usd or 0.0)
        self.has_native = True

    def record(self, usage: Optional[dict[str, Any]]) -> None:
        if not usage:
            return
        prompt = int(usage.get("prompt_tokens", 0) or 0)
        completion = int(usage.get("completion_tokens", 0) or 0)
        total = int(usage.get("total_tokens", prompt + completion) or 0)

        hit = usage.get("prompt_cache_hit_tokens")
        if hit is None:
            details = usage.get("prompt_tokens_details") or {}
            hit = details.get("cached_tokens", 0)
        hit = int(hit or 0)

        miss = usage.get("prompt_cache_miss_tokens")
        if miss is None:
            miss = max(prompt - hit, 0)
        miss = int(miss or 0)

        completion_details = usage.get("completion_tokens_details") or {}
        reasoning = int(completion_details.get("reasoning_tokens", 0) or 0)

        self.requests += 1
        self.prompt_tokens += prompt
        self.completion_tokens += completion
        self.total_tokens += total
        self.cache_hit_tokens += hit
        self.cache_miss_tokens += miss
        self.reasoning_tokens += reasoning
        logger.debug(
            "Cost record: +%d prompt (+%d hit/+%d miss), +%d completion",
            prompt,
            hit,
            miss,
            completion,
        )

    def cost_usd(self) -> dict[str, float]:
        cache_hit = self.cache_hit_tokens / PER_MILLION * self.pricing.cache_hit_per_m
        cache_miss = (
            self.cache_miss_tokens / PER_MILLION * self.pricing.cache_miss_per_m
        )
        output = self.completion_tokens / PER_MILLION * self.pricing.output_per_m
        total = (
            self.native_cost_usd
            if self.has_native
            else (cache_hit + cache_miss + output)
        )
        return {
            "cache_hit_input": round(cache_hit, 8),
            "cache_miss_input": round(cache_miss, 8),
            "output": round(output, 8),
            "total": round(total, 8),
        }

    def stats(self) -> dict[str, Any]:
        elapsed = max((_now() - self.started_at).total_seconds(), 0.0)
        costs = self.cost_usd()
        requests = self.requests or 1
        cache_hit_rate = (
            (self.cache_hit_tokens / self.prompt_tokens) if self.prompt_tokens else 0.0
        )
        per_minute = (self.requests / elapsed * 60.0) if elapsed > 0 else 0.0
        return {
            "model": self.pricing.model,
            "pricing_usd_per_million": {
                "cache_hit_input": self.pricing.cache_hit_per_m,
                "cache_miss_input": self.pricing.cache_miss_per_m,
                "output": self.pricing.output_per_m,
            },
            "requests": self.requests,
            "tokens": {
                "prompt": self.prompt_tokens,
                "completion": self.completion_tokens,
                "total": self.total_tokens,
                "cache_hit": self.cache_hit_tokens,
                "cache_miss": self.cache_miss_tokens,
                "reasoning": self.reasoning_tokens,
            },
            "cache_hit_rate": round(cache_hit_rate, 4),
            "cost_usd": costs,
            "cost_usd_total": costs["total"],
            "cost_cents_total": round(costs["total"] * 100, 6),
            "averages": {
                "tokens_per_request": round(self.total_tokens / requests, 2),
                "prompt_tokens_per_request": round(self.prompt_tokens / requests, 2),
                "completion_tokens_per_request": round(
                    self.completion_tokens / requests, 2
                ),
                "cost_usd_per_request": round(costs["total"] / requests, 8),
            },
            "session": {
                "started_at": self.started_at.strftime("%Y-%m-%dT%H:%M:%SZ"),
                "elapsed_seconds": round(elapsed, 1),
                "requests_per_minute": round(per_minute, 2),
            },
            "note": (
                "Priced as DeepSeek V4 Flash (the cheapest official DeepSeek model). "
                "Token counts are reported by the model endpoint."
            ),
        }


def get_tracker() -> Optional[CostTracker]:
    return _active_tracker.get()


def set_tracker(tracker: CostTracker) -> contextvars.Token:
    return _active_tracker.set(tracker)


def reset_tracker(token: contextvars.Token) -> None:
    _active_tracker.reset(token)


def record_usage(usage: Optional[dict[str, Any]]) -> None:
    tracker = _active_tracker.get()
    if tracker is not None:
        tracker.record(usage)


def record_cost(cost_usd: float) -> None:
    tracker = _active_tracker.get()
    if tracker is not None:
        tracker.record_native(cost_usd)
