import os
import time
import base64
import json
import pytest
from unittest.mock import MagicMock
from typing import Any

from frisk_sdk import DecisionOutcome, NormalizedToolCall
from frisk_sdk.core.frisk import Frisk


# Helpers to produce a minimal valid JWT with future exp
def make_jwt(exp_epoch_s: int | None = None) -> str:
    header = {"alg": "none", "typ": "JWT"}
    if exp_epoch_s is None:
        exp_epoch_s = int(time.time()) + 3600
    payload = {"exp": exp_epoch_s}

    def b64url(d: dict) -> str:
        raw = json.dumps(d, separators=(",", ":")).encode("utf-8")
        s = base64.urlsafe_b64encode(raw).decode("utf-8")
        return s.rstrip("=")

    return f"{b64url(header)}.{b64url(payload)}."  # empty signature is fine for decode-only


def _mock_tracing_manager(monkeypatch, frisk_module):
    """Replace TracingManager with a cheap in-memory fake.

    Frisk constructs TracingManager and later calls:
    - update_auth_token(token)
    - shutdown()

    The real TracingManager wires OTLP/gRPC exporter which we don't want in unit tests.
    """

    tracing_manager = MagicMock(name="TracingManager")
    tracer_obj = object()
    tracing_manager.get_tracer.return_value = tracer_obj

    def _update_auth_token(token: str) -> None:
        tracing_manager.auth_token = token

    tracing_manager.update_auth_token.side_effect = _update_auth_token

    class _FakeTracingManager:
        def __init__(self, *args, **kwargs):
            self._kwargs = kwargs

        def get_tracer(self):
            return tracing_manager.get_tracer()

        def update_auth_token(self, token: str) -> None:
            return tracing_manager.update_auth_token(token)

        def shutdown(self) -> None:
            return tracing_manager.shutdown()

    monkeypatch.setattr(frisk_module, "TracingManager", _FakeTracingManager)
    return tracing_manager


@pytest.fixture(autouse=True)
def suppress_rust_warn_logs(monkeypatch):
    # Suppress tracing::warn from Rust by raising the log level
    monkeypatch.setenv("RUST_LOG", "error")
    yield


@pytest.fixture(autouse=True)
def mock_access_token_provider(monkeypatch):
    # Ensure issuer URL presence doesn't trigger real network
    os.environ.setdefault("FRISK_TOKEN_ISSUER_URL", "http://example.local/issuer")

    # Monkeypatch AccessTokenProvider.fetch_access_token to avoid HTTP calls
    from frisk_sdk.core.access_token_provider import AccessTokenProvider

    def fake_fetch(self):
        return make_jwt()

    monkeypatch.setattr(AccessTokenProvider, "fetch_access_token", fake_fetch)
    yield


@pytest.fixture
def frisk_instance():
    # Create Frisk; AccessTokenProvider is already mocked to return a valid JWT
    f = Frisk(api_key="test_api_key", api_base_url="https://example.frisk.com")
    yield f
    # Stop background refresher
    f.shutdown()


@pytest.mark.integration
def test_decide_tool_call_allow_when_no_policies(frisk_instance):
    # With no policies loaded yet, engine should allow by default
    res = frisk_instance.decide_tool_call(
        tool_call=NormalizedToolCall(id="call-1", name="unknown_tool", args={"a": 1}),
        agent_state={"user": {"id": 123}},
        trace_context_carrier=None,
    )

    assert res.outcome == DecisionOutcome.ALLOW
    assert res.rules_matched_count == 0
    assert res.reason == "No active policies for tool"


@pytest.mark.integration
def test_decide_tool_call_with_trace_context(frisk_instance):
    # Providing a trace context carrier should be accepted by rust side
    carrier = {}
    res = frisk_instance.decide_tool_call(
        tool_call=NormalizedToolCall(id="call-2", name="another_tool", args={"x": 42}),
        agent_state={"user": {"isAdmin": False}},
        trace_context_carrier=carrier,
    )
    assert res.outcome in {
        DecisionOutcome.ALLOW,
        DecisionOutcome.DENY,
    }  # depending on policies; with none it's allow
    assert isinstance(res.rules_matched_count, int)


@pytest.mark.integration
def test_decide_tool_call_none_args_state_raises(frisk_instance):
    # Passing None for tool_args/agent_state should be accepted and treated as nulls by the core.
    res = frisk_instance.decide_tool_call(
        tool_call=NormalizedToolCall(id="call-3", name="tool_x", args=None),
        agent_state=None,
        trace_context_carrier={},
    )
    assert res.outcome in {DecisionOutcome.ALLOW, DecisionOutcome.DENY}
    assert isinstance(res.rules_matched_count, int)
    # reason may be None or a string depending on policies
    assert getattr(res, "reason", None) is None or isinstance(res.reason, str)


