"""Tests for BeaconLangGraphConfig.

Tests the BeaconLangGraphConfig class (checkpoint resolution, checkpoint_metadata
property, get_config() method, langgraph_config property) and the
create_langgraph_beacon_config convenience wrapper.
"""

from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest

from lumenova_beacon import BeaconClient
from lumenova_beacon.governance.config import GovernanceConfig
from lumenova_beacon.governance.integrations.langchain import BeaconLangGraphGovernanceHandler
from lumenova_beacon.tracing.integrations.langchain import (
    BeaconLangGraphHandler,
    BeaconLangGraphConfig,
)


@pytest.fixture(autouse=True)
def mock_beacon_client():
    """Patch get_client so BeaconLangGraphHandler and GovernanceHandler don't need a real client."""
    client = MagicMock(spec=BeaconClient)
    client.config = MagicMock()
    client.config.session_id = None
    client.config.environment = None
    client.config.user_id = None
    client.config.user_email = None
    client.config.user_name = None
    client._timestamp_override = None
    with patch(
        "lumenova_beacon.tracing.integrations.langchain.get_client",
        return_value=client,
    ), patch(
        "lumenova_beacon.core.client.get_client",
        return_value=client,
    ), patch(
        "lumenova_beacon.governance.client.get_transport",
        return_value=MagicMock(
            headers={"Authorization": "Bearer test-key"},
            timeout=10.0,
            verify=True,
        ),
    ), patch(
        "lumenova_beacon.governance.client.get_base_url",
        return_value="https://api.test.com",
    ):
        yield client


def _make_graph(state=None):
    """Create a mock compiled LangGraph graph with sync get_state."""
    graph = MagicMock()
    graph.get_state = MagicMock(return_value=state)
    return graph


def _make_state_snapshot(
    metadata: dict | None = None,
    next_nodes: tuple = (),
    interrupts: tuple = (),
):
    """Create a mock StateSnapshot-like object."""
    return SimpleNamespace(
        values={"messages": []},
        next=next_nodes,
        metadata=metadata,
        created_at=None,
        parent_config=None,
        tasks=(),
        interrupts=interrupts,
    )


# =========================================================================
# BeaconLangGraphConfig initialization / checkpoint resolution
# =========================================================================

