from typing import Any, cast
from uuid import uuid4

import pytest

from frisk_sdk import DecisionOutcome
from frisk_sdk.adapters.strands.frisk_tool_hook import FriskStrandsToolHook
from frisk_sdk.core.session_registry import SessionRegistry
from frisk_sdk.errors.invocation_errors import (
    InvalidFriskSessionError,
    MissingFriskSessionContextError,
)


class DummyFrisk:
    def __init__(self):
        self.logger_name = "frisk"
        self._adapter = _Adapter()


class _Adapter:
    def normalize_tool_call(self, tool_use: Any):
        return DummyToolCallInfo(
            id=tool_use.get("toolUseId", "id1"),
            name=tool_use.get("name", "tool"),
            args=tool_use.get("input", {}),
        )


class DummyFriskSession:
    def __init__(self, result, session_id=None):
        self._result = result
        self.session_id = session_id or uuid4()

    def decide_tool_call(self, *, tool_call, agent_state):
        return self._result

    def init_tracing(self, remote_session_id, inputs):
        self.init_tracing_calls = getattr(self, "init_tracing_calls", [])
        self.init_tracing_calls.append(
            {"remote_session_id": remote_session_id, "inputs": inputs}
        )

    def end_tracing(self):
        self.end_tracing_calls = getattr(self, "end_tracing_calls", 0) + 1


class DummyToolCallInfo:
    def __init__(self, name="tool", id="id1", args=None):
        self.name = name
        self.id = id
        self.args = args or {}


class DummyProcessToolCallResult:
    def __init__(self, outcome, reason="reason"):
        self.outcome = outcome
        self.reason = reason


class DummyBeforeToolCallEvent:
    def __init__(
        self, tool_use: dict[str, Any], invocation_state: dict[str, Any] | None
    ):
        self.tool_use = tool_use
        self.invocation_state = invocation_state
        self.cancel_tool = False


class DummyBeforeInvocationEvent:
    def __init__(
        self,
        invocation_state: dict[str, Any] | None,
        messages: list[dict[str, Any]] | None = None,
    ):
        self.invocation_state = invocation_state
        self.messages = messages


class DummyAfterInvocationEvent:
    def __init__(self, invocation_state: dict[str, Any] | None):
        self.invocation_state = invocation_state


def test_before_tool_call_deny_sets_cancel_message():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)

    result = DummyProcessToolCallResult(
        outcome=DecisionOutcome.DENY, reason="policy denied"
    )
    session_id = str(uuid4())
    session = DummyFriskSession(result, session_id)

    registry = SessionRegistry()
    registry.register(session_id, cast(Any, session))

    event = DummyBeforeToolCallEvent(
        tool_use={
            "toolUseId": "tooluse-1",
            "name": "weather",
            "input": {"city": "sf"},
        },
        invocation_state={"frisk_session_id": session_id},
    )
    hook.before_tool_call(cast(Any, event))

    assert isinstance(event.cancel_tool, str)
    assert "denied" in event.cancel_tool
    assert "policy denied" in event.cancel_tool

    registry.unregister(session_id)


def test_before_tool_call_allow_keeps_cancel_false():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)

    result = DummyProcessToolCallResult(outcome=DecisionOutcome.ALLOW)
    session_id = str(uuid4())
    session = DummyFriskSession(result, session_id)

    registry = SessionRegistry()
    registry.register(session_id, cast(Any, session))

    event = DummyBeforeToolCallEvent(
        tool_use={
            "toolUseId": "tooluse-1",
            "name": "weather",
            "input": {"city": "sf"},
        },
        invocation_state={"frisk_session_id": session_id},
    )

    hook.before_tool_call(cast(Any, event))

    assert event.cancel_tool is False
    registry.unregister(session_id)


