"""Tests for the agent context graph — lateral movement analysis."""

from __future__ import annotations

from agent_bom.context_graph import (
    ContextGraph,
    EdgeKind,
    NodeKind,
    build_context_graph,
    compute_interaction_risks,
    find_lateral_paths,
    to_serializable,
)

# ---------------------------------------------------------------------------
# Helpers — mirror JSON shapes from to_json() output
# ---------------------------------------------------------------------------


def _agent(
    name: str = "agent-a",
    agent_type: str = "claude-desktop",
    servers: list | None = None,
) -> dict:
    return {
        "name": name,
        "type": agent_type,
        "status": "configured",
        "mcp_servers": servers or [],
    }


def _server(
    name: str = "filesystem",
    env: dict | None = None,
    tools: list | None = None,
    packages: list | None = None,
) -> dict:
    return {
        "name": name,
        "command": "npx",
        "transport": "stdio",
        "env": env or {},
        "tools": tools or [],
        "packages": packages or [],
    }


def _tool(name: str = "read_file", description: str = "Read a file") -> dict:
    return {"name": name, "description": description}


def _blast(
    vuln_id: str = "CVE-2025-0001",
    severity: str = "critical",
    agents: list | None = None,
    servers: list | None = None,
    package: str = "langchain@0.1.0",
    cvss_score: float = 9.8,
) -> dict:
    return {
        "vulnerability_id": vuln_id,
        "severity": severity,
        "cvss_score": cvss_score,
        "epss_score": 0.5,
        "is_kev": False,
        "risk_score": 8.0,
        "package": package,
        "affected_agents": agents or ["agent-a"],
        "affected_servers": servers or ["filesystem"],
        "exposed_credentials": ["GITHUB_TOKEN"],
        "exposed_tools": ["read_file"],
    }


# ---------------------------------------------------------------------------
# TestBuildGraph
# ---------------------------------------------------------------------------


class TestBuildGraph:
    def test_empty_graph(self):
        graph = build_context_graph([], [])
        assert len(graph.nodes) == 0
        assert len(graph.edges) == 0

    def test_single_agent_no_servers(self):
        graph = build_context_graph([_agent()], [])
        assert "agent:agent-a" in graph.nodes
        assert graph.nodes["agent:agent-a"].kind == NodeKind.AGENT
        assert len(graph.edges) == 0

    def test_agent_server_edges(self):
        agents = [_agent(servers=[_server()])]
        graph = build_context_graph(agents, [])
        assert "agent:agent-a" in graph.nodes
        assert "server:agent-a:filesystem" in graph.nodes
        uses_edges = [e for e in graph.edges if e.kind == EdgeKind.USES]
        assert len(uses_edges) == 1
        assert uses_edges[0].source == "agent:agent-a"

    def test_credential_nodes(self):
        agents = [_agent(servers=[_server(env={"GITHUB_TOKEN": "xxx", "PATH": "/usr/bin"})])]
        graph = build_context_graph(agents, [])
        assert "cred:GITHUB_TOKEN" in graph.nodes
        assert graph.nodes["cred:GITHUB_TOKEN"].kind == NodeKind.CREDENTIAL
        # PATH should not be a credential
        assert "cred:PATH" not in graph.nodes
        exposes = [e for e in graph.edges if e.kind == EdgeKind.EXPOSES]
        assert len(exposes) == 1

    def test_tool_nodes_with_capability(self):
        agents = [_agent(servers=[_server(tools=[_tool("execute_code", "Run arbitrary code")])])]
        graph = build_context_graph(agents, [])
        tool_id = "tool:server:agent-a:filesystem:execute_code"
        assert tool_id in graph.nodes
        assert graph.nodes[tool_id].kind == NodeKind.TOOL
        caps = graph.nodes[tool_id].metadata.get("capabilities", [])
        assert "execute" in caps

    def test_vulnerability_edges(self):
        agents = [_agent(servers=[_server()])]
        blast = [_blast()]
        graph = build_context_graph(agents, blast)
        assert "vuln:CVE-2025-0001" in graph.nodes
        vuln_edges = [e for e in graph.edges if e.kind == EdgeKind.VULNERABLE_TO]
        assert len(vuln_edges) == 1
        assert vuln_edges[0].source == "server:agent-a:filesystem"

    def test_shared_server_detection(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="shared-srv")]),
            _agent(name="agent-b", servers=[_server(name="shared-srv")]),
        ]
        graph = build_context_graph(agents, [])
        shares = [e for e in graph.edges if e.kind == EdgeKind.SHARES_SERVER]
        assert len(shares) == 1
        assert set([shares[0].source, shares[0].target]) == {"agent:agent-a", "agent:agent-b"}

    def test_shared_credential_detection(self):
        agents = [
            _agent(name="agent-a", servers=[_server(env={"API_KEY": "xxx"})]),
            _agent(name="agent-b", servers=[_server(env={"API_KEY": "yyy"})]),
        ]
        graph = build_context_graph(agents, [])
        shares = [e for e in graph.edges if e.kind == EdgeKind.SHARES_CREDENTIAL]
        assert len(shares) == 1
        assert shares[0].metadata["credential"] == "API_KEY"