@pytest.mark.integration
def test_update_access_token_delegates_to_core_handle(monkeypatch):
    """Boundary/unit: verify Python forwards the token to the underlying core handle."""

    from frisk_sdk.core import frisk as frisk_module

    # Avoid wiring real OpenTelemetry providers (Frisk now uses TracingManager)
    _tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    # Fake shared token runtime
    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    # Mock the Rust/PyO3 handle and capture initialization.
    core_handle = MagicMock(name="FriskCoreHandle")

    def _handle_factory(current_token: str, base_url: str, otlp_endpoint: str):
        assert current_token == make_jwt()
        return core_handle

    monkeypatch.setattr(frisk_module, "FriskCoreHandle", _handle_factory)

    f = Frisk(api_key="test_api_key", api_base_url="https://example.frisk.com")
    try:
        f.update_access_token("new-token")
        core_handle.update_auth_token.assert_called_once_with("new-token")
    finally:
        f.shutdown()


@pytest.mark.integration
def test_update_access_token_accepts_empty_string(monkeypatch):
    """Boundary/unit: empty string should be forwarded unchanged (core decides validity)."""

    from frisk_sdk.core import frisk as frisk_module

    _tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    core_handle = MagicMock(name="FriskCoreHandle")
    monkeypatch.setattr(frisk_module, "FriskCoreHandle", lambda *_a, **_kw: core_handle)

    f = Frisk(api_key="test_api_key", api_base_url="https://example.frisk.com")
    try:
        f.update_access_token("")
        core_handle.update_auth_token.assert_called_once_with("")
    finally:
        f.shutdown()


@pytest.mark.integration
def test_shutdown_unsubscribes_and_shuts_down_handle(monkeypatch):
    """Boundary/unit: shutdown should unsubscribe callbacks and shut down the core handle."""

    from frisk_sdk.core import frisk as frisk_module

    tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    core_handle = MagicMock(name="FriskCoreHandle")
    monkeypatch.setattr(frisk_module, "FriskCoreHandle", lambda *_a, **_kw: core_handle)

    f = Frisk(api_key="test_api_key", api_base_url="https://example.frisk.com")

    # Act
    f.shutdown()

    # Assert: 2 callbacks registered and later unsubscribed
    assert token_runtime.subscribe.call_count == 2
    assert token_runtime.unsubscribe.call_count == 2
    core_handle.shutdown.assert_called_once_with()

    # and tracer handle should also be shut down
    tracing_manager.shutdown.assert_called_once_with()


@pytest.mark.integration
def test_shutdown_is_idempotent(monkeypatch):
    """Boundary/unit: shutdown should be safe to call twice."""

    from frisk_sdk.core import frisk as frisk_module

    tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    core_handle = MagicMock(name="FriskCoreHandle")
    monkeypatch.setattr(frisk_module, "FriskCoreHandle", lambda *_a, **_kw: core_handle)

    f = Frisk(api_key="test_api_key", api_base_url="https://example.frisk.com")

    f.shutdown()
    f.shutdown()

    # First shutdown unsubscribes two callbacks; second shutdown tries again but shouldn't error.
    assert token_runtime.unsubscribe.call_count >= 2
    assert core_handle.shutdown.call_count == 2
    assert tracing_manager.shutdown.call_count == 2


@pytest.mark.integration
def test_shutdown_propagates_handle_shutdown_errors(monkeypatch):
    """Boundary/unit: errors from the core handle shutdown should bubble up."""

    from frisk_sdk.core import frisk as frisk_module

    _tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    core_handle = MagicMock(name="FriskCoreHandle")
    core_handle.shutdown.side_effect = RuntimeError("shutdown failed")
    monkeypatch.setattr(frisk_module, "FriskCoreHandle", lambda *_a, **_kw: core_handle)

    f = Frisk(api_key="test_api_key", api_base_url="https://example.frisk.com")

    with pytest.raises(RuntimeError, match="shutdown failed"):
        f.shutdown()


