"""Graph analysis tools — knowledge graph queries for the pentesting agent.

Provides ``analyze_graph`` and ``query_findings`` tools that let the
agent inspect the attack knowledge graph and search through results.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from ziran.application.knowledge_graph.graph import AttackKnowledgeGraph
    from ziran.domain.entities.attack import AttackResult

logger = logging.getLogger(__name__)


def analyze_graph(
    graph: AttackKnowledgeGraph,
    *,
    include_paths: bool = True,
    include_critical_nodes: bool = True,
    max_paths: int = 10,
) -> dict[str, Any]:
    """Analyze the attack knowledge graph for attack paths and critical nodes.

    Args:
        graph: The knowledge graph to analyze.
        include_paths: Include attack path analysis.
        include_critical_nodes: Include critical node analysis.
        max_paths: Maximum number of paths to return.

    Returns:
        Structured analysis result.
    """
    from ziran.application.knowledge_graph.chain_analyzer import ToolChainAnalyzer

    result: dict[str, Any] = {
        "node_count": graph.node_count,
        "edge_count": graph.edge_count,
    }

    if include_paths:
        all_paths = graph.find_all_attack_paths()
        result["attack_paths"] = [
            {"path": list(path), "length": len(path)} for path in all_paths[:max_paths]
        ]
        result["total_attack_paths"] = len(all_paths)

    if include_critical_nodes:
        critical = graph.get_critical_nodes(top_n=10)
        result["critical_nodes"] = [
            {"node": node, "centrality": round(score, 4)} for node, score in critical
        ]

    # Run chain analysis
    try:
        analyzer = ToolChainAnalyzer(graph)
        chains = analyzer.analyze()
        result["dangerous_chains"] = [
            {
                "tools": list(c.tools),
                "vulnerability_type": c.vulnerability_type,
                "risk_score": round(c.risk_score, 4),
                "chain_type": c.chain_type,
            }
            for c in chains[:10]
        ]
    except Exception:
        logger.debug("Chain analysis failed", exc_info=True)
        result["dangerous_chains"] = []

    # Graph state summary
    state = graph.export_state()
    node_types: dict[str, int] = {}
    for node_data in state.get("nodes", []):
        if isinstance(node_data, dict):
            nt = node_data.get("node_type", "unknown")
        else:
            # Fallback for (node_id, data) tuple format
            _node_id, data = node_data
            nt = data.get("node_type", "unknown")
        node_types[nt] = node_types.get(nt, 0) + 1
    result["node_type_summary"] = node_types

    return result


def query_findings(
    attack_results: list[AttackResult],
    *,
    successful_only: bool = False,
    category: str | None = None,
    severity: str | None = None,
    phase: str | None = None,
    owasp: str | None = None,
    limit: int = 50,
) -> dict[str, Any]:
    """Search and filter attack results.

    Args:
        attack_results: All attack results from the campaign.
        successful_only: Only return successful attacks.
        category: Filter by attack category.
        severity: Filter by severity (low/medium/high/critical).
        phase: Filter by target phase.
        owasp: Filter by OWASP category (e.g. 'LLM01').
        limit: Maximum results to return.

    Returns:
        Filtered and summarized findings.
    """
    filtered = list(attack_results)

    if successful_only:
        filtered = [r for r in filtered if r.successful]

    if category:
        filtered = [r for r in filtered if r.category.value == category]

    if severity:
        filtered = [r for r in filtered if r.severity == severity]

    if owasp:
        filtered = [r for r in filtered if any(o.value == owasp for o in r.owasp_mapping)]

    total = len(filtered)
    filtered = filtered[:limit]

    return {
        "total_matching": total,
        "returned": len(filtered),
        "findings": [
            {
                "vector_id": r.vector_id,
                "vector_name": r.vector_name,
                "category": r.category.value,
                "severity": r.severity,
                "successful": r.successful,
                "owasp": [o.value for o in r.owasp_mapping],
                "evidence_keys": list(r.evidence.keys()) if r.evidence else [],
                "has_extracted_data": r.extracted_data is not None,
            }
            for r in filtered
        ],
    }
