"""Unit tests for pentesting agent tools (scan, graph, vectors, explain, external)."""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from ziran.domain.entities.attack import AttackCategory, AttackResult

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


def _mock_llm_client() -> MagicMock:
    """Create a mock LLM client (BaseLLMClient)."""
    client = MagicMock()
    response = MagicMock()
    response.content = "LLM response text"
    response.prompt_tokens = 10
    response.completion_tokens = 5
    response.total_tokens = 15
    client.complete = AsyncMock(return_value=response)
    return client


def _make_attack_result(
    vector_id: str = "vec_001",
    vector_name: str = "Test Vector",
    category: AttackCategory = AttackCategory.PROMPT_INJECTION,
    severity: str = "high",
    successful: bool = True,
    agent_response: str = "Leaked data",
    owasp: list | None = None,
    evidence: dict | None = None,
) -> AttackResult:
    return AttackResult(
        vector_id=vector_id,
        vector_name=vector_name,
        category=category,
        severity=severity,
        successful=successful,
        agent_response=agent_response,
        evidence=evidence or {"response": agent_response},
        owasp_mapping=owasp or [],
    )


# ──────────────────────────────────────────────────────────────────────
# tools/graph.py — analyze_graph
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestAnalyzeGraph:
    """Tests for the graph analysis tool."""

    def test_returns_node_edge_counts(self) -> None:
        from ziran.application.pentesting.tools.graph import analyze_graph

        graph = MagicMock()
        graph.node_count = 5
        graph.edge_count = 8
        graph.find_all_attack_paths.return_value = []
        graph.get_critical_nodes.return_value = []
        graph.export_state.return_value = {"nodes": []}

        result = analyze_graph(graph, include_paths=True, include_critical_nodes=True)

        assert result["node_count"] == 5
        assert result["edge_count"] == 8

    def test_includes_attack_paths(self) -> None:
        from ziran.application.pentesting.tools.graph import analyze_graph

        graph = MagicMock()
        graph.node_count = 3
        graph.edge_count = 2
        graph.find_all_attack_paths.return_value = [["a", "b", "c"], ["d", "e"]]
        graph.get_critical_nodes.return_value = []
        graph.export_state.return_value = {"nodes": []}

        result = analyze_graph(graph, include_paths=True, include_critical_nodes=False)

        assert result["total_attack_paths"] == 2
        assert len(result["attack_paths"]) == 2
        assert result["attack_paths"][0]["path"] == ["a", "b", "c"]
        assert result["attack_paths"][0]["length"] == 3

    def test_includes_critical_nodes(self) -> None:
        from ziran.application.pentesting.tools.graph import analyze_graph

        graph = MagicMock()
        graph.node_count = 3
        graph.edge_count = 2
        graph.find_all_attack_paths.return_value = []
        graph.get_critical_nodes.return_value = [("node_a", 0.95), ("node_b", 0.42)]
        graph.export_state.return_value = {"nodes": []}

        result = analyze_graph(graph, include_paths=False, include_critical_nodes=True)

        assert len(result["critical_nodes"]) == 2
        assert result["critical_nodes"][0]["node"] == "node_a"
        assert result["critical_nodes"][0]["centrality"] == 0.95

    def test_node_type_summary(self) -> None:
        from ziran.application.pentesting.tools.graph import analyze_graph

        graph = MagicMock()
        graph.node_count = 3
        graph.edge_count = 2
        graph.find_all_attack_paths.return_value = []
        graph.get_critical_nodes.return_value = []
        graph.export_state.return_value = {
            "nodes": [
                ("n1", {"node_type": "tool"}),
                ("n2", {"node_type": "tool"}),
                ("n3", {"node_type": "vulnerability"}),
            ]
        }

        result = analyze_graph(graph, include_paths=False, include_critical_nodes=False)
        assert result["node_type_summary"] == {"tool": 2, "vulnerability": 1}


