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

from __future__ import annotations

import logging
from typing import Any, FrozenSet, Sequence

from cyberbrain26.security_plane.egress import EgressGate

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,
    ) -> None:
        self.egress_policy = egress_policy
        self.audit_log = audit_log

    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,
        )
        # In a real implementation, this would call the underlying EgressGate
        # and verify the classification meets policy requirements.
        # For now we simulate successful authorization.
        pass

    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 hash if not already provided
        if isinstance(payload, str):
            import hashlib
            payload = payload.encode("utf-8")
        elif not isinstance(payload, bytes):
            raise TypeError("Payload must be bytes or a string")

        # Authorize the call
        self.authorize_model_call(
            agent_id=agent_id,
            provider_id=provider_id,
            model_id=model_id,
            payload_sha256=hashlib.sha256(payload).hexdigest(),
            classifications=classifications,
        )

        # Execute the call
        return await coro_factory()