# ---------------------------------------------------------------------------
# TestLateralPaths
# ---------------------------------------------------------------------------


class TestLateralPaths:
    def test_basic_bfs_through_shared_server(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="shared-srv")]),
            _agent(name="agent-b", servers=[_server(name="shared-srv")]),
        ]
        graph = build_context_graph(agents, [])
        paths = find_lateral_paths(graph, "agent:agent-a")
        assert len(paths) >= 1
        targets = {p.target for p in paths}
        assert "agent:agent-b" in targets

    def test_max_depth_limit(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="shared-srv")]),
            _agent(name="agent-b", servers=[_server(name="shared-srv")]),
        ]
        graph = build_context_graph(agents, [])
        # depth=0 means no expansion beyond start node
        paths = find_lateral_paths(graph, "agent:agent-a", max_depth=0)
        assert len(paths) == 0

    def test_no_cycles(self):
        # Two agents sharing server and credential — should not loop
        agents = [
            _agent(name="agent-a", servers=[_server(name="shared-srv", env={"API_KEY": "x"})]),
            _agent(name="agent-b", servers=[_server(name="shared-srv", env={"API_KEY": "y"})]),
        ]
        graph = build_context_graph(agents, [])
        paths = find_lateral_paths(graph, "agent:agent-a", max_depth=4)
        # Should terminate, not infinite loop
        assert isinstance(paths, list)

    def test_risk_scoring(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="srv", env={"SECRET_KEY": "x"})]),
            _agent(name="agent-b", servers=[_server(name="srv")]),
        ]
        blast = [_blast(severity="critical", agents=["agent-a", "agent-b"], servers=["srv"])]
        graph = build_context_graph(agents, blast)
        paths = find_lateral_paths(graph, "agent:agent-a")
        # Should have non-zero risk from credentials + vulns
        if paths:
            assert paths[0].composite_risk > 0

    def test_credential_along_path(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="srv", env={"TOKEN": "x"})]),
            _agent(name="agent-b", servers=[_server(name="srv")]),
        ]
        graph = build_context_graph(agents, [])
        paths = find_lateral_paths(graph, "agent:agent-a")
        cred_paths = [p for p in paths if p.credential_exposure]
        assert len(cred_paths) >= 1

    def test_execute_tool_along_path(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="srv", tools=[_tool("run_shell", "Execute shell commands")])]),
            _agent(name="agent-b", servers=[_server(name="srv")]),
        ]
        graph = build_context_graph(agents, [])
        paths = find_lateral_paths(graph, "agent:agent-a")
        tool_paths = [p for p in paths if p.tool_exposure]
        assert len(tool_paths) >= 1

    def test_path_summary_format(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="shared")]),
            _agent(name="agent-b", servers=[_server(name="shared")]),
        ]
        graph = build_context_graph(agents, [])
        paths = find_lateral_paths(graph, "agent:agent-a")
        for p in paths:
            assert " → " in p.summary
            assert p.summary.startswith("agent-a")

    def test_cap_at_100_paths(self):
        # Many agents sharing the same server
        agents = [_agent(name=f"agent-{i}", servers=[_server(name="shared")]) for i in range(20)]
        graph = build_context_graph(agents, [])
        paths = find_lateral_paths(graph, "agent:agent-0", max_depth=4)
        assert len(paths) <= 100

    def test_nonexistent_source(self):
        graph = build_context_graph([], [])
        paths = find_lateral_paths(graph, "agent:nonexistent")
        assert paths == []

    def test_bfs_queue_bound_on_dense_graph(self):
        """BFS terminates without OOM on a dense graph with many interconnections."""
        from agent_bom.context_graph import _MAX_QUEUE_SIZE

        # 15 agents sharing the same server — creates O(n^2) cross-links
        agents = [
            _agent(
                name=f"agent-{i}",
                servers=[
                    _server(name="shared", env={"KEY": "x"}, tools=[_tool("exec", "run code")]),
                ],
            )
            for i in range(15)
        ]
        graph = build_context_graph(agents, [])
        # Must complete quickly without hanging or OOM
        paths = find_lateral_paths(graph, "agent:agent-0", max_depth=4)
        assert isinstance(paths, list)
        assert len(paths) <= 100
        # Verify the constant is defined (regression guard)
        assert _MAX_QUEUE_SIZE > 0


