from __future__ import annotations

import hashlib
import json
import logging
import uuid
from typing import Optional

from .models import RBC, AdmissionDecision, Policy, now_s
from .rbc import jws_sign_hs256
from .settings import get_gateway_settings

_LOG = logging.getLogger(__name__)


def _policy_hash(policy: Optional[Policy]) -> str:
    if policy is None:
        return "p0:none"
    # Deterministic hash of relevant fields
    payload = (
        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]


class AdmissionController:
    """Deterministic, minimal admission controller (P0).

    - Filters providers against policy and a provided candidate set
    - Computes enforced params (currently pass-through)
    - Issues an RBC (short-lived JWS) for budget enforcement
    - Emits a Signed Decision Record (SDR) as a JWS over the decision
    """

    def __init__(self, price_book_version: str | None = None):
        # Compute a deterministic price book version if not provided
        if price_book_version is None:
            try:
                from ember.models.pricing import Pricing

                pb = Pricing().get_all_pricing()
                blob = json.dumps(pb, sort_keys=True)
                self.price_book_version = (
                    "pb:" + hashlib.sha256(blob.encode("utf-8")).hexdigest()[:16]
                )
            except Exception as exc:  # pragma: no cover
                _LOG.exception("Failed to load pricing information for admission controller")
                raise RuntimeError("Unable to load pricing data for admission controller") from exc
        else:
            self.price_book_version = price_book_version

    def admit(
        self,
        org_id: str,
        request_id: str,
        model: str,
        candidate_providers: list[str],
        policy: Optional[Policy],
        max_prompt_tokens: int,
        max_completion_tokens: int,
        ttl_seconds: int = 120,
    ) -> tuple[AdmissionDecision, RBC, str]:
        # Filter providers by policy constraints
        admissible = list(candidate_providers)
        enforced: dict[str, float] = {}
        lossiness: list[str] = []
        deny_reasons: list[str] = []

        if policy and policy.allowed_providers:
            admissible = [p for p in admissible if p in set(policy.allowed_providers)]
        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:
            # Optional: capability gating (flagged) will prune providers that lack strict JSON
            # Otherwise we declare lossiness and allow router to pick, leading to 409 later.
            gating = get_gateway_settings().features.capability_gating
            if gating:
                try:
                    from .capabilities import list_capabilities

                    caps = list_capabilities()
                    admissible = [p for p in admissible if p in caps and caps[p]["json_strict"]]
                except Exception as exc:
                    _LOG.exception("Capability registry unavailable during admission decision")
                    raise RuntimeError("Capability registry unavailable during admission") from exc
            else:
                # P0: we don't implement strict JSON uniformly; declare lossiness if requested
                lossiness.append("json_strict_not_supported_p0")
        if policy and policy.max_cost_usd is not None:
            # P0: cannot guarantee exact cost cap pre-run
            enforced["max_cost_usd_hint"] = 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,
        )

        # P0 Signed Decision Record (SDR)
        from ember.types import JSONValue

        sdr_claims: dict[str, JSONValue] = {
            "type": "SDR_P0",
            "org_id": org_id,
            "request_id": request_id,
            "policy_hash": _policy_hash(policy),
            "admissible": admissible,
            "enforced": enforced,
            "lossiness": lossiness,
        }
        decision.sdr_jws = jws_sign_hs256(sdr_claims)

        # RBC issuance
        rbc = RBC(
            id=f"rbc_{uuid.uuid4()}",
            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=self.price_book_version,
            exp=now_s() + max(1, ttl_seconds),
        )
        rbc_jws = jws_sign_hs256(rbc.as_claims())

        return decision, rbc, rbc_jws
