from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Iterable, Optional

from ember.models.providers.utils import ResponsesRequest

JSON = Any


@dataclass
class Message:
    role: str
    content: str
    tool_calls: list[dict[str, object]] | None = None
    tool_call_id: str | None = None


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


@dataclass
class GatewayRequest:
    org_id: str
    responses: ResponsesRequest
    messages_for_estimate: list[Message] = field(default_factory=list)
    session_id: Optional[str] = None
    idempotency_key: Optional[str] = None
    policy: Optional[Policy] = None

    @property
    def model(self) -> str:
        return self.responses.model

    @property
    def stream(self) -> bool:
        return bool(self.responses.stream)


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


@dataclass
class Choice:
    index: int
    message: Message
    finish_reason: Optional[str] = None


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


@dataclass
class RBC:
    id: str
    org_id: str
    request_id: str
    admissible: list[str]
    region_policy: Optional[str]
    max_prompt_tokens: int
    max_completion_tokens: int
    price_book_version: str
    exp: int

    def claims(self) -> dict[str, JSON]:
        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,
        }


@dataclass
class AdmissionDecision:
    allow: bool
    admissible_providers: list[str]
    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: Optional[str] = None


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


@dataclass
class StreamingResponse:
    id: str
    org_id: str
    model: str
    provider: str
    provider_model: str
    router: RouterInfo
    admission: AdmissionDecision
    rbc: RBC
    rbc_token: str
    stream: Iterable[str]
    created: int
    fbr_id: str
    billing_mode: str
    price_book_version: str
    plan: dict[str, JSON]
    idempotency_key: Optional[str]
    api_key_id: Optional[str]
    prompt_estimate: int
    completion_budget: int
    provider_metadata: dict[str, JSON] = field(default_factory=dict)
    stream_source: Iterable[str] | None = None
    responses_request: ResponsesRequest | None = None
    api_key_override: Optional[str] = None
