from __future__ import annotations

import time
import uuid
from dataclasses import dataclass, field
from typing import Any

from ember.types import JSONValue

# Canonical request/response envelopes (minimal P0)


@dataclass
class Message:
    """A message in a chat completion request."""

    role: str
    content: str | None
    tool_calls: list[dict[str, Any]] | None = None
    tool_call_id: str | None = None


@dataclass
class Policy:
    allowed_providers: list[str] | None = None
    allowed_models: list[str] | None = None
    region: str | None = None
    max_cost_usd: float | None = None
    json_strict: bool | None = None
    retention_ttl: int | None = None
    training_opt_out: bool | None = None


@dataclass
class ChatCompletionRequest:
    model: str
    messages: list[Message]
    stream: bool = False
    session_id: str | None = None
    idempotency_key: str | None = None
    policy: Policy | None = None
    extra_params: dict[str, Any] = field(default_factory=dict)


@dataclass
class Usage:
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    cost_usd: float | None = None


@dataclass
class Choice:
    index: int
    message: Message
    finish_reason: str | None = None


@dataclass
class RouterInfo:
    session_affinity: bool = False
    score_components: dict[str, float] = field(default_factory=dict)


@dataclass
class ChatCompletionResponse:
    id: str
    model: str
    provider: str
    region: str | None
    created: int
    choices: list[Choice]
    usage: Usage
    billing_record_id: str
    router: RouterInfo
    raw: dict[str, JSONValue] = field(default_factory=dict)


@dataclass
class AdmissionDecision:
    allow: bool
    admissible_providers: list[str] = field(default_factory=list)
    enforced_params: dict[str, float] = field(default_factory=dict)
    lossiness: list[str] = field(default_factory=list)
    deny_reasons: list[str] = field(default_factory=list)
    sdr_jws: str | None = None  # signed decision record


@dataclass
class RBC:
    """Request Budget Certificate (RBC).

    An RBC is a short-lived, signed token that authorizes a specific request
    to proceed with pre-approved resource limits. It encapsulates admission
    control decisions and budget constraints.

    Attributes:
        id: Unique identifier for this RBC
        org_id: Organization that owns this request
        request_id: The request this RBC authorizes
        admissible: List of provider names allowed for this request
        region_policy: Optional region constraint
        max_prompt_tokens: Maximum input tokens allowed
        max_completion_tokens: Maximum output tokens allowed
        price_book_version: Version of pricing data used
        exp: Expiration time (Unix epoch seconds)
    """

    id: str
    org_id: str
    request_id: str
    admissible: list[str]
    region_policy: str | None
    max_prompt_tokens: int
    max_completion_tokens: int
    price_book_version: str
    exp: int

    def as_claims(self) -> dict[str, JSONValue]:
        return {
            "rbc_id": self.id,
            "org_id": self.org_id,
            "request_id": self.request_id,
            "admissible": self.admissible,
            "region_policy": self.region_policy,
            "max_prompt_tokens": self.max_prompt_tokens,
            "max_completion_tokens": self.max_completion_tokens,
            "price_book_version": self.price_book_version,
            "exp": self.exp,
        }


def new_request_id() -> str:
    return f"req_{uuid.uuid4()}"


def now_s() -> int:
    return int(time.time())
