"""Tests for AgentX governance integration hooks."""
from types import SimpleNamespace
from unittest.mock import Mock

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from agendex.integrations.agentx import (  # noqa: E402
    AgentXGovernanceConfig,
    AgentXGovernanceHooks,
    GovernedAgentBuilder,
    enable_agentx_governance,
)


def _mock_client(mode: str = "shadow") -> Mock:
    client = Mock()
    client.agent_id = "agentx_test"
    client.mode = mode
    client.ingest_events = Mock(return_value={"ingested": 1})
    return client


def test_enable_agentx_governance_returns_hooks():
    client = _mock_client()
    hooks = enable_agentx_governance(client)
    assert isinstance(hooks, AgentXGovernanceHooks)


def test_wrap_tool_shadow_executes_local_tool_and_emits_events():
    client = _mock_client(mode="shadow")

    def invoke_side_effect(**kwargs):
        on_event = kwargs.get("on_event")
        if on_event:
            on_event({"type": "tool_call", "action": kwargs["action"], "task": kwargs["task"]})
            on_event({"type": "decision", "outcome": "would_allow", "reason": "shadow"})
            on_event({"type": "invoke_result", "result": {"ok": True}, "mode": "shadow"})
        return {"_agendex": "shadow"}

    client.invoke = Mock(side_effect=invoke_side_effect)

    hooks = AgentXGovernanceHooks(client, config=AgentXGovernanceConfig())
    state = {
        "user_query": "Compute sum",
        "manager_reasoning": "Need a calculator",
        "last_assistant_routed": "math_assistant",
        "user_id": "user-123",
        "chat_id": "chat-xyz",
    }

    wrapped = hooks.wrap_tool(
        tool_name="add_numbers",
        tool_func=lambda x, y: x + y,
        state_getter=lambda: state,
    )

    result = wrapped(2, 5)

    assert result == 7
    assert client.invoke.call_count == 1
    invoke_kwargs = client.invoke.call_args.kwargs
    assert invoke_kwargs["action"] == "tool.add_numbers"
    assert invoke_kwargs["trace_id"] == state["trace_id"]
    assert invoke_kwargs["context"]["user_prompt"] == "Compute sum"
    assert invoke_kwargs["context"]["manager_reasoning"] == "Need a calculator"
    assert invoke_kwargs["context"]["user_id"] == "user-123"
    assert invoke_kwargs["context"]["session_id"] == "chat-xyz"
    assert client.ingest_events.call_count >= 3
    emitted_events = [call.args[0][0] for call in client.ingest_events.call_args_list if call.args and call.args[0]]
    tool_local = next(evt for evt in emitted_events if evt.get("event_type") == "tool_result_local")
    assert isinstance(tool_local.get("latency_ms"), float)


def test_manager_and_assistant_hooks_emit_trace_events():
    client = _mock_client()
    client.invoke = Mock(return_value={"_agendex": "shadow"})

    hooks = AgentXGovernanceHooks(client, config=AgentXGovernanceConfig())
    state = {
        "user_query": "Find top customers",
        "manager_reasoning": "Route to sql assistant",
        "manager_routing_decision": "sql_assistant",
    }

    manager_decision = SimpleNamespace(
        reasoning="Need SQL lookup",
        next_assistant="sql_assistant",
        required_tool_names=["retrieve_exact_content"],
    )
    hooks.on_manager_decision(
        state=state,
        llm_input_messages=[SimpleNamespace(content="manager input")],
        decision=manager_decision,
    )

    assistant_response = SimpleNamespace(
        content="I should run execute_sql_queries",
        tool_calls=[{"name": "execute_sql_queries"}],
    )
    hooks.on_assistant_llm_response(
        assistant_name="sql_assistant",
        state=state,
        llm_input_messages=[SimpleNamespace(content="assistant input")],
        response=assistant_response,
    )

    assert client.ingest_events.call_count >= 2
    event_types = [
        call.args[0][0]["event_type"]
        for call in client.ingest_events.call_args_list
        if call.args and call.args[0]
    ]
    assert "manager_decision" in event_types
    assert "assistant_llm_result" in event_types


def test_governed_agent_builder_passes_hooks_to_supported_builder():
    client = _mock_client()
    client.invoke = Mock(return_value={"_agendex": "shadow"})

    class FakeBuilder:
        def __init__(self, *, governance_hooks=None):
            self.governance_hooks = governance_hooks
            self.manager = SimpleNamespace(_governance_hooks=None)
            self.assistants = [SimpleNamespace(_governance_hooks=None)]

        def build_agent(self):
            return "graph-ready"

    wrapped = GovernedAgentBuilder(
        agendex=client,
        agent_builder_cls=FakeBuilder,
        install_runtime_shims_flag=False,
    )

    assert isinstance(wrapped.hooks, AgentXGovernanceHooks)
    assert wrapped.builder.governance_hooks is wrapped.hooks
    assert wrapped.build_agent() == "graph-ready"


def test_governed_agent_builder_falls_back_for_legacy_builder():
    client = _mock_client()
    client.invoke = Mock(return_value={"_agendex": "shadow"})

    class LegacyBuilder:
        def __init__(self):
            self.manager = SimpleNamespace(_governance_hooks=None)
            self.assistants = [SimpleNamespace(_governance_hooks=None)]

        def build_agent(self):
            return {"ok": True}

    wrapped = GovernedAgentBuilder(
        agendex=client,
        agent_builder_cls=LegacyBuilder,
        install_runtime_shims_flag=False,
    )

    assert wrapped.builder.governance_hooks is wrapped.hooks
    assert wrapped.builder.manager._governance_hooks is wrapped.hooks
    assert wrapped.builder.assistants[0]._governance_hooks is wrapped.hooks
    assert wrapped.build_agent() == {"ok": True}