class TestBeaconLangGraphConfigInit:
    """Tests for BeaconLangGraphConfig checkpoint resolution."""

    def test_fresh_thread_generates_new_ids(self):
        """Fresh thread (no checkpoint) generates new trace_id and root_span_id."""
        graph = _make_graph(state=None)

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-1",
            agent_name="test_agent",
        )

        assert beacon.trace_id is not None
        assert beacon.root_span_id is not None
        assert len(beacon.trace_id) == 32
        assert len(beacon.root_span_id) == 16
        assert beacon.is_resume is False
        assert beacon.handler._agent_name == "test_agent"
        graph.get_state.assert_called_once()

    def test_interrupted_thread_resumes_trace(self):
        """Interrupted thread (next non-empty) resumes with stored trace_id."""
        state = _make_state_snapshot(
            metadata={
                "beacon_trace_id": "a" * 32,
                "beacon_root_span_id": "b" * 16,
                "beacon_session_id": "session-123",
            },
            next_nodes=("plan",),
        )
        graph = _make_graph(state=state)

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-2",
            agent_name="test_agent",
        )

        assert beacon.is_resume is True
        assert beacon.trace_id == "a" * 32
        assert beacon.root_span_id == "b" * 16
        assert beacon.handler._session_id == "session-123"

    def test_interrupted_thread_with_interrupts_field(self):
        """Thread with state.interrupts also triggers resume."""
        interrupt = SimpleNamespace(value="Please confirm", id="int-1")
        state = _make_state_snapshot(
            metadata={
                "beacon_trace_id": "c" * 32,
                "beacon_root_span_id": "d" * 16,
            },
            next_nodes=(),
            interrupts=(interrupt,),
        )
        graph = _make_graph(state=state)

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-3",
            agent_name="test_agent",
        )

        assert beacon.is_resume is True
        assert beacon.trace_id == "c" * 32

    def test_completed_thread_starts_fresh_trace(self):
        """Completed thread (next empty, no interrupts) starts fresh trace."""
        state = _make_state_snapshot(
            metadata={
                "beacon_trace_id": "e" * 32,
                "beacon_root_span_id": "f" * 16,
                "beacon_session_id": "session-456",
            },
            next_nodes=(),
            interrupts=(),
        )
        graph = _make_graph(state=state)

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-4",
            agent_name="test_agent",
        )

        assert beacon.is_resume is False
        assert beacon.trace_id != "e" * 32
        assert len(beacon.trace_id) == 32
        # Session ID IS reused from previous checkpoint
        assert beacon.handler._session_id == "session-456"

    def test_session_id_reused_from_checkpoint(self):
        """Session ID from previous checkpoint is reused when not provided."""
        state = _make_state_snapshot(
            metadata={"beacon_session_id": "reused-session"},
        )
        graph = _make_graph(state=state)

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-5",
            agent_name="test_agent",
        )

        assert beacon.handler._session_id == "reused-session"

    def test_explicit_session_id_overrides_checkpoint(self):
        """Explicitly provided session_id overrides the one from checkpoint."""
        state = _make_state_snapshot(
            metadata={"beacon_session_id": "old-session"},
        )
        graph = _make_graph(state=state)

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-6",
            agent_name="test_agent",
            session_id="explicit-session",
        )

        assert beacon.handler._session_id == "explicit-session"

    def test_get_state_error_falls_back_to_fresh(self):
        """If get_state raises, falls back to fresh trace gracefully."""
        graph = MagicMock()
        graph.get_state = MagicMock(side_effect=RuntimeError("no checkpointer"))

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-7",
            agent_name="test_agent",
        )

        assert beacon.is_resume is False
        assert beacon.trace_id is not None
        assert len(beacon.trace_id) == 32

    def test_no_metadata_in_state(self):
        """State snapshot with metadata=None falls back to fresh trace."""
        state = _make_state_snapshot(metadata=None)
        graph = _make_graph(state=state)

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-8",
            agent_name="test_agent",
        )

        assert beacon.is_resume is False
        assert beacon.trace_id is not None

    def test_session_id_generated_when_not_available(self):
        """A new session_id UUID is generated when none exists."""
        graph = _make_graph(state=None)

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-9",
            agent_name="test_agent",
        )

        session_id = beacon.handler._session_id
        assert session_id is not None
        assert len(session_id) == 36
        assert session_id.count("-") == 4

    def test_interrupted_but_no_stored_trace_id_starts_fresh(self):
        """If graph is interrupted but no beacon_trace_id in metadata, start fresh."""
        state = _make_state_snapshot(
            metadata={"source": "loop", "step": 2},
            next_nodes=("execute",),
        )
        graph = _make_graph(state=state)

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-10",
            agent_name="test_agent",
        )

        assert beacon.is_resume is False
        assert beacon.trace_id is not None
        assert len(beacon.trace_id) == 32

    def test_environment_forwarded_to_handler(self):
        """Environment is forwarded to BeaconLangGraphHandler."""
        graph = _make_graph(state=None)

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-11",
            agent_name="test_agent",
            environment="production",
        )

        assert beacon.handler._environment == "production"

    def test_handler_is_beacon_callback_handler(self):
        """The handler property returns a BeaconLangGraphHandler."""
        graph = _make_graph(state=None)

        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-12",
            agent_name="test_agent",
        )

        assert isinstance(beacon.handler, BeaconLangGraphHandler)


# =========================================================================
# checkpoint_metadata property
# =========================================================================

