"""Unit tests for the LangGraph pentesting agent."""

from __future__ import annotations

import json
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from ziran.domain.entities.pentest import (
    PentestPlan,
    PentestSession,
    PentestStatus,
)

# ──────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────


def _mock_llm_response(content: str = '{"tool":"run_scan","args":{}}') -> MagicMock:
    resp = MagicMock()
    resp.content = content
    resp.prompt_tokens = 10
    resp.completion_tokens = 5
    resp.total_tokens = 15
    return resp


def _mock_llm_client(
    responses: list[str] | None = None,
) -> MagicMock:
    """Create a mock LLM client that returns configured responses in sequence."""
    client = MagicMock()
    default_responses = responses or ['{"tool":"run_scan","args":{}}']
    side_effects = [_mock_llm_response(r) for r in default_responses]
    client.complete = AsyncMock(side_effect=side_effects)
    return client


def _mock_scanner() -> MagicMock:
    scanner = MagicMock()
    scanner.graph = MagicMock()
    scanner.graph.node_count = 5
    scanner.graph.edge_count = 3
    scanner.graph.graph = MagicMock()
    scanner.graph.graph.nodes.return_value = [
        ("tool_search", {"node_type": "capability"}),
    ]
    scanner._attack_results = []
    return scanner


# ──────────────────────────────────────────────────────────────────────
# PentestAgent._parse_action
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestParseAction:
    """Tests for LLM response parsing into tool actions."""

    @pytest.fixture
    def agent(self) -> Any:
        from ziran.application.pentesting.agent import PentestAgent

        return PentestAgent(
            scanner=_mock_scanner(),
            llm_client=_mock_llm_client(),
            max_iterations=5,
        )

    def test_parse_json_block(self, agent: Any) -> None:
        response = '```json\n{"tool": "run_phase", "args": {"phase": "reconnaissance"}}\n```'
        action = agent._parse_action(response)
        assert action["tool"] == "run_phase"
        assert action["args"]["phase"] == "reconnaissance"

    def test_parse_inline_json(self, agent: Any) -> None:
        response = 'I suggest we use {"tool": "analyze_graph", "args": {}} to inspect.'
        action = agent._parse_action(response)
        assert action["tool"] == "analyze_graph"

    def test_parse_heuristic_scan(self, agent: Any) -> None:
        response = "We should run a full scan campaign to discover vulnerabilities."
        action = agent._parse_action(response)
        assert action["tool"] == "run_scan"

    def test_parse_heuristic_graph(self, agent: Any) -> None:
        response = "Let me analyze the knowledge graph for attack paths."
        action = agent._parse_action(response)
        assert action["tool"] == "analyze_graph"

    def test_parse_heuristic_vectors(self, agent: Any) -> None:
        response = "We should generate dynamic custom attacks based on capabilities."
        action = agent._parse_action(response)
        assert action["tool"] == "generate_vectors"

    def test_parse_fallback_to_scan(self, agent: Any) -> None:
        response = "I'm not sure what to do next, let me think about it."
        action = agent._parse_action(response)
        assert action["tool"] == "run_scan"