# ──────────────────────────────────────────────────────────────────────
# tools/graph.py — query_findings
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestQueryFindings:
    """Tests for the finding query tool."""

    def test_returns_all_by_default(self) -> None:
        from ziran.application.pentesting.tools.graph import query_findings

        results = [_make_attack_result(vector_id=f"v{i}") for i in range(3)]
        out = query_findings(results)
        assert out["total_matching"] == 3
        assert out["returned"] == 3

    def test_filter_successful_only(self) -> None:
        from ziran.application.pentesting.tools.graph import query_findings

        results = [
            _make_attack_result(vector_id="v1", successful=True),
            _make_attack_result(vector_id="v2", successful=False),
        ]
        out = query_findings(results, successful_only=True)
        assert out["total_matching"] == 1
        assert out["findings"][0]["vector_id"] == "v1"

    def test_filter_by_category(self) -> None:
        from ziran.application.pentesting.tools.graph import query_findings

        results = [
            _make_attack_result(category=AttackCategory.PROMPT_INJECTION),
            _make_attack_result(category=AttackCategory.DATA_EXFILTRATION),
        ]
        out = query_findings(results, category="prompt_injection")
        assert out["total_matching"] == 1

    def test_filter_by_severity(self) -> None:
        from ziran.application.pentesting.tools.graph import query_findings

        results = [
            _make_attack_result(severity="high"),
            _make_attack_result(severity="low"),
        ]
        out = query_findings(results, severity="high")
        assert out["total_matching"] == 1

    def test_limit(self) -> None:
        from ziran.application.pentesting.tools.graph import query_findings

        results = [_make_attack_result(vector_id=f"v{i}") for i in range(10)]
        out = query_findings(results, limit=3)
        assert out["total_matching"] == 10
        assert out["returned"] == 3

    def test_finding_dict_shape(self) -> None:
        from ziran.application.pentesting.tools.graph import query_findings

        results = [_make_attack_result()]
        out = query_findings(results)
        finding = out["findings"][0]
        assert "vector_id" in finding
        assert "vector_name" in finding
        assert "category" in finding
        assert "severity" in finding
        assert "successful" in finding
        assert "evidence_keys" in finding


# ──────────────────────────────────────────────────────────────────────
# tools/vectors.py — generate_vectors
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestGenerateVectors:
    """Tests for the dynamic vector generation tool."""

    def test_generate_with_capabilities(self) -> None:
        from ziran.application.pentesting.tools.vectors import generate_vectors

        result = generate_vectors(capabilities=["search", "send_email"], count=5)
        assert "generated_count" in result
        assert "vectors" in result
        assert isinstance(result["vectors"], list)
        assert result["capabilities_used"] == ["search", "send_email"]

    def test_generate_with_focus_categories(self) -> None:
        from ziran.application.pentesting.tools.vectors import generate_vectors

        result = generate_vectors(
            capabilities=["shell_execute"],
            focus_categories=["prompt_injection"],
            count=3,
        )
        # All returned vectors should match the focus category
        for v in result["vectors"]:
            assert v["category"] == "prompt_injection"

    def test_generate_limits_count(self) -> None:
        from ziran.application.pentesting.tools.vectors import generate_vectors

        result = generate_vectors(capabilities=["search"], count=1)
        assert result["generated_count"] <= 1

    def test_generate_handles_error(self) -> None:
        from ziran.application.pentesting.tools.vectors import generate_vectors

        with patch(
            "ziran.application.dynamic_vectors.generator.DynamicVectorGenerator"
        ) as mock_gen:
            mock_gen.return_value.generate.side_effect = RuntimeError("boom")
            result = generate_vectors(capabilities=["x"])
            assert result["generated_count"] == 0
            assert "error" in result


# ──────────────────────────────────────────────────────────────────────
# tools/explain.py — explain_finding
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestExplainFinding:
    """Tests for the LLM-powered explain tool."""

    async def test_explain_returns_explanation(self) -> None:
        from ziran.application.pentesting.tools.explain import explain_finding

        llm = _mock_llm_client()
        finding = {
            "vector_id": "vec_001",
            "vector_name": "PI Test",
            "category": "prompt_injection",
            "severity": "high",
        }
        result = await explain_finding(llm, finding=finding)

        assert result["explanation"] == "LLM response text"
        assert result["vector_id"] == "vec_001"
        assert result["token_usage"]["total_tokens"] == 15
        llm.complete.assert_called_once()

    async def test_explain_handles_error(self) -> None:
        from ziran.application.pentesting.tools.explain import explain_finding

        llm = _mock_llm_client()
        llm.complete.side_effect = RuntimeError("LLM unavailable")

        result = await explain_finding(llm, finding={"vector_id": "v1"})
        assert "error" in result
        assert "LLM unavailable" in result["error"]


# ──────────────────────────────────────────────────────────────────────
# tools/explain.py — suggest_remediation
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestSuggestRemediation:
    """Tests for the remediation suggestion tool."""

    async def test_remediation_returns_suggestion(self) -> None:
        from ziran.application.pentesting.tools.explain import suggest_remediation

        llm = _mock_llm_client()
        finding = {"vector_id": "v1", "category": "data_exfiltration", "severity": "critical"}

        result = await suggest_remediation(llm, finding=finding, framework="langchain")
        assert result["remediation"] == "LLM response text"
        assert result["framework"] == "langchain"
        assert result["vector_id"] == "v1"

    async def test_remediation_handles_error(self) -> None:
        from ziran.application.pentesting.tools.explain import suggest_remediation

        llm = _mock_llm_client()
        llm.complete.side_effect = RuntimeError("fail")

        result = await suggest_remediation(llm, finding={"vector_id": "v1"})
        assert "error" in result


