from dataclasses import dataclass
from unittest.mock import MagicMock, call

import pytest

from frisk_sdk.framework_adapter.base_framework_adapter import BaseFrameworkAdapter
from frisk_sdk.framework_adapter.framework_adapter import NormalizedToolCall
from frisk_sdk.core.frisk_session import FriskSession
from frisk_sdk.telemetry.constants import (
    SPAN_NAME_FRISK_SESSION,
    ATTRIBUTE_NAME_SESSION_PROMPT,
    ATTRIBUTE_NAME_SESSION_ID,
    ATTRIBUTE_NAME_REMOTE_SESSION_ID,
)
from frisk_sdk.core.tool_call_decision import DecisionOutcome


class _DummyAdapter(BaseFrameworkAdapter):
    def noop(self): ...


@dataclass
class _DummyResult:
    decision: str = "allow"
    rules_matched_count: int = 0
    reason: str | None = None


@pytest.mark.unit
def test_init_tracing_sets_root_span_and_user_query_attribute(
    monkeypatch: pytest.MonkeyPatch,
):
    adapter = _DummyAdapter()

    span = MagicMock(name="Span")
    tracer = MagicMock(name="Tracer")
    tracer.start_span.return_value = span

    frisk_handle = MagicMock(name="Frisk")
    frisk_handle._adapter = adapter

    s = FriskSession(frisk=frisk_handle, tracer=tracer)

    inputs = {"prompt": "hi"}
    langchain_run_id = "run-1"
    returned_span = s.init_tracing(remote_session_id=langchain_run_id, inputs=inputs)

    assert returned_span is span
    assert s.get_root_span() is span

    tracer.start_span.assert_called_once_with(SPAN_NAME_FRISK_SESSION)
    span.set_attribute.assert_has_calls(
        [
            call(ATTRIBUTE_NAME_SESSION_ID, str(s.id)),
            call(ATTRIBUTE_NAME_REMOTE_SESSION_ID, str(langchain_run_id)),
            call(ATTRIBUTE_NAME_SESSION_PROMPT, inputs.get("prompt")),
        ]
    )  # Default behavior for BaseFrameworkAdapter is to extract user query from a field called "prompt"


@pytest.mark.unit
def test_end_tracing_noop_when_no_root_span(capsys):
    from frisk_sdk.core.session_registry import SessionRegistry

    adapter = _DummyAdapter()
    tracer = MagicMock(name="Tracer")
    frisk_handle = MagicMock(name="Frisk")
    frisk_handle._adapter = adapter

    # Register session in registry so unregister doesn't fail
    s = FriskSession(
        frisk=frisk_handle,
        tracer=tracer,
    )
    registry = SessionRegistry()
    registry.register(s.id, s)

    s.end_tracing()

    out = capsys.readouterr().out
    assert out == ""

    # Verify session was unregistered
    from frisk_sdk.core.session_registry import SessionNotFoundError

    with pytest.raises(SessionNotFoundError):
        registry.get(s.id)


@pytest.mark.unit
def test_generate_tool_call_id_uses_cuid_generator(monkeypatch: pytest.MonkeyPatch):
    adapter = _DummyAdapter()
    tracer = MagicMock(name="Tracer")
    frisk_handle = MagicMock(name="Frisk")
    frisk_handle._adapter = adapter

    s = FriskSession(frisk=frisk_handle, tracer=tracer)
    s._cuid_generator = lambda: "cuid-123"  # boundary: deterministic ID

    assert s.generate_tool_call_id() == "cuid-123"


