import pytest
import asyncio
from unittest.mock import MagicMock
from uuid import uuid4
from typing import cast, Any, TYPE_CHECKING
from langchain_core.tools import BaseTool
from langchain.tools.tool_node import ToolCallRequest

from frisk_sdk import DecisionOutcome
from frisk_sdk.adapters.langchain.frisk_tool_middleware import FriskToolMiddleware
from frisk_sdk.errors.invocation_errors import (
    MissingFriskSessionContextError,
    InvalidFriskSessionError,
)
from frisk_sdk.core.session_registry import SessionRegistry


# FriskToolMiddleware expects a `FriskLangchain` instance, but in tests we only
# need the small surface area used by the middleware.
class DummyFrisk:
    def __init__(self, *args: Any, **kwargs: Any):
        self.logger_name = "frisk"
        self._tracer = None
        self._adapter = MagicMock()
        # Instance-level MagicMock so tests can set `.return_value`.
        self.normalize_tool_call: MagicMock = MagicMock()


if TYPE_CHECKING:
    from typing import Protocol

    class DummyFriskProtocol(Protocol):
        logger_name: str
        _tracer: Any
        _adapter: Any
        normalize_tool_call: MagicMock

    # NOTE: don't alias DummyFrisk to the protocol; it breaks construction and confuses pyright.


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


class DummyToolCall:
    def __init__(self, name="tool", id="id1"):
        self.name = name
        self.id = id


class DummyToolCallRequest(ToolCallRequest):
    def __init__(
        self,
        tool: str | BaseTool | None = None,
        tool_call: Any = None,
        state: Any = None,
        runtime: Any = None,
    ):
        self.tool = cast(Any, tool)
        self.tool_call = tool_call
        self.state = state
        self.runtime = runtime


class DummyToolRuntime:
    def __init__(self, context: dict[str, Any]):
        self.context = context


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


class DummyToolMessage:
    def __init__(self, content, tool_call_id, name):
        self.content = content
        self.tool_call_id = tool_call_id
        self.name = name


class DummyCommand:
    def __init__(self, update):
        self.update = update


@pytest.mark.asyncio
async def test_awrap_tool_call_deny(monkeypatch):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall()
    frisk.normalize_tool_call.return_value = tool_call
    result = DummyProcessToolCallResult(outcome=DecisionOutcome.DENY, reason="nope")

    session_id = uuid4()
    session = DummyFriskSession(result, session_id)

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

    runtime = DummyToolRuntime({"frisk_session_id": str(session_id)})
    request = DummyToolCallRequest(
        tool="tool", tool_call="call", state={}, runtime=runtime
    )

    async def handler(_request):
        return DummyToolMessage("ok", tool_call.id, tool_call.name)

    # Patch ToolMessage
    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.ToolMessage",
        DummyToolMessage,
    )
    mw = FriskToolMiddleware(cast(Any, frisk))
    out = await mw.awrap_tool_call(request, cast(Any, handler))
    assert isinstance(out, DummyToolMessage)
    assert "denied" in out.content
    assert out.tool_call_id == tool_call.id
    assert out.name == tool_call.name

    # Cleanup
    registry.unregister(str(session_id))


@pytest.mark.asyncio
async def test_awrap_tool_call_allow_sync(monkeypatch):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall()
    frisk.normalize_tool_call.return_value = tool_call
    result = DummyProcessToolCallResult(outcome="allow")

    session_id = uuid4()
    session = DummyFriskSession(result, session_id)

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

    runtime = DummyToolRuntime({"frisk_session_id": str(session_id)})
    request = DummyToolCallRequest(
        tool="tool", tool_call="call", state={}, runtime=runtime
    )
    handler_result = DummyToolMessage("ok", tool_call.id, tool_call.name)

    async def handler(_request):
        return handler_result

    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.ToolMessage",
        DummyToolMessage,
    )
    mw = FriskToolMiddleware(cast(Any, frisk))
    out = await mw.awrap_tool_call(request, cast(Any, handler))
    assert out is handler_result

    # Cleanup
    registry.unregister(str(session_id))


