"""Risk gates for agent access control.

Provides configurable gates that assess an agent's or team's risk
before allowing an action. Useful for automated trust decisions.
"""

from __future__ import annotations

from typing import Literal, Optional

from pydantic import BaseModel, Field

from mnemom_types import RiskAssessment, RiskContext, RiskLevel, TeamRiskAssessment

from .api import RiskClient

_RISK_LEVEL_ORDINALS: dict[str, int] = {
    "low": 0,
    "medium": 1,
    "high": 2,
    "critical": 3,
}


class RiskGateConfig(BaseModel):
    """Configuration for a risk gate."""

    max_risk_level: Optional[RiskLevel] = Field(
        None, description="Maximum allowed risk level"
    )
    max_risk_score: Optional[float] = Field(
        None, description="Maximum allowed risk score (0-1)"
    )
    base_url: str = Field(
        "https://risk.mnemom.ai", description="Risk API base URL"
    )
    api_key: str = Field(..., description="API key for authentication")


class RiskGateResult(BaseModel):
    """Result of a risk gate check."""

    allowed: bool = Field(..., description="Whether the action passed the gate")
    assessment: Optional[RiskAssessment] = Field(
        None, description="The risk assessment (None on fetch error)"
    )
    reason: Optional[str] = Field(
        None, description="Reason for denial (None if allowed)"
    )


class TeamRiskGateResult(BaseModel):
    """Result of a team risk gate check."""

    allowed: bool = Field(..., description="Whether the team action passed the gate")
    assessment: Optional[TeamRiskAssessment] = Field(
        None, description="The team risk assessment (None on fetch error)"
    )
    reason: Optional[str] = Field(
        None, description="Reason for denial (None if allowed)"
    )


class RiskGate:
    """Gate that checks agent risk before allowing an action.

    Example::

        gate = RiskGate(RiskGateConfig(
            api_key="sk-...",
            max_risk_level="medium",
            max_risk_score=0.5,
        ))
        result = await gate.check("agent-123", financial_context(1000))
        if not result.allowed:
            print(f"Denied: {result.reason}")
    """

    def __init__(self, config: RiskGateConfig) -> None:
        self.config = config

    async def check(self, agent_id: str, context: RiskContext) -> RiskGateResult:
        """Check whether an agent's risk is within acceptable bounds.

        Args:
            agent_id: Agent identifier to check.
            context: Action context for the risk assessment.

        Returns:
            RiskGateResult indicating whether the action is allowed and why.
        """
        async with RiskClient(
            api_key=self.config.api_key,
            base_url=self.config.base_url,
        ) as client:
            try:
                assessment = await client.assess_risk(agent_id, context)
            except Exception as e:
                return RiskGateResult(allowed=False, assessment=None, reason=str(e))

        if (
            self.config.max_risk_score is not None
            and assessment.risk_score > self.config.max_risk_score
        ):
            return RiskGateResult(
                allowed=False,
                assessment=assessment,
                reason=(
                    f"Risk score {assessment.risk_score} exceeds "
                    f"maximum {self.config.max_risk_score}"
                ),
            )

        if self.config.max_risk_level is not None:
            max_ordinal = _RISK_LEVEL_ORDINALS.get(self.config.max_risk_level, 0)
            actual_ordinal = _RISK_LEVEL_ORDINALS.get(assessment.risk_level, 0)
            if actual_ordinal > max_ordinal:
                return RiskGateResult(
                    allowed=False,
                    assessment=assessment,
                    reason=(
                        f"Risk level {assessment.risk_level} exceeds "
                        f"maximum {self.config.max_risk_level}"
                    ),
                )

        return RiskGateResult(allowed=True, assessment=assessment)


class TeamRiskGate:
    """Gate that checks team risk before allowing an action.

    Example::

        gate = TeamRiskGate(RiskGateConfig(
            api_key="sk-...",
            max_risk_level="medium",
        ))
        result = await gate.check(["agent-1", "agent-2"], delegation_context("audit"))
        if not result.allowed:
            print(f"Denied: {result.reason}")
    """

    def __init__(self, config: RiskGateConfig) -> None:
        self.config = config

    async def check(
        self,
        agent_ids: list[str],
        context: RiskContext,
    ) -> TeamRiskGateResult:
        """Check whether a team's risk is within acceptable bounds.

        Args:
            agent_ids: List of agent identifiers to check.
            context: Action context for the risk assessment.

        Returns:
            TeamRiskGateResult indicating whether the action is allowed and why.
        """
        async with RiskClient(
            api_key=self.config.api_key,
            base_url=self.config.base_url,
        ) as client:
            try:
                assessment = await client.assess_team_risk(agent_ids, context)
            except Exception as e:
                return TeamRiskGateResult(
                    allowed=False, assessment=None, reason=str(e)
                )

        if (
            self.config.max_risk_score is not None
            and assessment.team_risk_score > self.config.max_risk_score
        ):
            return TeamRiskGateResult(
                allowed=False,
                assessment=assessment,
                reason=(
                    f"Team risk score {assessment.team_risk_score} exceeds "
                    f"maximum {self.config.max_risk_score}"
                ),
            )

        if self.config.max_risk_level is not None:
            max_ordinal = _RISK_LEVEL_ORDINALS.get(self.config.max_risk_level, 0)
            actual_ordinal = _RISK_LEVEL_ORDINALS.get(
                assessment.team_risk_level, 0
            )
            if actual_ordinal > max_ordinal:
                return TeamRiskGateResult(
                    allowed=False,
                    assessment=assessment,
                    reason=(
                        f"Team risk level {assessment.team_risk_level} exceeds "
                        f"maximum {self.config.max_risk_level}"
                    ),
                )

        return TeamRiskGateResult(allowed=True, assessment=assessment)
