from __future__ import annotations
import asyncio

from langchain.agents.middleware import AgentMiddleware
from langchain.agents.middleware.types import StateT
from langchain_core.messages import ToolMessage, AIMessage, ToolCall

from typing import Optional, Any, TYPE_CHECKING, TypedDict, cast
from uuid import UUID

from langgraph.runtime import Runtime
from langgraph.types import interrupt

from cuid2 import cuid_wrapper

from .interrupt_response_schema import (
    interrupt_response_schema,
    InterruptResponseSchema,
    ToolCallResumeAction,
)
from ...core.logging_options import LoggingOptions
from ...core.tool_approval_request import (
    ToolApprovalRequest,
    ToolApprovalStatus,
    ToolApprovalRequestStatusUpdateResult,
)
from ...core.tool_call_decision import (
    ToolCallDecision,
    DecisionOutcome,
)
from ...core import FriskSession
from ...core.session_registry import SessionRegistry
from frisk_sdk.framework_adapter.framework_adapter import NormalizedToolCall
from ...errors.invocation_errors import (
    MissingFriskSessionContextError,
    InvalidFriskSessionError,
)
from ...logging import get_logger
from ...type_utils import normalize_object_to_dictionary

if TYPE_CHECKING:
    from .frisk_langchain import FriskLangchain


class FriskGuardControls(TypedDict, total=False):
    tool_policies: Optional[bool]


