"""
Billing - Token usage extraction and per-request reporting via heartbeat.

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

Each request is reported individually (no aggregation) to preserve request_id
for downstream audit and reconciliation.
"""

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


@dataclass
class UsageRecord:
    """Token usage for a single request."""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    reasoning_tokens: int = 0
    request_id: str = ""

    def to_dict(self, org_id: str, model: str, timestamp: int) -> Dict[str, Any]:
        """Serialize for heartbeat payload (field names compatible with proxy-server UsageData)."""
        d: Dict[str, Any] = {
            "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": 1,
            "timestamp": timestamp,
        }
        if self.request_id:
            d["request_id"] = self.request_id
        return d


@dataclass
class _PendingItem:
    """Internal queue item: usage record with its billing context."""
    org_id: str
    model: str
    usage: UsageRecord
    timestamp: int = field(default_factory=lambda: int(time.time()))


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, request_id: str = "") -> 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_id=request_id,
    )


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.

        The response-level 'id' field (e.g. 'chatcmpl-xxx') is the canonical
        request_id assigned by the inference engine (vLLM/SGLang). It uniquely
        identifies the inference job and appears in the client-facing response,
        making it suitable for billing reconciliation.
        """
        try:
            data = json.loads(body)
            request_id = data.get("id", "")
            return _parse_usage(data.get("usage"), request_id=request_id)
        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).

        The chunk-level 'id' field is set by the inference engine and is
        consistent across all chunks of the same request.
        """
        if not line.startswith("data: "):
            return None
        payload = line[6:].strip()
        if payload == "[DONE]":
            return None
        try:
            data = json.loads(payload)
            request_id = data.get("id", "")
            return _parse_usage(data.get("usage"), request_id=request_id)
        except (json.JSONDecodeError, AttributeError, TypeError):
            return None


class BillingTracker:
    """
    Thread-safe per-request usage tracker.

    record() enqueues each request individually. flush() drains the queue
    and returns all pending records — called by heartbeat every 2s.
    """

    def __init__(self):
        self._lock = threading.Lock()
        self._queue: List[_PendingItem] = []
        self._total_records = 0
        self._total_flushes = 0

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

    def flush(self) -> List[Dict[str, Any]]:
        """Drain and return all pending per-request records, then clear."""
        with self._lock:
            if not self._queue:
                return []
            items = self._queue
            self._queue = []
            self._total_flushes += 1
        return [item.usage.to_dict(item.org_id, item.model, item.timestamp) for item in items]

    @property
    def pending_count(self) -> int:
        with self._lock:
            return len(self._queue)

    @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