# ---------------------------------------------------------------------------
# TestInteractionRisks
# ---------------------------------------------------------------------------


class TestInteractionRisks:
    def test_shared_credential_pattern(self):
        agents = [
            _agent(name="agent-a", servers=[_server(env={"DB_PASSWORD": "x"})]),
            _agent(name="agent-b", servers=[_server(env={"DB_PASSWORD": "y"})]),
        ]
        graph = build_context_graph(agents, [])
        risks = compute_interaction_risks(graph)
        cred_risks = [r for r in risks if r.pattern == "shared_credential"]
        assert len(cred_risks) == 1
        assert "DB_PASSWORD" in cred_risks[0].description

    def test_shared_server_pattern(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="common-srv")]),
            _agent(name="agent-b", servers=[_server(name="common-srv")]),
        ]
        graph = build_context_graph(agents, [])
        risks = compute_interaction_risks(graph)
        srv_risks = [r for r in risks if r.pattern == "shared_server"]
        assert len(srv_risks) == 1

    def test_tool_overlap_execute(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="srv-a", tools=[_tool("execute_code", "run code")])]),
            _agent(name="agent-b", servers=[_server(name="srv-b", tools=[_tool("execute_code", "run code")])]),
        ]
        graph = build_context_graph(agents, [])
        risks = compute_interaction_risks(graph)
        exec_risks = [r for r in risks if r.pattern == "tool_overlap_execute"]
        assert len(exec_risks) == 1

    def test_multi_hop_vuln(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="shared-srv")]),
            _agent(name="agent-b", servers=[_server(name="shared-srv")]),
        ]
        blast = [_blast(severity="critical", agents=["agent-a", "agent-b"], servers=["shared-srv"])]
        graph = build_context_graph(agents, blast)
        risks = compute_interaction_risks(graph)
        mhv_risks = [r for r in risks if r.pattern == "multi_hop_vuln"]
        assert len(mhv_risks) >= 1

    def test_no_risks_for_isolated_agents(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="srv-a")]),
            _agent(name="agent-b", servers=[_server(name="srv-b")]),
        ]
        graph = build_context_graph(agents, [])
        risks = compute_interaction_risks(graph)
        # No shared servers/creds, so no risks
        assert len(risks) == 0

    def test_owasp_tag_assignment(self):
        agents = [
            _agent(name="agent-a", servers=[_server(env={"TOKEN": "x"})]),
            _agent(name="agent-b", servers=[_server(env={"TOKEN": "y"})]),
        ]
        graph = build_context_graph(agents, [])
        risks = compute_interaction_risks(graph)
        cred_risks = [r for r in risks if r.pattern == "shared_credential"]
        assert cred_risks[0].owasp_agentic_tag == "ASI07"


# ---------------------------------------------------------------------------
# TestSerialization
# ---------------------------------------------------------------------------


class TestSerialization:
    def test_roundtrip(self):
        agents = [_agent(servers=[_server(env={"KEY": "v"}, tools=[_tool()])])]
        blast = [_blast()]
        graph = build_context_graph(agents, blast)
        data = to_serializable(graph)
        assert "nodes" in data
        assert "edges" in data
        assert "stats" in data
        assert data["stats"]["total_nodes"] > 0

    def test_stats_correctness(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="shared")]),
            _agent(name="agent-b", servers=[_server(name="shared")]),
        ]
        graph = build_context_graph(agents, [])
        paths = find_lateral_paths(graph, "agent:agent-a")
        risks = compute_interaction_risks(graph)
        data = to_serializable(graph, paths, risks)
        assert data["stats"]["agent_count"] == 2
        assert data["stats"]["shared_server_count"] >= 1
        assert data["stats"]["lateral_path_count"] == len(paths)
        assert data["stats"]["interaction_risk_count"] == len(risks)

    def test_empty_graph(self):
        graph = ContextGraph()
        data = to_serializable(graph)
        assert data["stats"]["total_nodes"] == 0
        assert data["lateral_paths"] == []
        assert data["interaction_risks"] == []

    def test_paths_included(self):
        agents = [
            _agent(name="agent-a", servers=[_server(name="shared")]),
            _agent(name="agent-b", servers=[_server(name="shared")]),
        ]
        graph = build_context_graph(agents, [])
        paths = find_lateral_paths(graph, "agent:agent-a")
        data = to_serializable(graph, paths)
        assert len(data["lateral_paths"]) > 0
        assert "summary" in data["lateral_paths"][0]
        assert "composite_risk" in data["lateral_paths"][0]