# ──────────────────────────────────────────────────────────────────────
# tools/external.py — web_search
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestWebSearch:
    """Tests for the LLM-based knowledge search tool."""

    async def test_search_returns_results(self) -> None:
        from ziran.application.pentesting.tools.external import web_search

        llm = _mock_llm_client()
        result = await web_search(llm, query="LangChain CVEs")

        assert result["query"] == "LangChain CVEs"
        assert result["results"] == "LLM response text"
        assert "token_usage" in result

    async def test_search_with_context(self) -> None:
        from ziran.application.pentesting.tools.external import web_search

        llm = _mock_llm_client()
        result = await web_search(llm, query="test", context="RAG agent")
        assert result["results"] == "LLM response text"

    async def test_search_handles_error(self) -> None:
        from ziran.application.pentesting.tools.external import web_search

        llm = _mock_llm_client()
        llm.complete.side_effect = RuntimeError("fail")

        result = await web_search(llm, query="test")
        assert result["results"] == ""
        assert "error" in result


# ──────────────────────────────────────────────────────────────────────
# tools/external.py — read_target_docs
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestReadTargetDocs:
    """Tests for the documentation fetcher tool."""

    async def test_fetch_success(self) -> None:
        from ziran.application.pentesting.tools.external import read_target_docs

        mock_resp = MagicMock()
        mock_resp.text = "# API Documentation\n\nEndpoints listed here..."
        mock_resp.raise_for_status = MagicMock()

        mock_client = AsyncMock()
        mock_client.get.return_value = mock_resp

        mock_cm = AsyncMock()
        mock_cm.__aenter__ = AsyncMock(return_value=mock_client)
        mock_cm.__aexit__ = AsyncMock(return_value=False)

        with patch("httpx.AsyncClient", return_value=mock_cm):
            result = await read_target_docs(url="https://example.com/docs")

        assert result["url"] == "https://example.com/docs"
        assert result["content_length"] > 0

    async def test_fetch_failure(self) -> None:
        from ziran.application.pentesting.tools.external import read_target_docs

        mock_client = AsyncMock()
        mock_client.get.side_effect = RuntimeError("connection refused")

        mock_cm = AsyncMock()
        mock_cm.__aenter__ = AsyncMock(return_value=mock_client)
        mock_cm.__aexit__ = AsyncMock(return_value=False)

        with patch("httpx.AsyncClient", return_value=mock_cm):
            result = await read_target_docs(url="https://bad.example.com")

        assert "error" in result
        assert result["content"] == ""


# ──────────────────────────────────────────────────────────────────────
# tools/scan.py — run_scan / run_phase
# ──────────────────────────────────────────────────────────────────────


@pytest.mark.unit
class TestScanTools:
    """Tests for run_scan and run_phase tools."""

    def _mock_scanner(self) -> MagicMock:
        """Create a mock AgentScanner that returns a plausible campaign result."""
        scanner = MagicMock()

        # Mock PhaseResult
        phase_result = MagicMock()
        phase_result.vulnerabilities_found = ["v1"]
        phase_result.trust_score = 0.7
        phase_result.duration_seconds = 2.5
        phase_result.token_usage = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}

        # Mock CampaignResult
        campaign = MagicMock()
        campaign.campaign_id = "test_campaign"
        campaign.phases_executed = [phase_result]
        campaign.total_vulnerabilities = 1
        campaign.final_trust_score = 0.7
        campaign.critical_paths = []
        campaign.attack_results = [
            {
                "vector_id": "v1",
                "vector_name": "Test",
                "category": "prompt_injection",
                "severity": "high",
                "successful": True,
            }
        ]
        campaign.token_usage = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}

        scanner.run_campaign = AsyncMock(return_value=campaign)
        return scanner

    async def test_run_scan_default(self) -> None:
        from ziran.application.pentesting.tools.scan import run_scan

        scanner = self._mock_scanner()
        result = await run_scan(scanner)

        assert result["campaign_id"] == "test_campaign"
        assert result["total_vulnerabilities"] == 1
        assert result["total_attacks_executed"] == 1
        scanner.run_campaign.assert_called_once()

    async def test_run_scan_with_phases(self) -> None:
        from ziran.application.pentesting.tools.scan import run_scan

        scanner = self._mock_scanner()
        await run_scan(scanner, phases=["reconnaissance"])

        scanner.run_campaign.assert_called_once()
        call_kwargs = scanner.run_campaign.call_args[1]
        assert call_kwargs["phases"] is not None
        assert len(call_kwargs["phases"]) == 1

    async def test_run_phase(self) -> None:
        from ziran.application.pentesting.tools.scan import run_phase

        scanner = self._mock_scanner()
        result = await run_phase(scanner, phase="reconnaissance")

        assert result["phase"] == "reconnaissance"
        assert "trust_score" in result
        assert "duration_seconds" in result