# ──────────────────────────────────────────────────────────────────────
# PentestAgent._summarize_result
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestSummarizeResult:
    """Tests for tool result summarization."""

    @pytest.fixture
    def agent(self) -> Any:
        from ziran.application.pentesting.agent import PentestAgent

        return PentestAgent(
            scanner=_mock_scanner(),
            llm_client=_mock_llm_client(),
            max_iterations=5,
        )

    def test_scan_summary(self, agent: Any) -> None:
        result = {"total_vulnerabilities": 3, "total_attacks_executed": 10}
        summary = agent._summarize_result("run_scan", result)
        assert "3 vulns" in summary
        assert "10 attacks" in summary

    def test_phase_summary(self, agent: Any) -> None:
        result = {"vulnerabilities_found": ["a", "b"]}
        summary = agent._summarize_result("run_phase", result)
        assert "2 vulns" in summary

    def test_graph_summary(self, agent: Any) -> None:
        result = {"total_attack_paths": 5, "dangerous_chains": [1, 2]}
        summary = agent._summarize_result("analyze_graph", result)
        assert "5 attack paths" in summary
        assert "2 dangerous chains" in summary

    def test_query_summary(self, agent: Any) -> None:
        result = {"total_matching": 7}
        summary = agent._summarize_result("query_findings", result)
        assert "7" in summary

    def test_vectors_summary(self, agent: Any) -> None:
        result = {"generated_count": 4}
        summary = agent._summarize_result("generate_vectors", result)
        assert "4" in summary

    def test_explain_summary(self, agent: Any) -> None:
        summary = agent._summarize_result("explain_finding", {})
        assert "explained" in summary.lower()

    def test_remediate_summary(self, agent: Any) -> None:
        summary = agent._summarize_result("suggest_remediation", {})
        assert "remediation" in summary.lower()

    def test_search_summary(self, agent: Any) -> None:
        result = {"query": "LangChain CVEs"}
        summary = agent._summarize_result("web_search", result)
        assert "LangChain CVEs" in summary

    def test_unknown_tool_summary(self, agent: Any) -> None:
        result = {}
        summary = agent._summarize_result("unknown_tool", result)
        assert "unknown_tool" in summary


# ──────────────────────────────────────────────────────────────────────
# PentestAgent routing
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestRouting:
    """Tests for the conditional routing logic."""

    @pytest.fixture
    def agent(self) -> Any:
        from ziran.application.pentesting.agent import PentestAgent

        return PentestAgent(
            scanner=_mock_scanner(),
            llm_client=_mock_llm_client(),
            max_iterations=5,
        )

    def test_route_after_executor_to_reasoner(self, agent: Any) -> None:
        session = PentestSession(status=PentestStatus.EXECUTING)
        state = {"session": session, "should_stop": False}
        assert agent._route_after_executor(state) == "reasoner"

    def test_route_after_executor_to_reporter_on_failure(self, agent: Any) -> None:
        session = PentestSession(status=PentestStatus.FAILED)
        state = {"session": session, "should_stop": False}
        assert agent._route_after_executor(state) == "reporter"

    def test_route_after_executor_to_reporter_on_stop(self, agent: Any) -> None:
        session = PentestSession(status=PentestStatus.EXECUTING)
        state = {"session": session, "should_stop": True}
        assert agent._route_after_executor(state) == "reporter"

    def test_route_after_reasoner_to_planner(self, agent: Any) -> None:
        session = PentestSession(max_iterations=10, iteration_count=3)
        session.status = PentestStatus.EXECUTING
        state = {"session": session, "should_stop": False}
        assert agent._route_after_reasoner(state) == "planner"

    def test_route_after_reasoner_to_reporter_max_iterations(self, agent: Any) -> None:
        session = PentestSession(max_iterations=5, iteration_count=5)
        state = {"session": session, "should_stop": False}
        assert agent._route_after_reasoner(state) == "reporter"

    def test_route_after_reasoner_to_reporter_on_stop(self, agent: Any) -> None:
        session = PentestSession(max_iterations=10, iteration_count=1)
        state = {"session": session, "should_stop": True}
        assert agent._route_after_reasoner(state) == "reporter"


# ──────────────────────────────────────────────────────────────────────
# PentestAgent._execute_tool
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestExecuteTool:
    """Tests for the tool dispatch function."""

    @pytest.fixture
    def agent(self) -> Any:
        from ziran.application.pentesting.agent import PentestAgent

        return PentestAgent(
            scanner=_mock_scanner(),
            llm_client=_mock_llm_client(),
            max_iterations=5,
        )

    async def test_dispatch_run_scan(self, agent: Any) -> None:
        with patch("ziran.application.pentesting.tools.scan.run_scan", new_callable=AsyncMock) as m:
            m.return_value = {"total_vulnerabilities": 0}
            result = await agent._execute_tool("run_scan", {})
            m.assert_called_once()
            assert result["total_vulnerabilities"] == 0

    async def test_dispatch_scan_alias(self, agent: Any) -> None:
        with patch("ziran.application.pentesting.tools.scan.run_scan", new_callable=AsyncMock) as m:
            m.return_value = {"total_vulnerabilities": 0}
            await agent._execute_tool("scan", {})
            m.assert_called_once()

    async def test_dispatch_analyze_graph(self, agent: Any) -> None:
        with patch("ziran.application.pentesting.tools.graph.analyze_graph") as m:
            m.return_value = {"node_count": 5}
            result = await agent._execute_tool("analyze_graph", {})
            m.assert_called_once()
            assert result["node_count"] == 5

    async def test_dispatch_explain(self, agent: Any) -> None:
        with patch(
            "ziran.application.pentesting.tools.explain.explain_finding",
            new_callable=AsyncMock,
        ) as m:
            m.return_value = {"explanation": "test"}
            await agent._execute_tool("explain_finding", {"finding": {}})
            m.assert_called_once()

    async def test_dispatch_unknown_defaults_to_scan(self, agent: Any) -> None:
        with patch("ziran.application.pentesting.tools.scan.run_scan", new_callable=AsyncMock) as m:
            m.return_value = {"total_vulnerabilities": 0}
            await agent._execute_tool("nonexistent_tool", {})
            m.assert_called_once()


