"""Risk assessment API methods.

Fetches risk assessments from the Mnemom Risk API for
individual agents and multi-agent teams.
"""

from __future__ import annotations

from urllib.parse import quote

import httpx

from mnemom_types import RiskAssessment, RiskContext, TeamRiskAssessment

DEFAULT_BASE_URL = "https://risk.mnemom.ai"


class RiskClient:
    """Async client for the Mnemom Risk API.

    Example::

        async with RiskClient(api_key="sk-...") as client:
            assessment = await client.assess_risk("agent-123", context)
            print(assessment.risk_level, assessment.risk_score)
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = DEFAULT_BASE_URL,
    ) -> None:
        self.base_url = base_url
        self._client = httpx.AsyncClient(
            base_url=base_url,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
        )

    async def __aenter__(self) -> RiskClient:
        return self

    async def __aexit__(self, *exc: object) -> None:
        await self.close()

    async def close(self) -> None:
        """Close the underlying HTTP client."""
        await self._client.aclose()

    async def assess_risk(
        self,
        agent_id: str,
        context: RiskContext,
    ) -> RiskAssessment:
        """Assess risk for a single agent performing an action.

        Args:
            agent_id: Agent identifier.
            context: Action context describing what the agent intends to do.

        Returns:
            The risk assessment result.

        Raises:
            PermissionError: If the feature is gated (HTTP 403).
            LookupError: If the agent is not found (HTTP 404).
            httpx.HTTPStatusError: For other non-2xx responses.
        """
        url = f"/v1/risk/{quote(agent_id, safe='')}/assess"
        response = await self._client.post(url, json=context.model_dump(exclude_none=True))

        if response.status_code == 403:
            raise PermissionError("Risk assessment feature is gated")
        if response.status_code == 404:
            raise LookupError("Agent not found")

        response.raise_for_status()
        return RiskAssessment(**response.json())

    async def assess_team_risk(
        self,
        agent_ids: list[str],
        context: RiskContext,
    ) -> TeamRiskAssessment:
        """Assess risk for a team of agents performing an action together.

        Args:
            agent_ids: List of agent identifiers.
            context: Action context describing the team operation.

        Returns:
            The team risk assessment result.

        Raises:
            PermissionError: If the feature is gated (HTTP 403).
            httpx.HTTPStatusError: For other non-2xx responses.
        """
        response = await self._client.post(
            "/v1/risk/team/assess",
            json={
                "agent_ids": agent_ids,
                "context": context.model_dump(exclude_none=True),
            },
        )

        if response.status_code == 403:
            raise PermissionError("Risk assessment feature is gated")

        response.raise_for_status()
        return TeamRiskAssessment(**response.json())

    async def get_risk_history(
        self,
        agent_id: str,
        limit: int = 20,
    ) -> list[RiskAssessment]:
        """Retrieve the risk assessment history for an agent.

        Args:
            agent_id: Agent identifier.
            limit: Maximum number of assessments to return.

        Returns:
            List of past risk assessments.

        Raises:
            PermissionError: If the feature is gated (HTTP 403).
            LookupError: If the agent is not found (HTTP 404).
            httpx.HTTPStatusError: For other non-2xx responses.
        """
        url = f"/v1/risk/{quote(agent_id, safe='')}/history"
        response = await self._client.get(url, params={"limit": limit})

        if response.status_code == 403:
            raise PermissionError("Risk assessment feature is gated")
        if response.status_code == 404:
            raise LookupError("Agent not found")

        response.raise_for_status()
        return [RiskAssessment(**item) for item in response.json()]

    async def get_assessment(self, assessment_id: str) -> RiskAssessment:
        """Retrieve a specific risk assessment by ID.

        Args:
            assessment_id: Assessment identifier (prefix: ra-).

        Returns:
            The risk assessment.

        Raises:
            PermissionError: If the feature is gated (HTTP 403).
            LookupError: If the assessment is not found (HTTP 404).
            httpx.HTTPStatusError: For other non-2xx responses.
        """
        url = f"/v1/risk/assessments/{quote(assessment_id, safe='')}"
        response = await self._client.get(url)

        if response.status_code == 403:
            raise PermissionError("Risk assessment feature is gated")
        if response.status_code == 404:
            raise LookupError("Assessment not found")

        response.raise_for_status()
        return RiskAssessment(**response.json())

    async def get_team_assessment(self, assessment_id: str) -> TeamRiskAssessment:
        """Retrieve a specific team risk assessment by ID.

        Args:
            assessment_id: Assessment identifier (prefix: tra-).

        Returns:
            The team risk assessment.

        Raises:
            PermissionError: If the feature is gated (HTTP 403).
            LookupError: If the assessment is not found (HTTP 404).
            httpx.HTTPStatusError: For other non-2xx responses.
        """
        url = f"/v1/risk/assessments/team/{quote(assessment_id, safe='')}"
        response = await self._client.get(url)

        if response.status_code == 403:
            raise PermissionError("Risk assessment feature is gated")
        if response.status_code == 404:
            raise LookupError("Team assessment not found")

        response.raise_for_status()
        return TeamRiskAssessment(**response.json())