@pytest.mark.integration
def test_decide_tool_call_injects_trace_context_and_delegates(
    monkeypatch: pytest.MonkeyPatch,
):
    adapter = _DummyAdapter()

    tracer = MagicMock(name="Tracer")
    frisk_handle = MagicMock(name="Frisk")
    expected = _DummyResult(decision="allow", rules_matched_count=2, reason=None)
    frisk_handle.decide_tool_call.return_value = expected
    frisk_handle._adapter = adapter

    # Patch ToolCallSpan to avoid OpenTelemetry and to assert wiring.
    import frisk_sdk.core.frisk_session as session_module

    tool_span_cm = MagicMock(name="ToolCallSpanCM")
    tool_span_cm.__enter__.return_value = tool_span_cm
    tool_span_cm.__exit__.return_value = False
    tool_span_cm.span = object()

    tool_call_span_ctor = MagicMock(return_value=tool_span_cm)
    monkeypatch.setattr(session_module, "ToolCallSpan", tool_call_span_ctor)

    # Make tracing context injection deterministic.
    inject = MagicMock(name="inject")
    monkeypatch.setattr(session_module.TraceContextTextMapPropagator, "inject", inject)

    span_ctx = object()
    monkeypatch.setattr(
        session_module.trace, "set_span_in_context", lambda _span: span_ctx
    )

    s = FriskSession(frisk=frisk_handle, tracer=tracer)

    normalized_tool_call = NormalizedToolCall(
        id="toolcall-1", name="tool_a", args={"x": 1}
    )
    agent_state = {"user": {"id": 1}}

    res = s.decide_tool_call(tool_call=normalized_tool_call, agent_state=agent_state)

    assert res is expected

    tool_call_span_ctor.assert_called_once()
    frisk_handle.decide_tool_call.assert_called_once()

    called_args = frisk_handle.decide_tool_call.call_args.kwargs
    assert called_args["tool_call"] == NormalizedToolCall(
        id="toolcall-1", name="tool_a", args={"x": 1}
    )
    assert called_args["agent_state"] == agent_state

    # A carrier dict should be passed through and injected into.
    carrier = called_args["trace_context_carrier"]
    assert isinstance(carrier, dict)

    inject.assert_called_once()
    assert inject.call_args.kwargs["context"] is span_ctx

    tool_span_cm.save_result.assert_called_once_with(expected)


@pytest.mark.unit
def test_decide_tool_call_when_frisk_raises_saves_error_and_returns_error_result(
    monkeypatch: pytest.MonkeyPatch,
):
    adapter = _DummyAdapter()

    tracer = MagicMock(name="Tracer")
    frisk_handle = MagicMock(name="Frisk")
    frisk_exc = RuntimeError("kaboom")
    frisk_handle.decide_tool_call.side_effect = frisk_exc
    frisk_handle._adapter = adapter

    # Patch ToolCallSpan to avoid OpenTelemetry and to assert wiring.
    import frisk_sdk.core.frisk_session as session_module

    tool_span_cm = MagicMock(name="ToolCallSpanCM")
    tool_span_cm.__enter__.return_value = tool_span_cm
    tool_span_cm.__exit__.return_value = False
    tool_span_cm.span = object()

    tool_call_span_ctor = MagicMock(return_value=tool_span_cm)
    monkeypatch.setattr(session_module, "ToolCallSpan", tool_call_span_ctor)

    # Make tracing context injection deterministic.
    inject = MagicMock(name="inject")
    monkeypatch.setattr(session_module.TraceContextTextMapPropagator, "inject", inject)

    span_ctx = object()
    monkeypatch.setattr(
        session_module.trace, "set_span_in_context", lambda _span: span_ctx
    )

    s = FriskSession(frisk=frisk_handle, tracer=tracer)

    normalized_tool_call = NormalizedToolCall(
        id="toolcall-1", name="tool_a", args={"x": 1}
    )
    agent_state = {"user": {"id": 1}}

    res = s.decide_tool_call(tool_call=normalized_tool_call, agent_state=agent_state)

    # Returned value should be a ToolCallDecision in ERROR state.
    assert res.outcome is DecisionOutcome.ERROR
    assert res.rules_matched_count == 0
    assert "kaboom" in (res.reason or "")

    tool_call_span_ctor.assert_called_once()
    frisk_handle.decide_tool_call.assert_called_once()

    # Tool span should record the error wrapper (not the raw exception).
    tool_span_cm.save_error.assert_called_once()
    saved_error = tool_span_cm.save_error.call_args.args[0]
    assert "kaboom" in str(saved_error)

    # Ensure we did not incorrectly treat this as success.
    tool_span_cm.save_result.assert_not_called()