@pytest.mark.unit
def test_create_session_returns_uuid_and_registers_in_registry(monkeypatch):
    """Boundary/unit: frisk.session() should return a UUID and register the session in the registry."""

    from frisk_sdk.core import frisk as frisk_module
    from frisk_sdk.core.session_registry import SessionRegistry

    _tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    # Make tracer deterministic and easy to compare.
    tracer = object()
    _tracing_manager.get_tracer.return_value = tracer

    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    # Core handle stub
    core_handle = MagicMock(name="FriskCoreHandle")
    monkeypatch.setattr(frisk_module, "FriskCoreHandle", lambda *_a, **_kw: core_handle)

    f = Frisk(api_key="test_api_key", api_base_url="https://example.frisk.com")
    try:
        # Clear registry before test
        registry = SessionRegistry()
        registry.clear()

        session_created = f.session()

        from frisk_sdk.core.frisk_session import FriskSession

        assert isinstance(session_created, FriskSession)

        session_retrieved = registry.get(session_created.id)
        assert session_retrieved.frisk is f
        assert session_retrieved.tracer is tracer
        # Default path should use BaseFrameworkAdapter
        from frisk_sdk.framework_adapter.base_framework_adapter import (
            BaseFrameworkAdapter,
        )

        assert isinstance(session_retrieved.adapter, BaseFrameworkAdapter)
    finally:
        registry = SessionRegistry()
        registry.clear()
        f.shutdown()


class _DummyAdapter:
    def __init__(self, tool_call=None, *, raise_exc: Exception | None = None):
        self._tool_call = tool_call
        self._raise_exc = raise_exc
        self.calls: list[object] = []

    def serialize_tool_args(self, tool_args):
        return json.dumps(tool_args)

    def serialize_agent_state(self, agent_state):
        return json.dumps(agent_state)

    def normalize_tool_call(self, tool_call):
        self.calls.append(tool_call)
        if self._raise_exc is not None:
            raise self._raise_exc
        return self._tool_call


@pytest.mark.unit
def test_create_session_uses_provided_framework_adapter(monkeypatch):
    """Boundary/unit: a provided adapter should be used by the created session."""

    from frisk_sdk.core import frisk as frisk_module
    from frisk_sdk.core.session_registry import SessionRegistry

    _tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    tracer = object()
    _tracing_manager.get_tracer.return_value = tracer

    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    core_handle = MagicMock(name="FriskCoreHandle")
    monkeypatch.setattr(frisk_module, "FriskCoreHandle", lambda *_a, **_kw: core_handle)

    from frisk_sdk.framework_adapter.base_framework_adapter import BaseFrameworkAdapter

    class _CustomAdapter(BaseFrameworkAdapter):
        def __init__(self):
            super().__init__()
            self.calls: list[object] = []

        def normalize_tool_call(self, tool_call):
            self.calls.append(tool_call)
            raise NotImplementedError

    custom_adapter = _CustomAdapter()

    f = Frisk(
        api_key="test_api_key",
        adapter=custom_adapter,
        api_base_url="https://example.frisk.com",
    )
    try:
        # Clear registry before test
        registry = SessionRegistry()
        registry.clear()

        session = f.session()

        assert session.adapter is custom_adapter
        assert session.frisk is f
        assert session.tracer is tracer
    finally:
        f.shutdown()


@pytest.mark.unit
def test_normalize_tool_call_delegates_to_framework_adapter(monkeypatch):
    """Boundary/unit: normalize_tool_call should delegate to the configured adapter."""

    from frisk_sdk.core import frisk as frisk_module

    _tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    core_handle = MagicMock(name="FriskCoreHandle")
    monkeypatch.setattr(frisk_module, "FriskCoreHandle", lambda *_a, **_kw: core_handle)

    from frisk_sdk.framework_adapter.framework_adapter import NormalizedToolCall
    from frisk_sdk.framework_adapter.base_framework_adapter import BaseFrameworkAdapter

    expected = NormalizedToolCall(id="id-1", name="tool", args={"x": 1})

    class _Adapter(BaseFrameworkAdapter):
        def __init__(self):
            super().__init__()
            self.calls: list[object] = []

        def normalize_tool_call(self, tool_call):
            self.calls.append(tool_call)
            return expected

    adapter = _Adapter()

    tool_call = object()

    f = Frisk(
        api_key="test_api_key",
        adapter=adapter,
        api_base_url="https://example.frisk.com",
    )
    try:
        res = f.normalize_tool_call(tool_call)

        assert res == expected
        assert adapter.calls == [tool_call]
    finally:
        f.shutdown()