# ---------------------------------------------------------------------------
# TestEdgeCases
# ---------------------------------------------------------------------------


class TestEdgeCases:
    def test_agent_with_no_servers(self):
        graph = build_context_graph([_agent(servers=[])], [])
        assert len(graph.nodes) == 1
        assert len(graph.edges) == 0

    def test_server_with_no_packages(self):
        graph = build_context_graph([_agent(servers=[_server(packages=[])])], [])
        assert "server:agent-a:filesystem" in graph.nodes

    def test_duplicate_server_names_same_agent(self):
        # Two servers with same name in one agent — should not create SHARES_SERVER
        agents = [_agent(servers=[_server(name="dup"), _server(name="dup")])]
        graph = build_context_graph(agents, [])
        shares = [e for e in graph.edges if e.kind == EdgeKind.SHARES_SERVER]
        assert len(shares) == 0

    def test_large_graph_performance(self):
        # 50 agents each with 3 servers — should complete quickly
        agents = [_agent(name=f"a{i}", servers=[_server(name=f"s{j}") for j in range(3)]) for i in range(50)]
        graph = build_context_graph(agents, [])
        assert graph.nodes
        assert len(graph.edges) > 0


# ---------------------------------------------------------------------------
# TestIntegration — CLI flag and API endpoint
# ---------------------------------------------------------------------------


class TestOutputIntegration:
    def test_context_graph_in_to_json(self):
        """to_json() includes context_graph when set on report."""
        from agent_bom.output import to_json

        # Build minimal report with context_graph_data
        agents_data = [
            _agent(name="a1", servers=[_server(name="shared")]),
            _agent(name="a2", servers=[_server(name="shared")]),
        ]
        graph = build_context_graph(agents_data, [])
        paths = find_lateral_paths(graph, "agent:a1")
        risks = compute_interaction_risks(graph)
        cg_data = to_serializable(graph, paths, risks)

        from agent_bom.models import AIBOMReport

        report = AIBOMReport()
        report.context_graph_data = cg_data

        data = to_json(report)
        assert "context_graph" in data
        assert data["context_graph"]["stats"]["agent_count"] == 2
        assert data["context_graph"]["stats"]["shared_server_count"] >= 1


class TestAPIContextGraph:
    @staticmethod
    def _make_job(job_id, agents_data, blast_data=None):
        from datetime import datetime, timezone

        from agent_bom.api.server import JobStatus, ScanJob, ScanRequest

        return ScanJob(
            job_id=job_id,
            status=JobStatus.DONE,
            created_at=datetime.now(timezone.utc).isoformat(),
            completed_at=datetime.now(timezone.utc).isoformat(),
            request=ScanRequest(),
            result={
                "agents": agents_data,
                "blast_radius": blast_data or [],
            },
        )

    def test_api_endpoint_returns_graph(self):
        """GET /v1/scan/{id}/context-graph returns valid graph structure."""
        import asyncio

        from agent_bom.api.server import _get_store, app

        store = _get_store()
        job_id = "test-cg-api"
        job = self._make_job(
            job_id,
            [
                _agent(name="a1", servers=[_server(name="shared")]),
                _agent(name="a2", servers=[_server(name="shared")]),
            ],
        )
        store.put(job)

        from httpx import ASGITransport, AsyncClient

        async def _call():
            transport = ASGITransport(app=app)
            async with AsyncClient(transport=transport, base_url="http://test") as client:
                return await client.get(f"/v1/scan/{job_id}/context-graph")

        resp = asyncio.run(_call())
        assert resp.status_code == 200
        data = resp.json()
        assert "nodes" in data
        assert "stats" in data
        assert data["stats"]["agent_count"] == 2

    def test_api_agent_filter(self):
        """GET /v1/scan/{id}/context-graph?agent=a1 filters lateral paths."""
        import asyncio

        from agent_bom.api.server import _get_store, app

        store = _get_store()
        job_id = "test-cg-filter"
        job = self._make_job(
            job_id,
            [
                _agent(name="a1", servers=[_server(name="shared")]),
                _agent(name="a2", servers=[_server(name="shared")]),
            ],
        )
        store.put(job)

        from httpx import ASGITransport, AsyncClient

        async def _call():
            transport = ASGITransport(app=app)
            async with AsyncClient(transport=transport, base_url="http://test") as client:
                return await client.get(f"/v1/scan/{job_id}/context-graph?agent=a1")

        resp = asyncio.run(_call())
        assert resp.status_code == 200
        data = resp.json()
        # Lateral paths should only originate from a1
        for p in data["lateral_paths"]:
            assert p["source"] == "agent:a1"
