from __future__ import annotations
import asyncio

from langchain.agents.middleware import (
    AgentMiddleware,
)
from langchain_core.messages import ToolMessage
from langchain.tools.tool_node import ToolCallRequest

from typing import Callable, Optional, Awaitable, Any
from uuid import UUID

from langgraph.types import Command

from cuid2 import cuid_wrapper

from ...core.process_tool_call_result import (
    ProcessToolCallResult,
    ProcessToolCallDecision,
)
from ...core import FriskSession
from ...core.session_registry import SessionRegistry
from frisk_sdk.framework_adapter.framework_adapter import ToolCallInfo
from ...errors.invocation_errors import (
    MissingFriskSessionContextError,
    InvalidFriskSessionError,
)
from ...logging import get_logger
from ...type_utils import normalize_object_to_dictionary


class FriskToolMiddleware(AgentMiddleware):
    """LangChain middleware that blocks or allows tool calls using Frisk decisions."""

    def __init__(self, frisk: Any, *, logger_name: Optional[str] = None):
        self._frisk_handle: Any = frisk
        logger_name = logger_name or f"{frisk.logger_name}.tool_middleware"
        self._logger = get_logger(logger_name)
        self._cuid_generator = cuid_wrapper()
        self._tracer = frisk._tracer

    def wrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[Any], ToolMessage | Command[Any]],
    ) -> ToolMessage | Command[Any]:
        """Synchronous middleware entrypoint that delegates to async logic."""

        async def async_handler(r: Any) -> ToolMessage | Command[Any]:
            return handler(r)

        return asyncio.run(self.awrap_tool_call(request, async_handler))

    async def awrap_tool_call(
        self,
        request: ToolCallRequest,
        handler: Callable[[Any], Awaitable[ToolMessage | Command[Any]]],
    ) -> ToolMessage | Command[Any]:
        """Evaluate a tool call and either short-circuit or execute the handler."""
        initial_tool = request.tool

        if initial_tool is None:
            """
              todo: When combined with the LLMOutputMiddleware, there's always an extra "empty" tool call. Figure out why this happens and fix it.
              https://linear.app/friskai/issue/SDK-24/fix-issues-in-llm-output-middleware
            """
            return Command(update={})

        initial_tool_call = request.tool_call
        tool_call_info = self._frisk_handle._framework_adapter.get_tool_call_info(
            initial_tool_call
        )
        try:
            frisk_session = self.get_frisk_session_from_runtime(request.runtime)
            process_tool_call_result = frisk_session.evaluate_tool_call(
                tool_call_info, request.state
            )

            self._logger.debug(
                "Tool call result: Tool name=%s, tool call ID=%s, decision=%s, reason=%s",
                tool_call_info.name,
                tool_call_info.id,
                process_tool_call_result.decision,
                process_tool_call_result.reason,
            )

            return await self.block_or_execute_tool_call(
                tool_call_info, request, handler, process_tool_call_result
            )
        except Exception as e:
            self._logger.error(
                "Tool call exception: Tool name=%s, tool call ID=%s, error=%s",
                tool_call_info.name,
                tool_call_info.id,
                str(e),
                exc_info=True,
            )
            raise

    def get_frisk_session_from_runtime(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 session_id is a UUID or string
        if not isinstance(frisk_session_id, (UUID, str)):
            raise InvalidFriskSessionError()

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

    async def block_or_execute_tool_call(
        self,
        tool_call_info: ToolCallInfo,
        request: ToolCallRequest,
        handler: Callable[[Any], Awaitable[ToolMessage | Command[Any]]],
        process_tool_call_result: ProcessToolCallResult,
    ) -> ToolMessage | Command[Any]:
        """Return denial response for blocked calls, otherwise execute tool handler."""
        if process_tool_call_result.decision == ProcessToolCallDecision.DENY:
            return ToolMessage(
                content=f"Call to {tool_call_info.name} denied. Reason: {process_tool_call_result.reason}",
                tool_call_id=tool_call_info.id,
                name=tool_call_info.name,
            )
        handler_result = await handler(request)
        return handler_result
