"""Egress gate wrapper with mandatory classification enforcement."""

from __future__ import annotations

import hashlib
import logging
from typing import Any

from cyberbrain26.security_plane.egress import EgressGate, EgressRequest
from cyberbrain26.security_plane.permission import SecurityError

logger = logging.getLogger(__name__)


class ModelEgressGate:
    """Wrapper around EgressGate that enforces mandatory classification."""

    def __init__(
        self,
        *,
        egress_policy: Any | None = None,
        audit_log: Any | None = None,
        egress_gate: EgressGate | None = None,
    ) -> None:
        self.egress_policy = egress_policy
        self.audit_log = audit_log
        self._egress_gate = egress_gate

    def authorize_model_call(
        self,
        *,
        agent_id: str,
        provider_id: str,
        model_id: str,
        payload_sha256: str,
        classifications: frozenset[Any],
    ) -> None:
        """Authorize a model call with mandatory classification.

        Args:
            agent_id: Identifier of the agent making the call.
            provider_id: Identifier of the model provider.
            model_id: Identifier of the model being called.
            payload_sha256: SHA-256 hash of the request payload.
            classifications: Mandatory set of classifications for the call.

        Raises:
            SecurityError: If classifications is empty or None.
        """
        if not classifications:
            raise SecurityError(
                "Model egress requires an explicit classification; refusing to fail open"
            )
        # Log the authorization attempt
        logger.info(
            "Authorizing model call: agent=%s, provider=%s, model=%s, payload_hash=%s, classifications=%s",
            agent_id,
            provider_id,
            model_id,
            payload_sha256,
            classifications,
        )
        if self._egress_gate is None:
            raise SecurityError("egress gate not configured; refusing to fail open")

        # Construct EgressRequest from the provided parameters
        request = EgressRequest(
            agent_id=agent_id,
            provider_id=provider_id,
            payload_sha256=payload_sha256,
            classifications=classifications,
            destination="",
        )
        self._egress_gate.authorize(request)

    async def guarded_model_call(
        self,
        coro_factory: Any,
        *,
        agent_id: str,
        provider_id: str,
        model_id: str,
        payload: bytes | str,
        classifications: frozenset[Any],
    ) -> Any:
        """Execute a model call after authorization.

        Args:
            coro_factory: Async callable that performs the actual model inference.
            agent_id: Identifier of the agent making the call.
            provider_id: Identifier of the model provider.
            model_id: Identifier of the model being called.
            payload: The request payload.
            classifications: Mandatory set of classifications for the call.

        Returns:
            The result of the coro_factory call.

        Raises:
            SecurityError: If classifications is empty or None.
            Exception: Propagated from coro_factory if it fails.
        """
        # Validate classification
        if not classifications:
            raise SecurityError("Model egress requires an explicit classification; refusing to fail open")

        # Compute payload digest before anything else; the raw payload is never
        # logged and never leaves this method.
        if isinstance(payload, str):
            payload_bytes = payload.encode("utf-8")
        elif isinstance(payload, bytes):
            payload_bytes = payload
        else:
            raise TypeError("Payload must be bytes or a string")

        self.authorize_model_call(
            agent_id=agent_id,
            provider_id=provider_id,
            model_id=model_id,
            payload_sha256=hashlib.sha256(payload_bytes).hexdigest(),
            classifications=classifications,
        )

        return await coro_factory()