# ──────────────────────────────────────────────────────────────────────
# PentestAgent.run
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestAgentRun:
    """Tests for the top-level run method."""

    async def test_run_creates_session_and_plan(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent

        # Build an agent with a mock graph that returns immediately
        scanner = _mock_scanner()
        llm = _mock_llm_client()
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=1)

        # Mock the LangGraph compiled graph to shortcut execution
        mock_graph = MagicMock()
        final_session = PentestSession(
            status=PentestStatus.COMPLETE,
            max_iterations=1,
        )
        final_session.plan = PentestPlan(goal="test goal")
        mock_graph.ainvoke = AsyncMock(return_value={"session": final_session})
        agent._graph = mock_graph

        result = await agent.run(goal="test goal")

        assert result.status == PentestStatus.COMPLETE
        assert result.plan is not None
        assert result.plan.goal == "test goal"
        mock_graph.ainvoke.assert_called_once()

    async def test_run_handles_exception(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent

        scanner = _mock_scanner()
        llm = _mock_llm_client()
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=1)

        mock_graph = MagicMock()
        mock_graph.ainvoke = AsyncMock(side_effect=RuntimeError("LangGraph crash"))
        agent._graph = mock_graph

        result = await agent.run(goal="fail test")

        assert result.status == PentestStatus.FAILED
        assert len(result.steps) == 1
        assert "error" in result.steps[0].reasoning.lower()


# ──────────────────────────────────────────────────────────────────────
# PentestAgent.send_directive
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestSendDirective:
    """Tests for the interactive directive processing."""

    async def test_send_directive_adds_message_and_runs(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent

        scanner = _mock_scanner()
        llm = _mock_llm_client()
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=5)

        final_session = PentestSession(max_iterations=5)
        from ziran.domain.entities.pentest import PentestStep, PentestStepType

        final_session.add_step(
            PentestStep(step_type=PentestStepType.PLAN, reasoning="test", result_summary="Done")
        )
        mock_graph = MagicMock()
        mock_graph.ainvoke = AsyncMock(return_value={"session": final_session})
        agent._graph = mock_graph

        session = PentestSession(max_iterations=5)
        result = await agent.send_directive(session, "Focus on prompt injection")

        assert len(result.conversation_history) >= 1

    async def test_send_directive_handles_error(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent

        scanner = _mock_scanner()
        llm = _mock_llm_client()
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=5)

        mock_graph = MagicMock()
        mock_graph.ainvoke = AsyncMock(side_effect=RuntimeError("boom"))
        agent._graph = mock_graph

        session = PentestSession(max_iterations=5)
        result = await agent.send_directive(session, "test")

        # Should add error message instead of crashing
        assert any(
            "Error" in m.content or "error" in m.content for m in result.conversation_history
        )


