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, TYPE_CHECKING
from uuid import UUID

from langgraph.types import Command

from cuid2 import cuid_wrapper

from ...core.logging_options import LoggingOptions
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 FriskToolMiddleware(AgentMiddleware):
    """LangChain middleware that blocks or allows tool calls using Frisk decisions."""

    def __init__(
        self, frisk: FriskLangchain, *, 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()

    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
        normalized_tool_call = self._frisk_handle.normalize_tool_call(initial_tool_call)
        try:
            frisk_session = self.resolve_frisk_session(request.runtime)
            tool_call_decision = frisk_session.decide_tool_call(
                tool_call=normalized_tool_call, agent_state=request.state
            )

            self._logger.debug(
                "Tool call result: Tool name=%s, tool call ID=%s, outcome=%s, reason=%s",
                normalized_tool_call.name,
                normalized_tool_call.id,
                tool_call_decision.outcome,
                tool_call_decision.reason,
            )

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

    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()

        # 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: NormalizedToolCall,
        request: ToolCallRequest,
        handler: Callable[[Any], Awaitable[ToolMessage | Command[Any]]],
        tool_call_decision: ToolCallDecision,
    ) -> ToolMessage | Command[Any]:
        """Return denial response for blocked calls, otherwise execute tool handler."""
        if tool_call_decision.outcome == DecisionOutcome.DENY:
            return ToolMessage(
                content=f"Call to {tool_call.name} denied. Reason: {tool_call_decision.reason}",
                tool_call_id=tool_call.id,
                name=tool_call.name,
            )
        handler_result = await handler(request)
        return handler_result
