"""LangGraph workflow construction for penetration testing."""

from typing import Literal, Any
from langgraph.graph import StateGraph, START, END
from langgraph.types import interrupt, RetryPolicy

from .state import PenTestState
from .nodes import (
    reconnaissance_search,
    coordinate_agents,
    generate_attack,
    synthesize_attack,
    execute_multi_turn,
    sandbox_enhance_attack,
    self_evaluate_attack,
    execute_attack,
    sandbox_execute_attack,
    sandbox_probe_attack,
    analyze_response,
    record_metrics,
    trigger_agent_learning,
    extract_phase_intelligence,
    summarize_if_needed,
    update_session_summary,
    calculate_score,
    generate_report,
)
from src.utils.logging import get_logger

logger = get_logger(__name__)


# ---------------------------------------------------------------------------
# Routing
# ---------------------------------------------------------------------------


def route_after_analysis(
    state: PenTestState,
) -> Literal["critical_finding", "continue_testing", "complete"]:
    """
    Route workflow after response analysis and scoring.

    Determines whether to:
    - Pause for critical finding review (human-in-the-loop)
    - Continue testing
    - Complete the test

    Matches the routing logic used by test_orchestrated.py's main loop.

    Args:
        state: Current test state

    Returns:
        Next node to execute
    """
    findings = state["security_findings"]

    # Check for critical findings – pause for human review
    critical_findings = [f for f in findings if f["severity"] == "critical"]
    if critical_findings and state["current_attempt"] >= 3:
        logger.info(
            "critical_findings_detected",
            test_session_id=state["test_session_id"],
            critical_count=len(critical_findings),
        )
        return "critical_finding"

    # Check if we've reached max attempts
    if state["current_attempt"] >= state["max_attempts"]:
        logger.info(
            "max_attempts_reached",
            test_session_id=state["test_session_id"],
            attempts=state["current_attempt"],
        )
        return "complete"

    # Continue testing
    return "continue_testing"


# ---------------------------------------------------------------------------
# Human-in-the-loop node
# ---------------------------------------------------------------------------


async def critical_finding_handler(state: PenTestState) -> dict[str, Any]:
    """
    Pause execution for critical findings requiring human review.

    Uses LangGraph's interrupt() to pause the workflow. The caller
    (CLI or API) can inspect the findings and resume with
    ``Command(resume={"continue": True})`` or stop the test with
    ``Command(resume={"continue": False})``.

    When no checkpointer is configured the interrupt is a no-op and
    the test continues automatically.

    Args:
        state: Current test state

    Returns:
        State update (may set test_status to "stopped" if reviewer declines)
    """
    critical = [f for f in state["security_findings"] if f["severity"] == "critical"]

    logger.info(
        "critical_finding_interrupt",
        test_session_id=state["test_session_id"],
        findings_count=len(critical),
    )

    # Pause for human review – yields control back to caller
    response = interrupt(
        {
            "message": "Critical vulnerability found — human review required",
            "findings_count": len(critical),
            "findings": [
                {
                    "severity": f["severity"],
                    "category": f["category"],
                    "description": f.get("description", "")[:200],
                    "confidence": f.get("confidence", 0),
                }
                for f in critical[:5]
            ],
        }
    )

    # Caller resumes with Command(resume={"continue": True/False})
    if isinstance(response, dict) and not response.get("continue", True):
        logger.info(
            "test_stopped_by_reviewer",
            test_session_id=state["test_session_id"],
        )
        return {"test_status": "stopped"}

    # Reviewer approved — continue to report generation
    return {}


# ---------------------------------------------------------------------------
# Workflow construction
# ---------------------------------------------------------------------------