# ──────────────────────────────────────────────────────────────────────
# PentestAgent._planner_node
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestPlannerNode:
    """Tests for the planner node."""

    async def test_planner_creates_step_and_action(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent

        scanner = _mock_scanner()
        resp_content = json.dumps({"tool": "run_phase", "args": {"phase": "reconnaissance"}})
        llm = _mock_llm_client(responses=[resp_content])
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=3)

        session = PentestSession(max_iterations=3)
        session.plan = PentestPlan(goal="test planner")
        state = {"session": session, "next_action": "", "should_stop": False}

        result = await agent._planner_node(state)

        assert result["session"].iteration_count == 1
        assert len(result["session"].steps) == 1
        assert result["session"].steps[0].step_type.value == "plan"
        action = json.loads(result["next_action"])
        assert action["tool"] == "run_phase"

    async def test_planner_extracts_capabilities(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent

        scanner = _mock_scanner()
        llm = _mock_llm_client(responses=['{"tool": "run_scan", "args": {}}'])
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=3)

        session = PentestSession(max_iterations=3)
        session.plan = PentestPlan(goal="test")
        state = {"session": session, "next_action": "", "should_stop": False, "attack_results": []}

        result = await agent._planner_node(state)
        assert "next_action" in result


# ──────────────────────────────────────────────────────────────────────
# PentestAgent._executor_node
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestExecutorNode:
    """Tests for the executor node."""

    async def test_executor_runs_tool_and_records_step(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent

        scanner = _mock_scanner()
        llm = _mock_llm_client()
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=3)

        session = PentestSession(max_iterations=3)
        action = json.dumps({"tool": "analyze_graph", "args": {}})

        with patch("ziran.application.pentesting.tools.graph.analyze_graph") as mock_graph:
            mock_graph.return_value = {
                "node_count": 5,
                "edge_count": 3,
                "total_attack_paths": 2,
                "dangerous_chains": [],
            }
            state = {"session": session, "next_action": action, "should_stop": False}
            result = await agent._executor_node(state)

        assert len(result["session"].steps) == 1

    async def test_executor_handles_tool_error(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent

        scanner = _mock_scanner()
        llm = _mock_llm_client()
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=3)

        session = PentestSession(max_iterations=3)
        action = json.dumps({"tool": "run_scan", "args": {}})

        with patch("ziran.application.pentesting.tools.scan.run_scan", new_callable=AsyncMock) as m:
            m.side_effect = RuntimeError("scan failed")
            state = {"session": session, "next_action": action, "should_stop": False}
            result = await agent._executor_node(state)

        assert "error" in result
        assert len(result["session"].steps) == 1
        assert "Error" in result["session"].steps[0].result_summary

    async def test_executor_tracks_tested_vectors(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent

        scanner = _mock_scanner()
        llm = _mock_llm_client()
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=3)

        session = PentestSession(max_iterations=3)
        action = json.dumps({"tool": "run_scan", "args": {}})

        scan_result = {
            "total_vulnerabilities": 2,
            "total_attacks_executed": 5,
            "attacks_executed": 5,
            "successful_attacks": [
                {"vector_id": "v1"},
                {"vector_id": "v2"},
            ],
        }
        ar1 = MagicMock()
        ar1.vector_id = "v1"
        ar2 = MagicMock()
        ar2.vector_id = "v3"
        scanner._attack_results = [ar1, ar2]

        with patch("ziran.application.pentesting.tools.scan.run_scan", new_callable=AsyncMock) as m:
            m.return_value = scan_result
            state = {"session": session, "next_action": action, "should_stop": False}
            await agent._executor_node(state)

        assert "v1" in agent._tested_vectors
        assert "v3" in agent._tested_vectors