class TestCheckpointMetadataProperty:
    """Tests for the checkpoint_metadata property."""

    def test_returns_all_beacon_keys(self):
        """Returns trace_id, root_span_id, and session_id."""
        graph = _make_graph(state=None)
        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-1",
            agent_name="test_agent",
        )

        meta = beacon.checkpoint_metadata

        assert "beacon_trace_id" in meta
        assert "beacon_root_span_id" in meta
        assert "beacon_session_id" in meta
        assert meta["beacon_trace_id"] == beacon.trace_id
        assert meta["beacon_root_span_id"] == beacon.root_span_id
        assert meta["beacon_session_id"] == beacon.handler._session_id

    def test_resume_ids_included(self):
        """On resume, returns the checkpoint-restored IDs."""
        state = _make_state_snapshot(
            metadata={
                "beacon_trace_id": "a" * 32,
                "beacon_root_span_id": "b" * 16,
                "beacon_session_id": "sess-1",
            },
            next_nodes=("plan",),
        )
        graph = _make_graph(state=state)
        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-2",
            agent_name="test_agent",
        )

        meta = beacon.checkpoint_metadata

        assert meta["beacon_trace_id"] == "a" * 32
        assert meta["beacon_root_span_id"] == "b" * 16
        assert meta["beacon_session_id"] == "sess-1"


# =========================================================================
# get_config() method
# =========================================================================

class TestGetConfig:
    """Tests for the get_config() method."""

    def test_basic_config_structure(self):
        """Config has configurable, metadata, and callbacks."""
        graph = _make_graph(state=None)
        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-1",
            agent_name="test_agent",
        )

        config = beacon.get_config()

        assert config["configurable"]["thread_id"] == "thread-1"
        assert "beacon_trace_id" in config["metadata"]
        assert "beacon_root_span_id" in config["metadata"]
        assert "beacon_session_id" in config["metadata"]
        assert config["callbacks"] == [beacon.handler]
        # Beacon keys NOT in configurable (only in metadata)
        assert "beacon_trace_id" not in config["configurable"]

    def test_custom_configurable_merged(self):
        """User-provided configurable keys are merged with thread_id."""
        graph = _make_graph(state=None)
        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-1",
            agent_name="test_agent",
        )

        config = beacon.get_config(configurable={"recursion_limit": 50})

        assert config["configurable"]["thread_id"] == "thread-1"
        assert config["configurable"]["recursion_limit"] == 50

    def test_custom_metadata_merged(self):
        """User-provided metadata is merged with beacon keys."""
        graph = _make_graph(state=None)
        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-1",
            agent_name="test_agent",
        )

        config = beacon.get_config(metadata={"custom_key": "value"})

        assert config["metadata"]["custom_key"] == "value"
        assert "beacon_trace_id" in config["metadata"]

    def test_tags_added(self):
        """Tags are added to config."""
        graph = _make_graph(state=None)
        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-1",
            agent_name="test_agent",
        )

        config = beacon.get_config(tags=["production", "v2"])

        assert config["tags"] == ["production", "v2"]

    def test_extra_kwargs_added(self):
        """Extra kwargs are included in config."""
        graph = _make_graph(state=None)
        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-1",
            agent_name="test_agent",
        )

        config = beacon.get_config(run_name="my_run")

        assert config["run_name"] == "my_run"

    def test_resume_ids_in_metadata(self):
        """On resume, config metadata contains the restored IDs."""
        state = _make_state_snapshot(
            metadata={
                "beacon_trace_id": "a" * 32,
                "beacon_root_span_id": "b" * 16,
                "beacon_session_id": "session-x",
            },
            next_nodes=("plan",),
        )
        graph = _make_graph(state=state)
        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-2",
            agent_name="test_agent",
        )

        config = beacon.get_config()

        assert config["metadata"]["beacon_trace_id"] == "a" * 32
        assert config["metadata"]["beacon_root_span_id"] == "b" * 16
        assert config["metadata"]["beacon_session_id"] == "session-x"


# =========================================================================
# langgraph_config property (delegates to get_config)
# =========================================================================

