from __future__ import annotations

from typing import Any, Optional
from uuid import UUID, uuid4

from ...core import FriskSession
from ...core.logging_options import LoggingOptions
from ...core.tool_call_decision import DecisionOutcome
from ...core.session_registry import SessionRegistry
from ...errors.invocation_errors import (
    InvalidFriskSessionError,
    MissingFriskSessionContextError,
)
from ...logging import get_logger

from .utils import filter_state

from strands.hooks import (
    BeforeToolCallEvent,
    BeforeInvocationEvent,
    AfterInvocationEvent,
)


class FriskStrandsToolHook:
    """Strands hook provider that enforces Frisk decisions before tool execution."""

    def __init__(self, frisk: Any, *, logging: Optional[LoggingOptions] = None):
        self._frisk_handle = frisk
        logger_name = (logging or {}).get(
            "logger_name"
        ) or f"{frisk.logger_name}.tool_hook"
        self._logger = (logging or {}).get("logger") or get_logger(logger_name)
        self.logger_name = self._logger.name

    def register_hooks(self, registry: Any, **kwargs: Any) -> None:
        """Register this hook with a Strands HookRegistry."""
        registry.add_callback(BeforeToolCallEvent, self.before_tool_call)
        registry.add_callback(BeforeInvocationEvent, self.before_invocation)
        registry.add_callback(AfterInvocationEvent, self.after_invocation)

    def before_invocation(self, event: BeforeInvocationEvent) -> None:
        """BeforeInvocation callback that could be used to enforce session context on invocation."""
        frisk_session: FriskSession[Any] = self.get_frisk_session_from_invocation_state(
            event.invocation_state
        )

        run_id = str(uuid4())

        event_dict = {
            "messages": event.messages,
            "invocation_state": event.invocation_state,
        }

        frisk_session.init_tracing(remote_session_id=run_id, inputs=event_dict)

    def after_invocation(self, event: AfterInvocationEvent) -> None:
        """AfterInvocation callback that could be used to perform actions after invocation completes."""
        frisk_session: FriskSession[Any] = self.get_frisk_session_from_invocation_state(
            event.invocation_state
        )
        frisk_session.end_tracing()

    def before_tool_call(self, event: BeforeToolCallEvent) -> None:
        """BeforeToolCall callback that blocks denied tool calls."""
        normalized_tool_call = self._frisk_handle._adapter.normalize_tool_call(
            event.tool_use
        )
        try:
            frisk_session: FriskSession[Any] = (
                self.get_frisk_session_from_invocation_state(event.invocation_state)
            )

            state = self.get_context_from_invocation_state(event.invocation_state)
            tool_call_decision = frisk_session.decide_tool_call(
                tool_call=normalized_tool_call, agent_state=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,
            )

            if tool_call_decision.outcome == DecisionOutcome.DENY:
                deny_reason = (
                    tool_call_decision.reason
                    if tool_call_decision.reason
                    else "Denied by Frisk policy"
                )
                event.cancel_tool = (
                    f"Call to {normalized_tool_call.name} denied. Reason: {deny_reason}"
                )

        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 get_frisk_session_from_invocation_state(
        self, invocation_state: dict[str, Any] | None
    ) -> FriskSession[Any]:
        """Resolve and validate the Frisk session from Strands invocation_state."""
        if invocation_state is None:
            raise MissingFriskSessionContextError()

        if invocation_state.get("invocation_state", None) is not None:
            nested_invocation_state = invocation_state["invocation_state"]
            if nested_invocation_state is None:
                raise MissingFriskSessionContextError()
            if not isinstance(nested_invocation_state, dict):
                raise InvalidFriskSessionError()
            invocation_state = nested_invocation_state

        if invocation_state is None:
            raise MissingFriskSessionContextError()

        frisk_session_id = invocation_state.get("frisk_session_id", None)
        if frisk_session_id is None:
            raise MissingFriskSessionContextError()

        if not isinstance(frisk_session_id, (UUID, str)):
            raise InvalidFriskSessionError()

        registry = SessionRegistry()
        try:
            return registry.get(frisk_session_id)
        except Exception as exc:
            raise InvalidFriskSessionError() from exc

    def get_context_from_invocation_state(
        self, invocation_state: dict[str, Any] | None
    ) -> dict[str, Any]:
        """Extract context from Strands invocation_state."""
        if invocation_state is None:
            raise MissingFriskSessionContextError()

        context = invocation_state.get("context", None)

        if context is None:
            context = invocation_state

        if context is None:
            raise MissingFriskSessionContextError()

        if not isinstance(context, dict):
            raise InvalidFriskSessionError()

        # return new filtered context
        context = filter_state(context)

        return context