@pytest.mark.unit
def test_normalize_tool_call_propagates_adapter_errors(monkeypatch):
    """Boundary/unit: adapter failures should bubble up unchanged."""

    from frisk_sdk.core import frisk as frisk_module

    _tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    core_handle = MagicMock(name="FriskCoreHandle")
    monkeypatch.setattr(frisk_module, "FriskCoreHandle", lambda *_a, **_kw: core_handle)

    from frisk_sdk.framework_adapter.base_framework_adapter import BaseFrameworkAdapter

    class _Adapter(BaseFrameworkAdapter):
        def __init__(self):
            super().__init__()
            self.calls: list[object] = []

        def normalize_tool_call(self, tool_call):
            self.calls.append(tool_call)
            raise ValueError("bad tool call")

    adapter = _Adapter()
    tool_call = object()

    f = Frisk(
        api_key="test_api_key",
        adapter=adapter,
        api_base_url="https://example.frisk.com",
    )
    try:
        with pytest.raises(ValueError, match="bad tool call"):
            f.normalize_tool_call(tool_call)
        assert adapter.calls == [tool_call]
    finally:
        f.shutdown()


def _dummy_core_process_result(
    *, decision: str = "allow", rules_matched_count: int = 0, reason=None
):
    """Create a minimal object that mimics the PyO3 ProcessToolCallResultPy binding.

    `ToolCallDecision.from_rust_binding()` only relies on these three attributes.
    """

    class _Result:
        def __init__(self):
            self.decision = decision
            self.rules_matched_count = rules_matched_count
            self.reason = reason

    return _Result()


@pytest.mark.unit
def test_decide_tool_call_forwards_options_dict(monkeypatch):
    """Dict redaction should be merged with defaults and forwarded to rust."""

    from frisk_sdk.core import frisk as frisk_module

    _tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    core_handle = MagicMock(name="FriskCoreHandle")
    core_handle.process.return_value = _dummy_core_process_result()
    monkeypatch.setattr(frisk_module, "FriskCoreHandle", lambda *_a, **_kw: core_handle)

    raw_options: Any = {"redact_tool_args": True}

    f = Frisk(
        api_key="test_api_key",
        api_base_url="https://example.frisk.com",
        redact=raw_options,
    )
    try:
        f.decide_tool_call(
            tool_call=NormalizedToolCall(id="call-opt-2", name="tool_a", args={"a": 1}),
            agent_state={"user": {"id": 1}},
            trace_context_carrier={},
        )

        called = core_handle.process.call_args.args
        assert called[4] == {"redact_tool_args": True, "redact_agent_state": None}
    finally:
        f.shutdown()


@pytest.mark.unit
def test_decide_tool_call_keeps_unknown_option_keys(monkeypatch):
    """normalize_frisk_options treats mapping-like input as a dict and passes unknown keys through.

    This test locks in that behavior so adding strict validation later is a deliberate change.
    """

    from frisk_sdk.core import frisk as frisk_module

    _tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    core_handle = MagicMock(name="FriskCoreHandle")
    core_handle.process.return_value = _dummy_core_process_result()
    monkeypatch.setattr(frisk_module, "FriskCoreHandle", lambda *_a, **_kw: core_handle)

    raw_options: Any = {"redact_agent_state": True, "some_future_flag": "on"}

    f = Frisk(
        api_key="test_api_key",
        api_base_url="https://example.frisk.com",
        redact=raw_options,
    )
    try:
        f.decide_tool_call(
            tool_call=NormalizedToolCall(id="call-opt-4", name="tool_a", args={"a": 1}),
            agent_state={"user": {"id": 1}},
            trace_context_carrier={},
        )

        called = core_handle.process.call_args.args
        assert called[4]["redact_tool_args"] is None
        assert called[4]["redact_agent_state"] is True
        assert called[4]["some_future_flag"] == "on"
    finally:
        f.shutdown()


@pytest.mark.unit
def test_decide_tool_call_rejects_non_mapping_options(monkeypatch):
    """A non-dataclass, non-mapping redaction value should fail normalization at construction time."""

    from frisk_sdk.core import frisk as frisk_module

    _tracing_manager = _mock_tracing_manager(monkeypatch, frisk_module)

    token_runtime = MagicMock(name="AccessTokenProviderRuntime")
    token_runtime.get_access_token.return_value = make_jwt()
    token_runtime.subscribe.return_value = None
    token_runtime.unsubscribe.return_value = None
    monkeypatch.setattr(
        frisk_module,
        "get_or_init_access_token_provider_runtime",
        lambda **_kw: token_runtime,
    )

    raw_bad_options: Any = 123

    # No need to mock FriskCoreHandle; __init__ should fail before creating it.
    with pytest.raises(TypeError):
        Frisk(
            api_key="test_api_key",
            api_base_url="https://example.frisk.com",
            redact=raw_bad_options,
        )
