"""Cost attribution (架构文档 §13, FR-BILL-08, NFR-COST-01).

**每一次付费调用都必须落 ``ob_ai_call``，并可归因到用户与功能。**
V1 不收费让这张表的地位更高，不是更低 —— 它是 V2 定价的唯一数据来源，
不允许「先上线后补埋点」。

Rates are configuration, never constants at a call site: providers re-price
without warning and V1's whole cost model depends on these numbers being
correct at the time of the call.
"""

from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass, replace
from datetime import UTC, datetime
from decimal import ROUND_HALF_UP, Decimal
from typing import Protocol, runtime_checkable

from packages.llm.types import Usage

_MILLION = Decimal("1000000")
_QUANT = Decimal("0.000001")


@dataclass(frozen=True, slots=True)
class ModelRate:
    """Cents per 1M tokens. Cached input is billed at a lower rate by DeepSeek."""

    input_cents_per_mtok: Decimal
    cached_input_cents_per_mtok: Decimal
    output_cents_per_mtok: Decimal

    @classmethod
    def from_mapping(cls, data: Mapping[str, Decimal | str | int | float]) -> ModelRate:
        return cls(
            input_cents_per_mtok=Decimal(str(data["input_cents_per_mtok"])),
            cached_input_cents_per_mtok=Decimal(str(data["cached_input_cents_per_mtok"])),
            output_cents_per_mtok=Decimal(str(data["output_cents_per_mtok"])),
        )


class RateTable:
    """model → :class:`ModelRate`, with an explicit fallback so an unknown model
    is still costed (a zero-cost blind spot is worse than an approximate one)."""

    def __init__(self, rates: dict[str, ModelRate], default: ModelRate) -> None:
        self._rates = dict(rates)
        self._default = default

    def for_model(self, model: str) -> ModelRate:
        return self._rates.get(model, self._default)

    def known_models(self) -> tuple[str, ...]:
        return tuple(sorted(self._rates))


def compute_cost_cents(usage: Usage, rate: ModelRate) -> Decimal:
    """Cost of one call, in cents, with cache-hit tokens billed separately."""
    billable_input = Decimal(usage.billable_input) * rate.input_cents_per_mtok
    cached_input = Decimal(usage.tokens_cached) * rate.cached_input_cents_per_mtok
    output = Decimal(usage.tokens_out) * rate.output_cents_per_mtok
    total = (billable_input + cached_input + output) / _MILLION
    return total.quantize(_QUANT, rounding=ROUND_HALF_UP)


@dataclass(frozen=True, slots=True)
class AICallRecord:
    """One row of ``ob_ai_call``. Written for successes *and* failures."""

    call_id: int
    scene: str
    provider: str
    model: str
    success: bool
    usage: Usage
    latency_ms: int
    cost_cents: Decimal
    user_id: int | None = None
    prompt_id: str | None = None
    prompt_version: str | None = None
    error_code: str | None = None
    error_message: str | None = None
    ref_type: str | None = None
    ref_id: str | None = None
    created_at: datetime | None = None

    def with_defaults(self) -> AICallRecord:
        if self.created_at is not None:
            return self
        return replace(self, created_at=datetime.now(UTC))


@runtime_checkable
class AICallRecorder(Protocol):
    """Persistence port. The gateway never touches the ORM directly."""

    async def record(self, record: AICallRecord) -> None: ...


class NullAICallRecorder:
    """Explicitly opts out of cost attribution. Only valid in unit tests."""

    async def record(self, record: AICallRecord) -> None:
        return None


class InMemoryAICallRecorder:
    """Test double that keeps every record for assertions."""

    def __init__(self) -> None:
        self.records: list[AICallRecord] = []

    async def record(self, record: AICallRecord) -> None:
        self.records.append(record)


class CompositeAICallRecorder:
    """Fan out to several recorders (e.g. PostgreSQL + Prometheus).

    A failure in one recorder must not lose the others, and must never fail the
    user's request — cost bookkeeping is mandatory but not user-facing.
    """

    def __init__(self, *recorders: AICallRecorder) -> None:
        self._recorders = recorders

    async def record(self, record: AICallRecord) -> None:
        for recorder in self._recorders:
            await recorder.record(record)