def create_pentest_workflow(checkpointer=None):
    """
    Create and compile the penetration testing workflow.

    Graph topology::

        START -> reconnaissance_search -> coordinate_agents -> generate_attack
              -> synthesize_attack -> execute_multi_turn
              -> sandbox_enhance_attack -> self_evaluate_attack
              -> execute_attack -> sandbox_execute_attack -> sandbox_probe_attack
              -> analyze_response -> record_metrics
              -> [summarize_if_needed, update_session_summary, trigger_agent_learning, extract_phase_intelligence]  (parallel)
              -> calculate_score
              -> [continue_testing -> coordinate_agents | critical_finding -> handler | complete -> generate_report]
              -> END

    Args:
        checkpointer: A pre-initialized checkpointer instance.
                      - ``None`` (default): uses in-memory ``MemorySaver`` so
                        that ``interrupt()`` works out of the box.
                      - A checkpointer object (e.g. ``AsyncSqliteSaver``):
                        enables durable persistence / session resumption.
                      - ``False``: disables checkpointing entirely
                        (``interrupt()`` will **not** work).

    Returns:
        Compiled LangGraph application (CompiledGraph)

    Example:
        >>> # Simple (in-memory checkpointing, interrupt works):
        >>> app = create_pentest_workflow()
        >>>
        >>> # With SQLite persistence:
        >>> async with AsyncSqliteSaver.from_conn_string("db.sqlite") as saver:
        ...     app = create_pentest_workflow(checkpointer=saver)
    """
    logger.info("creating_pentest_workflow")

    # Create workflow
    workflow = StateGraph(PenTestState)

    # ── Add nodes (with retry policies on external-facing nodes) ──────
    workflow.add_node(
        "reconnaissance_search",
        reconnaissance_search,
        retry_policy=RetryPolicy(max_attempts=3, initial_interval=2.0, backoff_factor=2.0),
    )
    workflow.add_node("coordinate_agents", coordinate_agents)
    workflow.add_node("generate_attack", generate_attack)
    workflow.add_node("synthesize_attack", synthesize_attack)
    workflow.add_node("execute_multi_turn", execute_multi_turn)
    workflow.add_node(
        "sandbox_enhance_attack",
        sandbox_enhance_attack,
        retry_policy=RetryPolicy(max_attempts=1, initial_interval=1.0),
    )
    workflow.add_node(
        "self_evaluate_attack",
        self_evaluate_attack,
        retry_policy=RetryPolicy(max_attempts=1, initial_interval=1.0),
    )
    workflow.add_node(
        "execute_attack",
        execute_attack,
        retry_policy=RetryPolicy(max_attempts=2, initial_interval=1.0, backoff_factor=2.0),
    )
    workflow.add_node(
        "sandbox_execute_attack",
        sandbox_execute_attack,
        retry_policy=RetryPolicy(max_attempts=1, initial_interval=1.0),
    )
    workflow.add_node(
        "sandbox_probe_attack",
        sandbox_probe_attack,
        retry_policy=RetryPolicy(max_attempts=1, initial_interval=1.0),
    )
    workflow.add_node("analyze_response", analyze_response)
    workflow.add_node("record_metrics", record_metrics)
    workflow.add_node("trigger_agent_learning", trigger_agent_learning)
    workflow.add_node("extract_phase_intelligence", extract_phase_intelligence)
    workflow.add_node("summarize_if_needed", summarize_if_needed)
    workflow.add_node("update_session_summary", update_session_summary)
    workflow.add_node("calculate_score", calculate_score)
    workflow.add_node("critical_finding_handler", critical_finding_handler)
    workflow.add_node(
        "generate_report",
        generate_report,
        retry_policy=RetryPolicy(max_attempts=2, initial_interval=1.0),
    )

    # ── Build edges ────────────────────────────────────────────────────
    workflow.add_edge(START, "reconnaissance_search")
    workflow.add_edge("reconnaissance_search", "coordinate_agents")
    workflow.add_edge("coordinate_agents", "generate_attack")
    workflow.add_edge("generate_attack", "synthesize_attack")
    workflow.add_edge("synthesize_attack", "execute_multi_turn")
    workflow.add_edge("execute_multi_turn", "sandbox_enhance_attack")
    workflow.add_edge("sandbox_enhance_attack", "self_evaluate_attack")
    workflow.add_edge("self_evaluate_attack", "execute_attack")
    workflow.add_edge("execute_attack", "sandbox_execute_attack")
    workflow.add_edge("sandbox_execute_attack", "sandbox_probe_attack")
    workflow.add_edge("sandbox_probe_attack", "analyze_response")
    workflow.add_edge("analyze_response", "record_metrics")

    # ── Parallel fan-out: record_metrics -> four independent nodes ──────
    # summarize_if_needed        -> updates conversation_history
    # update_session_summary     -> updates session_summary
    # trigger_agent_learning     -> updates agent_learnings
    # extract_phase_intelligence -> updates phase_intelligence
    # They write to different state keys, so no conflicts.
    workflow.add_edge("record_metrics", "summarize_if_needed")
    workflow.add_edge("record_metrics", "update_session_summary")
    workflow.add_edge("record_metrics", "trigger_agent_learning")
    workflow.add_edge("record_metrics", "extract_phase_intelligence")

    # ── Fan-in: all four feed into calculate_score ───────────────────
    workflow.add_edge("summarize_if_needed", "calculate_score")
    workflow.add_edge("update_session_summary", "calculate_score")
    workflow.add_edge("trigger_agent_learning", "calculate_score")
    workflow.add_edge("extract_phase_intelligence", "calculate_score")

    # ── Conditional routing after scoring ──────────────────────────────
    workflow.add_conditional_edges(
        "calculate_score",
        route_after_analysis,
        {
            "critical_finding": "critical_finding_handler",
            "continue_testing": "coordinate_agents",  # Loop back
            "complete": "generate_report",
        },
    )

    workflow.add_edge("critical_finding_handler", "generate_report")
    workflow.add_edge("generate_report", END)

    # ── Compile ────────────────────────────────────────────────────────
    if checkpointer is False:
        # Explicitly disabled — interrupt() will not work
        app = workflow.compile()
        logger.info("pentest_workflow_created", checkpoint_enabled=False)
    else:
        if checkpointer is None:
            # Default: in-memory checkpointer (interrupt works, no persistence)
            from langgraph.checkpoint.memory import MemorySaver

            checkpointer = MemorySaver()
            logger.info("pentest_workflow_created", checkpoint_type="MemorySaver")
        else:
            logger.info("pentest_workflow_created", checkpoint_type=type(checkpointer).__name__)
        app = workflow.compile(checkpointer=checkpointer)

    return app
