from __future__ import annotations

import hashlib
from dataclasses import dataclass
from typing import Iterable, Mapping, Optional, Protocol

from .model import RBC, AdmissionDecision, Policy


@dataclass
class AdmissionResult:
    decision: AdmissionDecision
    rbc: RBC
    token: str


class RbcSigner(Protocol):
    def sign(self, claims: Mapping[str, object]) -> str: ...


class AdmissionController:
    """Policy enforcement, decision signing, and RBC issuance."""

    def __init__(self, signer: RbcSigner, *, capability_gating: bool = False) -> None:
        self._signer = signer
        self._capability_gating = capability_gating

    def admit(
        self,
        *,
        org_id: str,
        request_id: str,
        model: str,
        candidate_providers: Iterable[str],
        policy: Policy | None,
        max_prompt_tokens: int,
        max_completion_tokens: int,
        ttl_seconds: int,
        price_book_version: str = "p0-0",
        capabilities: Mapping[str, Mapping[str, bool]] | None = None,
        now_s: Optional[int] = None,
    ) -> AdmissionResult:
        admissible = list(candidate_providers)
        enforced: dict[str, float] = {}
        lossiness: list[str] = []
        deny_reasons: list[str] = []

        if policy and policy.allowed_providers:
            allowed = set(policy.allowed_providers)
            admissible = [p for p in admissible if p in allowed]

        if policy and policy.allowed_models and model not in set(policy.allowed_models):
            deny_reasons.append("model_not_allowed")

        if policy and policy.json_strict:
            if self._capability_gating:
                provider_caps = capabilities or {}
                admissible = [
                    provider
                    for provider in admissible
                    if provider_caps.get(provider, {}).get("json_strict", False)
                ]
            else:
                lossiness.append("json_strict_not_supported")

        if policy and policy.max_cost_usd is not None:
            enforced["max_cost_usd_hint"] = float(policy.max_cost_usd)

        decision = AdmissionDecision(
            allow=bool(admissible) and not deny_reasons,
            admissible_providers=admissible,
            enforced_params=enforced,
            lossiness=lossiness,
            deny_reasons=deny_reasons,
        )

        exp = (now_s or 0) + max(1, ttl_seconds) if now_s is not None else ttl_seconds
        rbc = RBC(
            id=f"rbc_{request_id}",
            org_id=org_id,
            request_id=request_id,
            admissible=admissible,
            region_policy=policy.region if policy else None,
            max_prompt_tokens=max_prompt_tokens,
            max_completion_tokens=max_completion_tokens,
            price_book_version=price_book_version,
            exp=exp,
        )

        claims = rbc.claims()
        token = self._signer.sign(claims)

        decision.sdr_jws = self._signer.sign(
            {
                "type": "SDR_P0",
                "org_id": org_id,
                "request_id": request_id,
                "policy_hash": _policy_hash(policy),
                "admissible": admissible,
                "enforced": enforced,
                "lossiness": lossiness,
                "deny_reasons": deny_reasons,
            }
        )

        return AdmissionResult(decision=decision, rbc=rbc, token=token)


def _policy_hash(policy: Policy | None) -> str:
    if policy is None:
        return "p0:none"
    payload = "|".join(
        [
            str(sorted(policy.allowed_providers or [])),
            str(sorted(policy.allowed_models or [])),
            str(policy.region),
            str(policy.max_cost_usd),
            str(policy.json_strict),
            str(policy.retention_ttl),
            str(policy.training_opt_out),
        ]
    )
    return "p0:" + hashlib.sha256(payload.encode("utf-8")).hexdigest()[:16]