def test_before_invocation_initializes_tracing():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)
    session_id = str(uuid4())
    session = DummyFriskSession(result=None, session_id=session_id)
    registry = SessionRegistry()
    registry.register(session_id, cast(Any, session))

    event = DummyBeforeInvocationEvent(
        invocation_state={"frisk_session_id": session_id}
    )
    hook.before_invocation(cast(Any, event))

    assert len(session.init_tracing_calls) == 1
    assert session.init_tracing_calls[0]["inputs"] == {
        "messages": None,
        "invocation_state": {"frisk_session_id": session_id},
    }
    assert isinstance(session.init_tracing_calls[0]["remote_session_id"], str)
    registry.unregister(session_id)


def test_after_invocation_ends_tracing():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)
    session_id = str(uuid4())
    session = DummyFriskSession(result=None, session_id=session_id)
    registry = SessionRegistry()
    registry.register(session_id, cast(Any, session))

    event = DummyAfterInvocationEvent(invocation_state={"frisk_session_id": session_id})
    hook.after_invocation(cast(Any, event))

    assert session.end_tracing_calls == 1
    registry.unregister(session_id)


def test_get_frisk_session_from_invocation_state_valid():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)

    session_id = str(uuid4())
    session = DummyFriskSession(None, session_id)
    registry = SessionRegistry()
    registry.register(session_id, cast(Any, session))

    out = hook.get_frisk_session_from_invocation_state({"frisk_session_id": session_id})
    assert out is session

    registry.unregister(session_id)


def test_get_frisk_session_from_nested_invocation_state_valid():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)

    session_id = str(uuid4())
    session = DummyFriskSession(None, session_id)
    registry = SessionRegistry()
    registry.register(session_id, cast(Any, session))

    out = hook.get_frisk_session_from_invocation_state(
        {"invocation_state": {"frisk_session_id": session_id}}
    )
    assert out is session

    registry.unregister(session_id)


def test_get_frisk_session_from_invocation_state_missing():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)
    with pytest.raises(MissingFriskSessionContextError):
        hook.get_frisk_session_from_invocation_state({})


def test_get_frisk_session_from_invocation_state_invalid():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)
    with pytest.raises(InvalidFriskSessionError):
        hook.get_frisk_session_from_invocation_state({"frisk_session_id": object()})


def test_before_tool_call_missing_session_raises():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)
    event = DummyBeforeToolCallEvent(
        tool_use={
            "toolUseId": "tooluse-1",
            "name": "weather",
            "input": {"city": "sf"},
        },
        invocation_state={},
    )
    with pytest.raises(MissingFriskSessionContextError):
        hook.before_tool_call(cast(Any, event))


def test_get_context_from_invocation_state_filters_framework_keys():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)
    invocation_state = {
        "frisk_session_id": "sid",
        "agent": object(),
        "model": object(),
        "messages": [],
        "tool_config": {},
        "request_state": {},
        "event_loop_cycle_id": "c1",
        "event_loop_cycle_span": {},
        "event_loop_cycle_trace": {},
        "system_prompt": "sp",
        "context": {"keep": 1, "frisk_session_id": "sid", "agent": "x"},
        "keep_top_level": 2,
    }

    ctx = hook.get_context_from_invocation_state(invocation_state)
    assert ctx == {"keep": 1}


def test_get_context_from_invocation_state_uses_top_level_when_no_context():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)
    invocation_state = {
        "frisk_session_id": "sid",
        "keep": 5,
        "agent": object(),
        "model": object(),
        "messages": [],
    }
    ctx = hook.get_context_from_invocation_state(invocation_state)
    assert ctx == {"keep": 5}


def test_get_context_from_invocation_state_invalid_context_type():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)
    with pytest.raises(InvalidFriskSessionError):
        hook.get_context_from_invocation_state({"context": "not-a-dict"})


def test_register_hooks_adds_before_tool_callback():
    frisk = DummyFrisk()
    hook = FriskStrandsToolHook(frisk)

    class _Registry:
        def __init__(self):
            self.calls = []

        def add_callback(self, event_type, callback):
            self.calls.append((event_type, callback))

    registry = _Registry()

    hook.register_hooks(registry)

    assert len(registry.calls) == 3
    callbacks = [callback for _, callback in registry.calls]
    assert hook.before_tool_call in callbacks
    assert hook.before_invocation in callbacks
    assert hook.after_invocation in callbacks
