"""
Billing - Token usage extraction and incremental reporting via heartbeat.

Flow: Worker response → UsageExtractor → BillingTracker.record()
     → Register heartbeat flush() → Gateway → VM/Redis → Billing service
"""

import json
import time
import threading
from collections import defaultdict
from dataclasses import dataclass
from typing import Optional, Dict, List, Any


@dataclass
class UsageRecord:
    """Token usage for a single request or accumulated period."""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    reasoning_tokens: int = 0
    request_count: int = 0

    def accumulate(self, other: "UsageRecord"):
        """Add another record's values into this one."""
        self.prompt_tokens += other.prompt_tokens
        self.completion_tokens += other.completion_tokens
        self.total_tokens += other.total_tokens
        self.reasoning_tokens += other.reasoning_tokens
        self.request_count += other.request_count

    def to_dict(self, org_id: str, model: str, flush_ts: int) -> Dict[str, Any]:
        """Serialize for heartbeat payload (field names compatible with proxy-server UsageData)."""
        return {
            "org_id": org_id,
            "model_name": model,
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
            "total_tokens": self.total_tokens,
            "reasoning_tokens": self.reasoning_tokens,
            "request_count": self.request_count,
            "timestamp": flush_ts,
        }


def _safe_int(val: Any) -> int:
    """Coerce to non-negative int. Returns 0 for invalid/negative values."""
    try:
        n = int(val)
        return n if n > 0 else 0
    except (TypeError, ValueError, OverflowError):
        return 0


def _parse_usage(usage: Any) -> Optional[UsageRecord]:
    """Parse a usage dict into UsageRecord. Returns None if invalid or zero."""
    if not isinstance(usage, dict):
        return None
    prompt = _safe_int(usage.get("prompt_tokens", 0))
    completion = _safe_int(usage.get("completion_tokens", 0))
    total = _safe_int(usage.get("total_tokens", 0)) or (prompt + completion)
    if total == 0:
        return None
    return UsageRecord(
        prompt_tokens=prompt, completion_tokens=completion,
        total_tokens=total, reasoning_tokens=_safe_int(usage.get("reasoning_tokens", 0)),
        request_count=1,
    )


class UsageExtractor:
    """Extracts token usage from OpenAI-compatible responses (vLLM/SGLang)."""

    @staticmethod
    def extract_from_json(body: bytes) -> Optional[UsageRecord]:
        """Extract from non-streaming JSON response body."""
        try:
            return _parse_usage(json.loads(body).get("usage"))
        except (json.JSONDecodeError, AttributeError, TypeError):
            return None

    @staticmethod
    def extract_from_sse_line(line: str) -> Optional[UsageRecord]:
        """Extract from a single SSE 'data: ...' line (usually the last chunk)."""
        if not line.startswith("data: "):
            return None
        payload = line[6:].strip()
        if payload == "[DONE]":
            return None
        try:
            return _parse_usage(json.loads(payload).get("usage"))
        except (json.JSONDecodeError, AttributeError, TypeError):
            return None


class BillingTracker:
    """
    Thread-safe incremental usage tracker.

    record() accumulates by (org_id, model). flush() returns the delta
    since the last flush and clears the buffer — called by heartbeat every 2s.
    """

    def __init__(self):
        self._lock = threading.Lock()
        self._pending: Dict[tuple, UsageRecord] = defaultdict(UsageRecord)
        self._total_records = 0
        self._total_flushes = 0

    def record(self, org_id: str, model: str, usage: Optional[UsageRecord]):
        """Record one request's usage (thread-safe). Skips None / zero-token records."""
        if not usage or usage.total_tokens == 0:
            return
        key = (org_id or "unknown", model or "unknown")
        with self._lock:
            self._pending[key].accumulate(usage)
            self._total_records += 1

    def flush(self) -> List[Dict[str, Any]]:
        """Return incremental data since last flush, then clear."""
        with self._lock:
            if not self._pending:
                return []
            now = int(time.time())
            data = [acc.to_dict(k[0], k[1], now) for k, acc in self._pending.items()]
            self._pending = defaultdict(UsageRecord)
            self._total_flushes += 1
            return data

    @property
    def pending_dimensions(self) -> int:
        with self._lock:
            return len(self._pending)

    @property
    def total_records(self) -> int:
        return self._total_records

    @property
    def total_flushes(self) -> int:
        return self._total_flushes


# ── Global singleton ──

_billing_tracker: Optional[BillingTracker] = None
_billing_lock = threading.Lock()


def get_billing_tracker() -> BillingTracker:
    global _billing_tracker
    if _billing_tracker is None:
        with _billing_lock:
            if _billing_tracker is None:
                _billing_tracker = BillingTracker()
    return _billing_tracker


def reset_billing_tracker():
    """Reset singleton (testing only)."""
    global _billing_tracker
    with _billing_lock:
        _billing_tracker = None