class TestLanggraphConfigProperty:
    """Tests for the langgraph_config property."""

    def test_config_contains_beacon_keys_in_metadata(self):
        """Config has beacon IDs in metadata, not configurable."""
        graph = _make_graph(state=None)
        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-1",
            agent_name="test_agent",
        )

        config = beacon.langgraph_config

        assert config["configurable"]["thread_id"] == "thread-1"
        assert config["metadata"]["beacon_trace_id"] == beacon.trace_id
        assert config["metadata"]["beacon_root_span_id"] == beacon.root_span_id
        assert config["metadata"]["beacon_session_id"] == beacon.handler._session_id
        assert config["callbacks"] == [beacon.handler]

    def test_config_preserves_resume_ids(self):
        """On resume, config contains the same trace_id from checkpoint."""
        state = _make_state_snapshot(
            metadata={
                "beacon_trace_id": "a" * 32,
                "beacon_root_span_id": "b" * 16,
                "beacon_session_id": "session-x",
            },
            next_nodes=("plan",),
        )
        graph = _make_graph(state=state)
        beacon = BeaconLangGraphConfig(
            graph=graph,
            thread_id="thread-2",
            agent_name="test_agent",
        )

        config = beacon.langgraph_config

        assert config["metadata"]["beacon_trace_id"] == "a" * 32
        assert config["metadata"]["beacon_root_span_id"] == "b" * 16
        assert config["metadata"]["beacon_session_id"] == "session-x"


# =========================================================================
# GovernanceConfig integration
# =========================================================================


class TestGovernanceConfigIntegration:
    """Tests for GovernanceConfig integration with BeaconLangGraphConfig."""

    def test_no_governance_by_default(self):
        """Without governance param, governance_handler is None."""
        graph = _make_graph(state=None)
        beacon = BeaconLangGraphConfig(
            graph=graph, thread_id="t-1", agent_name="agent",
        )

        assert beacon.governance_handler is None

    def test_governance_creates_handler(self):
        """Passing GovernanceConfig creates a BeaconLangGraphGovernanceHandler."""
        graph = _make_graph(state=None)
        gov = GovernanceConfig(stack_ids=["stack-1"], fail_open=False)

        beacon = BeaconLangGraphConfig(
            graph=graph, thread_id="t-2", agent_name="agent",
            governance=gov,
        )

        assert beacon.governance_handler is not None
        assert isinstance(beacon.governance_handler, BeaconLangGraphGovernanceHandler)

    def test_governance_handler_in_callbacks(self):
        """get_config() includes the governance handler in callbacks."""
        graph = _make_graph(state=None)
        gov = GovernanceConfig(stack_ids=["stack-1"])

        beacon = BeaconLangGraphConfig(
            graph=graph, thread_id="t-3", agent_name="agent",
            governance=gov,
        )
        config = beacon.get_config()

        assert len(config["callbacks"]) == 2
        assert config["callbacks"][0] is beacon.handler
        assert config["callbacks"][1] is beacon.governance_handler

    def test_no_governance_callbacks_single(self):
        """Without governance, callbacks list only has the tracing handler."""
        graph = _make_graph(state=None)
        beacon = BeaconLangGraphConfig(
            graph=graph, thread_id="t-4", agent_name="agent",
        )
        config = beacon.get_config()

        assert len(config["callbacks"]) == 1
        assert config["callbacks"][0] is beacon.handler

    def test_shared_params_propagated(self):
        """Shared params (agent_name, environment, metadata) propagated to governance handler."""
        graph = _make_graph(state=None)
        gov = GovernanceConfig(stack_ids=["stack-1"])

        beacon = BeaconLangGraphConfig(
            graph=graph, thread_id="t-5",
            agent_name="my-agent",
            environment="production",
            metadata={"key": "val"},
            governance=gov,
        )

        gh = beacon.governance_handler
        assert gh._agent_name == "my-agent"
        assert gh._environment == "production"
        assert gh._metadata == {"key": "val"}

    def test_session_id_shared(self):
        """Auto-generated session_id is shared between tracing and governance handlers."""
        graph = _make_graph(state=None)
        gov = GovernanceConfig(stack_ids=["stack-1"])

        beacon = BeaconLangGraphConfig(
            graph=graph, thread_id="t-6", agent_name="agent",
            governance=gov,
        )

        assert beacon.handler._session_id == beacon.governance_handler._session_id

    def test_explicit_session_id_shared(self):
        """Explicit session_id is shared between both handlers."""
        graph = _make_graph(state=None)
        gov = GovernanceConfig(stack_ids=["stack-1"])

        beacon = BeaconLangGraphConfig(
            graph=graph, thread_id="t-7", agent_name="agent",
            session_id="explicit-sess",
            governance=gov,
        )

        assert beacon.handler._session_id == "explicit-sess"
        assert beacon.governance_handler._session_id == "explicit-sess"

    def test_governance_specific_params_forwarded(self):
        """Governance-specific params are forwarded to the handler."""
        graph = _make_graph(state=None)
        gov = GovernanceConfig(
            stack_ids=["stack-1"],
            tags=["tag-a"],
            fail_open=False,
            enabled_hooks={"pre_tool", "post_tool"},
            timeout=5.0,
            debug=False,
            on_violation="stop",
            max_violations=3,
            agent_id="agent-id-1",
        )

        beacon = BeaconLangGraphConfig(
            graph=graph, thread_id="t-8", agent_name="agent",
            governance=gov,
        )

        gh = beacon.governance_handler
        assert gh._stack_ids == ["stack-1"]
        assert gh._tags == ["tag-a"]
        assert gh._fail_open is False
        assert gh._enabled_hooks == {"pre_tool", "post_tool"}
        assert gh._timeout == 5.0
        assert gh._on_violation == "stop"
        assert gh._max_violations == 3
        assert gh._agent_id == "agent-id-1"

    def test_checkpoint_session_id_synced_to_governance(self):
        """Checkpoint-restored session_id is synced to governance handler."""
        state = _make_state_snapshot(
            metadata={"beacon_session_id": "checkpoint-session"},
        )
        graph = _make_graph(state=state)
        gov = GovernanceConfig(stack_ids=["stack-1"])

        beacon = BeaconLangGraphConfig(
            graph=graph, thread_id="t-9", agent_name="agent",
            governance=gov,
        )
        # Trigger checkpoint resolution
        beacon.get_config()

        assert beacon.handler._session_id == "checkpoint-session"
        assert beacon.governance_handler._session_id == "checkpoint-session"

    def test_explicit_session_id_not_overridden_by_checkpoint(self):
        """Explicit session_id is NOT overridden by checkpoint for either handler."""
        state = _make_state_snapshot(
            metadata={"beacon_session_id": "old-session"},
        )
        graph = _make_graph(state=state)
        gov = GovernanceConfig(stack_ids=["stack-1"])

        beacon = BeaconLangGraphConfig(
            graph=graph, thread_id="t-10", agent_name="agent",
            session_id="user-session",
            governance=gov,
        )
        beacon.get_config()

        assert beacon.handler._session_id == "user-session"
        assert beacon.governance_handler._session_id == "user-session"