@pytest.mark.asyncio
async def test_awrap_tool_call_allow_async(monkeypatch):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall()
    frisk.normalize_tool_call.return_value = tool_call
    result = DummyProcessToolCallResult(outcome="allow")

    session_id = uuid4()
    session = DummyFriskSession(result, session_id)

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

    runtime = DummyToolRuntime({"frisk_session_id": str(session_id)})
    request = DummyToolCallRequest(
        tool="tool", tool_call="call", state={}, runtime=runtime
    )

    async def async_handler(req):
        return DummyToolMessage("async", tool_call.id, tool_call.name)

    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.ToolMessage",
        DummyToolMessage,
    )
    mw = FriskToolMiddleware(cast(Any, frisk))
    out = await mw.awrap_tool_call(request, cast(Any, async_handler))
    assert isinstance(out, DummyToolMessage)
    assert out.content == "async"

    # Cleanup
    registry.unregister(str(session_id))


@pytest.mark.asyncio
async def test_awrap_tool_call_missing_session(monkeypatch):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall()
    frisk.normalize_tool_call.return_value = tool_call
    runtime = DummyToolRuntime({})
    request = DummyToolCallRequest(
        tool="tool", tool_call="call", state={}, runtime=runtime
    )

    async def handler(_request):
        return DummyToolMessage("unused", tool_call.id, tool_call.name)

    mw = FriskToolMiddleware(cast(Any, frisk))
    with pytest.raises(MissingFriskSessionContextError):
        await mw.awrap_tool_call(request, cast(Any, handler))


@pytest.mark.asyncio
async def test_awrap_tool_call_invalid_session(monkeypatch):
    frisk = DummyFrisk("dummy_api_key")
    tool_call = DummyToolCall()
    frisk.normalize_tool_call.return_value = tool_call
    runtime = DummyToolRuntime({"frisk_session_id": object()})
    request = DummyToolCallRequest(
        tool="tool", tool_call="call", state={}, runtime=runtime
    )

    async def handler(_request):
        return DummyToolMessage("unused", tool_call.id, tool_call.name)

    mw = FriskToolMiddleware(cast(Any, frisk))
    with pytest.raises(InvalidFriskSessionError):
        await mw.awrap_tool_call(request, cast(Any, handler))


def test_get_frisk_session_from_runtime_valid():
    frisk = DummyFrisk("dummy_api_key")
    mw = FriskToolMiddleware(cast(Any, frisk))

    session_id = uuid4()
    session = DummyFriskSession(None, session_id)

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

    runtime = DummyToolRuntime({"frisk_session_id": str(session_id)})
    out = mw.resolve_frisk_session(runtime)
    assert out is session

    # Cleanup
    registry.unregister(str(session_id))


def test_get_frisk_session_from_runtime_missing():
    frisk = DummyFrisk("dummy_api_key")
    mw = FriskToolMiddleware(cast(Any, frisk))
    runtime = DummyToolRuntime({})
    with pytest.raises(MissingFriskSessionContextError):
        mw.resolve_frisk_session(runtime)


def test_get_frisk_session_from_runtime_invalid():
    frisk = DummyFrisk("dummy_api_key")
    mw = FriskToolMiddleware(cast(Any, frisk))
    runtime = DummyToolRuntime({"frisk_session_id": object()})
    with pytest.raises(InvalidFriskSessionError):
        mw.resolve_frisk_session(runtime)


def test_wrap_tool_call_sync(monkeypatch):
    frisk = DummyFrisk("dummy_api_key")
    mw = FriskToolMiddleware(cast(Any, frisk))

    async def coro(*a, **k):
        return "wrapped"

    monkeypatch.setattr(mw, "awrap_tool_call", coro)
    request = DummyToolCallRequest()
    handler = MagicMock()
    out = mw.wrap_tool_call(request, handler)
    assert out == "wrapped"


def test_awrap_tool_call_empty_tool(monkeypatch):
    frisk = DummyFrisk("dummy_api_key")
    mw = FriskToolMiddleware(cast(Any, frisk))
    request = DummyToolCallRequest(tool=None)

    async def handler(_request):
        return DummyToolMessage("unused", "id", "tool")

    monkeypatch.setattr(
        "frisk_sdk.adapters.langchain.frisk_tool_middleware.Command", DummyCommand
    )
    out = asyncio.run(mw.awrap_tool_call(request, cast(Any, handler)))
    assert isinstance(out, DummyCommand)
    assert out.update == {}