# todo: Rename FriskGuardMiddleware (fast-follow)
class FriskToolMiddleware(AgentMiddleware):
    """LangChain middleware that blocks or allows tool calls using Frisk decisions."""

    def __init__(
        self,
        frisk: FriskLangchain,
        *,
        controls: Optional[FriskGuardControls] = None,
        logging: Optional[LoggingOptions] = None,
    ):
        self._frisk_handle: Any = frisk

        logger_name = (logging or {}).get(
            "logger_name"
        ) or f"{frisk.logger_name}.tool_middleware"
        self._logger = (logging or {}).get("logger") or get_logger(logger_name)
        self._logger_name = self._logger.name

        self._cuid_generator = cuid_wrapper()
        self._tool_policies_enabled = (controls or {}).get(
            "tool_policies", True
        ) is not False

    def after_model(self, state: StateT, runtime: Runtime) -> dict[str, Any] | None:
        """Process AI model output and handle tool call escalations.

        This method intercepts tool calls from the AI model, evaluates them against
        Frisk policies, and handles escalations by interrupting for human approval.
        """
        return asyncio.run(self.aafter_model(state, runtime))

    async def aafter_model(
        self, state: StateT, runtime: Runtime
    ) -> dict[str, Any] | None:
        """Process AI model output and handle tool call escalations.

        This method intercepts tool calls from the AI model, evaluates them against
        Frisk policies, and handles escalations by interrupting for human approval.
        """
        # todo: It should be possible to disable policy blocking and still collect telemetry data. https://linear.app/friskai/issue/POL-200/feature-it-should-be-possible-to-disable-tool-blocking-while-still
        if not self._tool_policies_enabled:
            return None

        messages = state.get("messages", [])
        if not messages:
            return None

        # Find the last AI message with tool calls
        tool_calls: list[ToolCall] = self._get_tool_calls_from_state(messages)
        if len(tool_calls) == 0:
            return None

        # Normalize all tool calls
        normalized_tool_calls: list[NormalizedToolCall] = [
            self._frisk_handle.normalize_tool_call(tool_call)
            for tool_call in tool_calls
        ]

        normalized_tool_calls_by_id: dict[str, NormalizedToolCall] = {
            normalized_tool_call.id: normalized_tool_call
            for normalized_tool_call in normalized_tool_calls
        }

        frisk_session = self.resolve_frisk_session(runtime)

        # Create decision mapping for all tool calls
        tool_call_decisions_by_tool_call_id: dict[str, ToolCallDecision] = {}
        for tool_call in normalized_tool_calls:
            tool_call_decisions_by_tool_call_id[tool_call.id] = (
                frisk_session.decide_tool_call(
                    tool_call=tool_call, agent_state=cast(dict[str, Any], state)
                )
            )

        # Identify tool calls that need escalation
        tool_approval_request_inputs = [
            {
                "tool_call_id": tool_call_id,
                "policy_id": decision.policy_id or "unknown",
                "policy_version_id": decision.policy_version_id or "unknown",
                "tool_name": normalized_tool_calls_by_id[tool_call_id].name,
            }
            for tool_call_id, decision in tool_call_decisions_by_tool_call_id.items()
            if decision.outcome == DecisionOutcome.ESCALATE
        ]

        # Get or create approval requests for escalated calls
        approval_requests_by_tool_call_id: dict[str, ToolApprovalRequest | None] = dict(
            await self._frisk_handle.get_or_create_many_tool_approval_requests(
                tool_approval_request_inputs=tool_approval_request_inputs,
                session_id=frisk_session.id,
            )
        )

        # Filter for pending escalations
        pending_approval_requests_by_tool_call_id = (
            self._filter_pending_approval_requests_by_tool_call_id(
                approval_requests_by_tool_call_id
            )
        )

        interrupt_message = (
            "Provide a map of tool call IDs to either 'retry' or 'cancel' "
            "for each tool call that is pending approval."
        )
        interrupt_error = None

        # Loop while there are pending escalations
        while pending_approval_requests_by_tool_call_id:
            # Send interrupt and get response
            tool_calls_pending_approval = [
                normalized_tool_calls_by_id[tool_call_id]
                for tool_call_id in pending_approval_requests_by_tool_call_id.keys()
            ]
            interrupt_payload = self._make_interrupt_payload(
                interrupt_error=interrupt_error,
                interrupt_message=interrupt_message,
                tool_calls_pending_approval=tool_calls_pending_approval,
            )

            self._logger.debug(
                f"Some tool calls were escalated. Interrupt message: {interrupt_payload}"
            )
            caller_interrupt_response = interrupt(interrupt_payload)
            interrupt_error = None

            retry_or_cancel_actions_by_tool_id, interrupt_response_formatting_error = (
                self.get_tool_call_resume_actions_by_id_from_interrupt_response(
                    caller_interrupt_response, approval_requests_by_tool_call_id
                )
            )

            if interrupt_response_formatting_error:
                interrupt_error = str(interrupt_response_formatting_error)
                continue

            tool_call_ids_to_cancel = [
                tool_call_id
                for tool_call_id, action in retry_or_cancel_actions_by_tool_id.items()
                if action == ToolCallResumeAction.Cancel
            ]

            # Extract request IDs to cancel, skipping any missing or failed approval requests
            approval_request_ids_to_cancel = []
            for tool_call_id in tool_call_ids_to_cancel:
                approval_request = approval_requests_by_tool_call_id.get(tool_call_id)
                if approval_request is None:
                    interrupt_error = f"Approval request not found for tool call ID {tool_call_id}. Could not cancel. Try again."
                    continue
                if approval_request is not None:
                    approval_request_ids_to_cancel.append(approval_request.id)

            # Cancel the approval requests
            try:
                if approval_request_ids_to_cancel:
                    canceled_tool_approval_requests_by_request_id: (
                        dict[str, ToolApprovalRequestStatusUpdateResult | None] | None
                    ) = await self._frisk_handle.cancel_many_tool_approval_requests(
                        approval_request_ids=approval_request_ids_to_cancel
                    )
                else:
                    canceled_tool_approval_requests_by_request_id = {}
            except Exception as e:
                self._logger.error(e)
                interrupt_error = f"Failed to cancel approval requests. Error: {str(e)}"
                continue

            # Update approval requests with canceled ones
            for tool_call_id in tool_call_ids_to_cancel:
                approval_request_cancellation_successful = False
                approval_request = approval_requests_by_tool_call_id[tool_call_id]
                if approval_request:
                    approval_request_id = approval_request.id
                    if (
                        approval_request_id
                        in canceled_tool_approval_requests_by_request_id
                    ):
                        cancellation_result = (
                            canceled_tool_approval_requests_by_request_id[
                                approval_request_id
                            ]
                        )
                        canceled_approval_request = None
                        if cancellation_result:
                            canceled_approval_request = cancellation_result.result
                        if canceled_approval_request:
                            approval_requests_by_tool_call_id[tool_call_id] = (
                                canceled_approval_request
                            )
                            approval_request_cancellation_successful = True
                if not approval_request_cancellation_successful:
                    exc = f"Failed to cancel approval request for tool call ID {tool_call_id}. Try again."
                    self._logger.error(exc)
                    interrupt_error = str(exc)
                    continue  # Restart loop

            # Recalculate pending escalations
            pending_approval_requests_by_tool_call_id = (
                self._filter_pending_approval_requests_by_tool_call_id(
                    approval_requests_by_tool_call_id
                )
            )

        # Process all tool calls and separate allowed from denied
        denied_tool_call_messages: list[ToolMessage] = []
        for normalized_tool_call in normalized_tool_calls:
            tool_call_decision = tool_call_decisions_by_tool_call_id[
                normalized_tool_call.id
            ]

            # If explicitly denied by policy, deny.
            is_denied_by_policy = tool_call_decision.outcome == DecisionOutcome.DENY
            if is_denied_by_policy:
                denied_tool_call_messages.append(
                    self._create_denied_tool_call_message(
                        normalized_tool_call, tool_call_decision
                    )
                )
                continue

            # Check escalation status
            approval_request = approval_requests_by_tool_call_id.get(
                normalized_tool_call.id
            )
            status = approval_request.status if approval_request else None

            is_escalated_canceled = status == ToolApprovalStatus.CANCELED
            if is_escalated_canceled:
                denied_tool_call_messages.append(
                    self._create_escalated_canceled_tool_message(normalized_tool_call)
                )
                continue

            is_escalated_rejected = status == ToolApprovalStatus.REJECTED
            if is_escalated_rejected:
                denied_tool_call_messages.append(
                    self._create_escalated_rejected_tool_message(normalized_tool_call)
                )
                continue

        # Update the AI message with allowed tool calls
        message_updates = [*denied_tool_call_messages]
        return {"messages": message_updates}

    def _filter_pending_approval_requests_by_tool_call_id(
        self,
        approval_requests_by_tool_call_id: dict[str, ToolApprovalRequest | None],
    ) -> dict[str, ToolApprovalRequest | None]:
        return {
            tool_call_id: approval_request
            for tool_call_id, approval_request in approval_requests_by_tool_call_id.items()
            if (
                approval_request.status
                if approval_request
                else ToolApprovalStatus.PENDING
            )
            == ToolApprovalStatus.PENDING
        }

    def _make_interrupt_payload(
        self,
        *,
        interrupt_error: Optional[str] = None,
        interrupt_message: str,
        tool_calls_pending_approval: list[NormalizedToolCall],
    ):
        interrupt_payload = {
            "__frisk": True,
            "message": f"{interrupt_error} {interrupt_message}"
            if interrupt_error
            else interrupt_message,
            "escalated_tool_calls": {
                tool_call.id: {
                    "tool_name": tool_call.name,
                }
                for tool_call in tool_calls_pending_approval
            },
            "response_schema": interrupt_response_schema,
        }
        return interrupt_payload

    def _get_tool_calls_from_state(self, messages):
        last_ai_msg = next(
            (msg for msg in reversed(messages) if isinstance(msg, AIMessage)), None
        )
        if not last_ai_msg or not last_ai_msg.tool_calls:
            return []
        return last_ai_msg.tool_calls or []

    def resolve_frisk_session(self, runtime: Any) -> FriskSession[Any]:
        """Resolve and validate the Frisk session from LangGraph runtime context."""
        runtime_context = runtime.context
        if runtime_context is None:
            raise MissingFriskSessionContextError()

        normalized_runtime_context = normalize_object_to_dictionary(runtime_context)
        frisk_session_id = normalized_runtime_context.get("frisk_session_id", None)
        if frisk_session_id is None:
            raise MissingFriskSessionContextError()

        # Validate frisk_session_id is a UUID or string
        if not isinstance(frisk_session_id, (UUID, str)):
            raise InvalidFriskSessionError()

        # Normalize frisk_session_id to the string key used by SessionRegistry
        session_id_key = (
            str(frisk_session_id)
            if isinstance(frisk_session_id, UUID)
            else frisk_session_id
        )

        # Resolve session from registry
        registry = SessionRegistry()
        try:
            return registry.get(session_id_key)
        except Exception as exc:
            # Re-raise as InvalidFriskSessionError for backwards compatibility
            raise InvalidFriskSessionError() from exc

    def _create_escalated_canceled_tool_message(
        self, tool_call: NormalizedToolCall
    ) -> ToolMessage:
        """Create a tool message for a tool call that was canceled after escalation."""
        content = (
            f"The caller canceled the request to {tool_call.name} "
            f"after it was escalated for approval."
        )
        return self._create_synthetic_tool_call_message(
            tool_call, content, status="error"
        )

    def _create_escalated_rejected_tool_message(
        self, tool_call: NormalizedToolCall
    ) -> ToolMessage:
        """Create a tool message for a tool call that was rejected after escalation."""
        content = (
            f"Call to {tool_call.name} was denied after it was escalated for approval."
        )
        return self._create_synthetic_tool_call_message(
            tool_call, content, status="error"
        )

    def _create_denied_tool_call_message(
        self, tool_call: NormalizedToolCall, decision: ToolCallDecision
    ) -> ToolMessage:
        """Create a tool message for a tool call that was denied by policy."""
        content = f"Call to {tool_call.name} denied. Reason: {decision.reason}"
        return self._create_synthetic_tool_call_message(
            tool_call, content, status="error"
        )

    def _create_synthetic_tool_call_message(
        self, tool_call: NormalizedToolCall, content: str, status: Optional[str] = None
    ) -> ToolMessage:
        """Create a synthetic tool message with the given content."""
        if status:
            return ToolMessage(
                tool_call_id=tool_call.id,
                content=content,
                name=tool_call.name,
                status=status,
            )
        return ToolMessage(
            tool_call_id=tool_call.id,
            content=content,
            name=tool_call.name,
        )

    def get_tool_call_resume_actions_by_id_from_interrupt_response(
        self,
        caller_interrupt_response: Any,
        approval_requests_by_tool_call_id: dict[str, ToolApprovalRequest | None],
    ) -> tuple[InterruptResponseSchema, Optional[Exception]]:
        formatted_response: InterruptResponseSchema = {
            tool_call_id: ToolCallResumeAction.Retry
            for tool_call_id in approval_requests_by_tool_call_id.keys()
        }
        errors: list[Exception] = []

        if not isinstance(caller_interrupt_response, dict):
            return (
                formatted_response,
                ValueError(
                    "Received an invalid interrupt response. Expected a JSON object using the given response_schema."
                ),
            )

        for tool_call_id, decision_string in caller_interrupt_response.items():
            try:
                # Validate that all tool call IDs in decisions are valid
                if tool_call_id not in approval_requests_by_tool_call_id:
                    raise ValueError(
                        f"The given tool call ID does not map to a tool call with an escalation status: {tool_call_id}"
                    )
                # Per InterruptResponseSchema, omitted or null values mean "retry"
                if decision_string is None:
                    continue
                else:
                    decision = ToolCallResumeAction(decision_string)
                formatted_response[tool_call_id] = decision
            except Exception as e:
                errors.append(e)

        error = (
            None
            if len(errors) == 0
            else ValueError(
                "Received an invalid interrupt response: "
                + "\n".join(str(e) for e in errors)
            )
        )

        return (
            formatted_response,
            error,
        )