# =========================================================================
# GovernanceConfig dataclass
# =========================================================================


class TestGovernanceConfigDefaults:
    """Tests for GovernanceConfig default values."""

    def test_defaults(self):
        gov = GovernanceConfig()
        assert gov.stack_ids is None
        assert gov.tags is None
        assert gov.fail_open is True
        assert gov.enabled_hooks is None
        assert gov.result_parser is None
        assert gov.timeout is None
        assert gov.debug is False
        assert gov.on_violation == "raise"
        assert gov.max_violations is None
        assert gov.agent_id is None

    def test_custom_values(self):
        parser = lambda r: True
        gov = GovernanceConfig(
            stack_ids=["s1", "s2"],
            tags=["t1"],
            fail_open=False,
            enabled_hooks={"pre_tool"},
            result_parser=parser,
            timeout=3.0,
            debug=True,
            on_violation="stop",
            max_violations=5,
            agent_id="a1",
        )
        assert gov.stack_ids == ["s1", "s2"]
        assert gov.tags == ["t1"]
        assert gov.fail_open is False
        assert gov.enabled_hooks == {"pre_tool"}
        assert gov.result_parser is parser
        assert gov.timeout == 3.0
        assert gov.debug is True
        assert gov.on_violation == "stop"
        assert gov.max_violations == 5
        assert gov.agent_id == "a1"