# ──────────────────────────────────────────────────────────────────────
# PentestAgent._reasoner_node
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestReasonerNode:
    """Tests for the reasoner node."""

    async def test_reasoner_continues(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent
        from ziran.domain.entities.pentest import PentestStep, PentestStepType

        scanner = _mock_scanner()
        llm = _mock_llm_client(responses=["We should run_phase to find more vulnerabilities."])
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=5)

        session = PentestSession(max_iterations=5)
        session.add_step(
            PentestStep(
                step_type=PentestStepType.SCAN,
                reasoning="test",
                result_summary="Scan complete: 5 vulns",
            )
        )
        session.total_vulnerabilities = 5

        with patch("ziran.application.pentesting.tools.graph.analyze_graph") as m:
            m.return_value = {
                "node_count": 5,
                "edge_count": 3,
                "total_attack_paths": 2,
                "dangerous_chains": [],
            }
            state = {"session": session, "next_action": "", "should_stop": False}
            result = await agent._reasoner_node(state)

        assert result["should_stop"] is False

    async def test_reasoner_stops_on_keyword(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent
        from ziran.domain.entities.pentest import PentestStep, PentestStepType

        scanner = _mock_scanner()
        llm = _mock_llm_client(responses=["Session complete. Sufficient findings discovered."])
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=5)

        session = PentestSession(max_iterations=5)
        session.add_step(
            PentestStep(
                step_type=PentestStepType.SCAN,
                reasoning="test",
                result_summary="Test",
            )
        )
        session.total_vulnerabilities = 10

        with patch("ziran.application.pentesting.tools.graph.analyze_graph") as m:
            m.return_value = {
                "node_count": 5,
                "edge_count": 3,
                "total_attack_paths": 2,
                "dangerous_chains": [],
            }
            state = {"session": session, "next_action": "", "should_stop": False}
            result = await agent._reasoner_node(state)

        assert result["should_stop"] is True

    async def test_plateau_detection_forces_stop(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent
        from ziran.domain.entities.pentest import PentestStep, PentestStepType

        scanner = _mock_scanner()
        llm = _mock_llm_client(responses=["Continue testing with run_phase."])
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=10)

        # Simulate 2 prior iterations with same vuln count
        agent._vuln_history = [10, 10]

        session = PentestSession(max_iterations=10)
        session.add_step(
            PentestStep(
                step_type=PentestStepType.SCAN,
                reasoning="test",
                result_summary="Test",
            )
        )
        session.total_vulnerabilities = 10  # Same as last 2

        with patch("ziran.application.pentesting.tools.graph.analyze_graph") as m:
            m.return_value = {
                "node_count": 5,
                "edge_count": 3,
                "total_attack_paths": 2,
                "dangerous_chains": [],
            }
            state = {"session": session, "next_action": "", "should_stop": False}
            result = await agent._reasoner_node(state)

        assert result["should_stop"] is True  # Plateau detected

    async def test_vuln_trend_in_reasoning(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent
        from ziran.domain.entities.pentest import PentestStep, PentestStepType

        scanner = _mock_scanner()
        llm = _mock_llm_client(responses=["Continue with run_scan."])
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=10)
        agent._vuln_history = [5]

        session = PentestSession(max_iterations=10)
        session.add_step(
            PentestStep(
                step_type=PentestStepType.SCAN,
                reasoning="test",
                result_summary="Test",
            )
        )
        session.total_vulnerabilities = 10

        with patch("ziran.application.pentesting.tools.graph.analyze_graph") as m:
            m.return_value = {
                "node_count": 5,
                "edge_count": 3,
                "total_attack_paths": 2,
                "dangerous_chains": [],
            }
            state = {"session": session, "next_action": "", "should_stop": False}
            await agent._reasoner_node(state)

        # Verify vuln_history is updated
        assert agent._vuln_history == [5, 10]


# ──────────────────────────────────────────────────────────────────────
# PentestAgent._reporter_node
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestReporterNode:
    """Tests for the reporter node."""

    async def test_reporter_generates_report(self) -> None:
        from ziran.application.pentesting.agent import PentestAgent
        from ziran.domain.entities.pentest import PentestStep, PentestStepType

        scanner = _mock_scanner()
        llm = _mock_llm_client(responses=["## Final Report\nAll tests complete."])
        agent = PentestAgent(scanner=scanner, llm_client=llm, max_iterations=3)

        session = PentestSession(max_iterations=3)
        session.plan = PentestPlan(goal="test reporter")
        session.add_step(
            PentestStep(
                step_type=PentestStepType.SCAN,
                reasoning="found vulns",
                result_summary="5 vulns found",
                result_data={"total_vulnerabilities": 5},
            )
        )

        with patch("ziran.application.pentesting.tools.graph.analyze_graph") as m:
            m.return_value = {"dangerous_chains": []}
            state = {"session": session, "should_stop": True}
            result = await agent._reporter_node(state)

        assert result["session"].status == PentestStatus.COMPLETE
        assert len(result["session"].steps) == 2  # original + reporter step
