"""
Test runner for executing security tests via CLI.

Phase 5 switchover: Uses LangGraph workflow instead of test_orchestrated.py.
"""

import asyncio
import sys
from pathlib import Path

# Add root directory to path (still needed for configs)
ROOT_DIR = Path(__file__).parent.parent.parent
sys.path.insert(0, str(ROOT_DIR))

from datetime import UTC  # noqa: E402

from rich import box  # noqa: E402
from rich.console import Console  # noqa: E402
from rich.panel import Panel  # noqa: E402
from rich.table import Table  # noqa: E402

from src.cli.config_loader import load_config  # noqa: E402
from src.connectors.factory import create_connector  # noqa: E402

console = Console()


async def run_test(args):
    """
    Execute full orchestrated test using LangGraph workflow.

    Enhanced with:
    - --agents: Select specific agents
    - --dry-run: Preview without executing
    - --output: Specify report output
    - --phase: Start at specific phase
    - --verbose: Detailed output
    - --machine: JSONL streaming output for AI agents
    """
    from src.cli.machine_output import emit_error, emit_jsonl, is_machine

    verbose = getattr(args, "verbose", False)
    machine = is_machine(args)

    # ═══════════════════════════════════════════════════════════════════════════
    # Load Configuration
    # ═══════════════════════════════════════════════════════════════════════════
    if not machine:
        console.print()
        console.print(
            Panel.fit(
                "[bold cyan]PenBot Security Testing Framework[/bold cyan]\n"
                "[dim]Multi-agent adversarial testing for AI chatbots[/dim]",
                border_style="cyan",
            )
        )
        console.print()
        console.print(f"📂 Loading config: [cyan]{args.config}[/cyan]")

    try:
        config = load_config(args.config)
    except FileNotFoundError:
        if machine:
            emit_error("config_not_found", f"Config file not found: {args.config}")
        else:
            console.print(f"[red]❌ Config file not found: {args.config}[/red]")
        return
    except Exception as e:
        if machine:
            emit_error("config_error", str(e))
        else:
            console.print(f"[red]❌ Failed to load config: {e}[/red]")
        return

    target_name = config["target"]["name"]
    if not machine:
        console.print(f"✅ Target: [green]{target_name}[/green]")

    # ═══════════════════════════════════════════════════════════════════════════
    # Agent Selection
    # ═══════════════════════════════════════════════════════════════════════════
    if args.agents:
        from src.cli.agents_cmd import validate_agent_list

        valid_agents, invalid_agents = validate_agent_list(args.agents)

        if invalid_agents and not machine:
            console.print(f"[yellow]⚠️  Unknown agents: {', '.join(invalid_agents)}[/yellow]")

        if not valid_agents:
            if machine:
                emit_error("no_valid_agents", "No valid agents specified")
            else:
                console.print("[red]❌ No valid agents specified[/red]")
                console.print("[dim]Use 'penbot agents' to see available agents[/dim]")
            return

        config["test"]["agents"] = valid_agents
        if not machine:
            console.print(f"🤖 Agents: [cyan]{', '.join(valid_agents)}[/cyan]")
    else:
        agents_in_config = config["test"].get("agents", [])
        if not machine:
            console.print(
                f"🤖 Agents: [cyan]{', '.join(agents_in_config) if agents_in_config else 'All available'}[/cyan]"
            )

    # ═══════════════════════════════════════════════════════════════════════════
    # Determine Attack Count
    # ═══════════════════════════════════════════════════════════════════════════
    max_attacks = args.max_attacks
    if not max_attacks:
        max_attacks = config["test"].get("max_attacks", 30)

    if args.quick:
        max_attacks = 3
        if not machine:
            console.print(f"⚡ Quick mode: [yellow]{max_attacks} attacks[/yellow]")
    elif not machine:
        console.print(f"🎯 Max Attacks: [cyan]{max_attacks}[/cyan]")

    start_phase = getattr(args, "phase", None)
    if start_phase and not machine:
        console.print(f"📍 Starting Phase: [cyan]{start_phase}[/cyan]")

    output_dir = Path(args.output) if getattr(args, "output", None) else Path("reports")
    output_dir.mkdir(parents=True, exist_ok=True)

    if not machine:
        console.print(f"📄 Reports: [cyan]{output_dir}[/cyan]")
        console.print()

    # ═══════════════════════════════════════════════════════════════════════════
    # Dry Run Mode
    # ═══════════════════════════════════════════════════════════════════════════
    if getattr(args, "dry_run", False):
        await _dry_run(config, max_attacks)
        return

    # ═══════════════════════════════════════════════════════════════════════════
    # Initialize Connector (unified factory)
    # ═══════════════════════════════════════════════════════════════════════════
    conn_type = config["target"]["connection"].get("type", "rest")
    target_type = "playwright" if conn_type == "playwright" else "api"

    if not machine:
        if target_type == "playwright":
            console.print("🌐 Using Playwright browser automation")
        else:
            console.print("🔌 Using REST API connector")

    connector = create_connector(target_type, config["target"])

    if machine:
        await connector.initialize()
    else:
        with console.status("[bold green]Initializing connector..."):
            await connector.initialize()
        console.print("✅ Connector initialized")
        console.print()

    # ═══════════════════════════════════════════════════════════════════════════
    # Test Configuration Summary
    # ═══════════════════════════════════════════════════════════════════════════
    if not machine:
        summary_table = Table(box=box.ROUNDED, show_header=False, title="📋 Test Configuration")
        summary_table.add_column("Setting", style="bold")
        summary_table.add_column("Value", style="cyan")

        summary_table.add_row("Target", target_name)
        summary_table.add_row("Platform", config["target"].get("platform", "unknown"))
        summary_table.add_row("Max Attacks", str(max_attacks))
        summary_table.add_row("Agents", str(len(config["test"].get("agents", []))) + " active")
        summary_table.add_row("Engine", "LangGraph")
        if start_phase:
            summary_table.add_row("Start Phase", start_phase)

        console.print(summary_table)
        console.print()

    # ═══════════════════════════════════════════════════════════════════════════
    # Execute Campaign via LangGraph
    # ═══════════════════════════════════════════════════════════════════════════
    from langgraph.types import Command

    from src.utils.helpers import generate_uuid
    from src.workflow.graph import create_pentest_workflow
    from src.workflow.state import create_initial_state
    from src.workflow.streaming import print_node_output, stream_workflow_to_websocket

    if not machine:
        console.print("🚀 [bold green]Starting LangGraph campaign...[/bold green]")
        console.print()

    # Start operational metrics tracking
    from src.utils.config import settings as app_settings
    from src.utils.operational_metrics import get_metrics

    metrics = get_metrics()
    metrics.start_session()

    if app_settings.inter_round_delay > 0 and not machine:
        console.print(f"⏱  Inter-round delay: [cyan]{app_settings.inter_round_delay}s[/cyan]")

    session_id = generate_uuid()

    if machine:
        emit_jsonl(
            "test_started",
            {
                "session_id": session_id,
                "target_name": target_name,
                "max_attacks": max_attacks,
                "connector_type": target_type,
                "agents": config["test"].get("agents", []),
            },
        )

    # Build target config with vision/RAG/agentic flags
    target_config_enriched = {
        **config["target"].get("connection", {}),
        "name": target_name,
        "endpoint": config["target"].get("connection", {}).get("endpoint", ""),
        "supports_images": config["target"].get("supports_images", False),
        "image_format": config["target"].get("image_format", "openai"),
        "is_rag": config["target"].get("is_rag", False),
        "is_agentic": config["target"].get("is_agentic", False),
    }

    # Create initial state
    initial_state = create_initial_state(
        target_name=target_name,
        target_type=target_type,
        target_config=target_config_enriched,
        attack_group=config["test"].get("attack_group", "prompt_engineering"),
        max_attempts=max_attacks,
        attack_graph_enabled=True,
        session_id=session_id,
    )

    # ── Initialize checkpointer (SQLite for persistence) ──────────────
    # AsyncSqliteSaver.from_conn_string() is an async context manager in
    # langgraph-checkpoint-sqlite v3.x — must be entered before use.
    from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver

    async with AsyncSqliteSaver.from_conn_string("pentest_sessions.db") as checkpointer:
        # Create LangGraph workflow with durable SQLite checkpointer
        app = create_pentest_workflow(checkpointer=checkpointer)

        # Runtime config — carries non-serializable objects via configurable
        # These objects persist across all node calls within a single execution.
        langgraph_config = {
            "recursion_limit": max(max_attacks * 14 + 20, 250),
            "configurable": {
                "thread_id": session_id,
                "connector": connector,
                # Runtime objects (lazy-initialized by nodes on first access):
                # "canned_detector": <created by execute_attack>
                # "target_fingerprinter": <created by execute_attack>
                # "attack_graph": <created by coordinate_agents>
                # "attack_memory": <created by trigger_agent_learning>
            },
        }

        # Track the latest state snapshot for graceful interrupt saving
        last_state_snapshot: dict = {}

        try:
            # Emit test_started WebSocket event
            try:
                from src.workflow.streaming import emit_ws_event_http

                await emit_ws_event_http(
                    session_id,
                    "test_started",
                    {
                        "target_name": target_name,
                        "max_attacks": max_attacks,
                        "attack_group": config["test"].get("attack_group", "prompt_engineering"),
                        "agents": list(config["test"].get("agents", [])),
                        "session_id": session_id,
                    },
                )
            except Exception:
                pass  # Dashboard may not be running

            # Stream execution with interrupt handling for critical findings
            input_value = initial_state
            last_round_completed = 0  # Track for inter-round delay

            while True:
                has_interrupt = False

                async for chunk in app.astream(
                    input_value, langgraph_config, stream_mode="updates"
                ):
                    # Check for interrupt (critical findings pause)
                    if "__interrupt__" in chunk:
                        has_interrupt = True
                        interrupt_data = chunk["__interrupt__"]
                        console.print()
                        console.print(
                            "[bold yellow]⚠️  Critical vulnerability detected — "
                            "human review requested[/bold yellow]"
                        )
                        for item in interrupt_data:
                            value = item.value if hasattr(item, "value") else item
                            if isinstance(value, dict):
                                console.print(
                                    f"  Findings: {value.get('findings_count', '?')} critical"
                                )
                        console.print(
                            "[dim]Auto-continuing (use API/dashboard for manual review)[/dim]"
                        )
                        break

                    # Normal node output: chunk is {node_name: state_update}
                    for node_name, state_update in chunk.items():
                        if state_update is None:
                            continue

                        last_state_snapshot.update(state_update)

                        round_num = state_update.get(
                            "current_attempt",
                            initial_state.get("current_attempt", 0),
                        )

                        if machine:
                            # Emit JSONL event per node
                            emit_jsonl(
                                "node_output",
                                {
                                    "node": node_name,
                                    "round": round_num,
                                    "session_id": session_id,
                                    "data": _extract_machine_data(node_name, state_update),
                                },
                            )
                        else:
                            if node_name == "generate_attack":
                                if last_round_completed > 0 and app_settings.inter_round_delay > 0:
                                    await asyncio.sleep(app_settings.inter_round_delay)

                                max_attacks = initial_state.get("max_attempts", "?")
                                console.print()
                                console.print(f"[bold cyan]{'─' * 60}[/bold cyan]")
                                console.print(
                                    f"[bold cyan]  ⚔  ROUND {round_num}/{max_attacks}[/bold cyan]"
                                )
                                console.print(f"[bold cyan]{'─' * 60}[/bold cyan]")
                                last_round_completed = round_num

                            print_node_output(node_name, state_update)

                        # WebSocket streaming (regardless of mode)
                        try:
                            await stream_workflow_to_websocket(session_id, node_name, state_update)
                        except Exception:
                            pass

                if has_interrupt:
                    # Resume after interrupt — auto-continue in CLI mode
                    input_value = Command(resume={"continue": True})
                else:
                    break  # Graph completed normally

            if machine:
                emit_jsonl(
                    "test_completed",
                    {
                        "session_id": session_id,
                        "rounds": last_state_snapshot.get("current_attempt", 0),
                        "findings": len(last_state_snapshot.get("security_findings", [])),
                        "vulnerability_score": last_state_snapshot.get("vulnerability_score", 0),
                    },
                )
            else:
                console.print()
                console.print("[bold green]✅ Campaign completed successfully.[/bold green]")

        except (KeyboardInterrupt, asyncio.CancelledError):
            if not machine:
                console.print()
                console.print("[yellow]⚠️  Campaign interrupted by user.[/yellow]")

            # Cancel lingering tasks (MCP subprocesses, etc.) to avoid
            # "Attempted to exit cancel scope in a different task" tracebacks.
            for task in asyncio.all_tasks():
                if task is not asyncio.current_task() and not task.done():
                    task.cancel()
            await asyncio.sleep(0.1)

            # Graceful interrupt: Save partial session state
            if last_state_snapshot:
                try:
                    import json
                    from datetime import datetime

                    sessions_dir = Path("sessions")
                    sessions_dir.mkdir(exist_ok=True)
                    interrupt_path = sessions_dir / f"{session_id}_interrupted.json"

                    interrupt_data = {
                        "session_id": session_id,
                        "target_name": target_name,
                        "status": "interrupted",
                        "interrupted_at": datetime.now(UTC).isoformat(),
                        "rounds_completed": last_state_snapshot.get("current_attempt", 0),
                        "security_findings": last_state_snapshot.get("security_findings", []),
                        "vulnerability_score": last_state_snapshot.get("vulnerability_score", 0),
                        "attack_attempts_count": len(
                            last_state_snapshot.get("attack_attempts", [])
                        ),
                    }

                    def _json_serial(obj):
                        if hasattr(obj, "isoformat"):
                            return obj.isoformat()
                        if hasattr(obj, "model_dump"):
                            return obj.model_dump()
                        if hasattr(obj, "dict"):
                            return obj.dict()
                        return str(obj)

                    with open(interrupt_path, "w", encoding="utf-8") as f:
                        json.dump(interrupt_data, f, default=_json_serial, indent=2)

                    console.print(f"[dim]Partial session saved: {interrupt_path}[/dim]")
                except Exception as save_err:
                    console.print(f"[dim]Could not save partial state: {save_err}[/dim]")

        except Exception as e:
            if machine:
                emit_error("campaign_failed", str(e))
            else:
                console.print()
                console.print(f"[red]❌ Campaign failed: {e}[/red]")
            if verbose:
                import traceback

                traceback.print_exc()
        finally:
            await connector.close()
            # Allow background checkpoint writes to flush before
            # the async-with block closes the SQLite connection.
            await asyncio.sleep(0.5)

    # ═══════════════════════════════════════════════════════════════════════════
    # Post-Test Summary
    # ═══════════════════════════════════════════════════════════════════════════
    if not machine:
        console.print()
        console.print("[dim]Use 'penbot sessions' to view past sessions[/dim]")
        console.print("[dim]Use 'penbot report --latest' to generate a report[/dim]")


async def _dry_run(config: dict, max_attacks: int):
    """
    Preview what would be tested without actually executing.

    Shows:
    - Target configuration
    - Active agents
    - Sample attack patterns
    """
    console.print()
    console.print(
        Panel.fit(
            "[bold yellow]🔍 DRY RUN MODE[/bold yellow]\n"
            "[dim]Preview only - no attacks will be executed[/dim]",
            border_style="yellow",
        )
    )
    console.print()

    # Show target info
    target = config["target"]
    console.print("[bold]Target Configuration:[/bold]")
    console.print(f"  • Name: {target.get('name', 'Unknown')}")
    console.print(f"  • Platform: {target.get('platform', 'Unknown')}")
    console.print(f"  • Endpoint: {target.get('connection', {}).get('endpoint', 'Unknown')}")
    console.print()

    # Show agents
    agents = config["test"].get("agents", [])
    console.print("[bold]Active Agents:[/bold]")

    if agents:
        from src.cli.agents_cmd import AGENT_REGISTRY

        for agent_id in agents:
            if agent_id in AGENT_REGISTRY:
                info = AGENT_REGISTRY[agent_id]
                console.print(f"  • [cyan]{agent_id}[/cyan] - {info['description'][:50]}...")
            else:
                console.print(f"  • [cyan]{agent_id}[/cyan]")
    else:
        console.print("  [dim]All available agents will be used[/dim]")

    console.print()

    # Show sample patterns
    console.print("[bold]Sample Attack Patterns:[/bold]")
    console.print()

    try:
        from src.agents.jailbreak import JailbreakAgent

        agent = JailbreakAgent(llm_client=None, config={})
        patterns = agent.get_attack_patterns()[:5]  # First 5

        for i, pattern in enumerate(patterns, 1):
            name = pattern.get("name", "unknown")
            desc = pattern.get("description", "No description")[:60]
            severity = pattern.get("severity_if_success", "unknown")
            console.print(f"  {i}. [cyan]{name}[/cyan]")
            console.print(f"     {desc}")
            console.print(f"     Severity: {severity}")
            console.print()

    except Exception as e:
        console.print(f"  [dim]Could not load patterns: {e}[/dim]")

    # Estimated duration
    # Each round includes: agent coordination (~10-15s), subagent pipeline (~10-15s),
    # attack execution (~2-5s), response analysis (~2-5s), optional Think-MCP (~10-20s).
    # Realistic average: ~45-60 seconds per attack.
    avg_time_per_attack = 50  # seconds (conservative middle estimate)
    total_time = max_attacks * avg_time_per_attack
    hours = total_time // 3600
    minutes = (total_time % 3600) // 60
    seconds = total_time % 60

    console.print()
    if hours > 0:
        console.print(
            f"[bold]Estimated Duration:[/bold] ~{hours}h {minutes}m for {max_attacks} attacks"
        )
    else:
        console.print(
            f"[bold]Estimated Duration:[/bold] ~{minutes}m {seconds}s for {max_attacks} attacks"
        )
    console.print()
    console.print("[dim]To execute, run without --dry-run flag[/dim]")


def _extract_machine_data(node_name: str, state_update: dict) -> dict:
    """Extract a concise, serializable subset from a node's state update."""
    data: dict = {}

    if node_name == "generate_attack":
        attacks = state_update.get("attack_attempts", [])
        if attacks:
            last = attacks[-1] if isinstance(attacks, list) else attacks
            data["attack"] = {
                "agent": last.get("agent_name", ""),
                "type": last.get("attack_type", ""),
                "query_preview": str(last.get("query", ""))[:200],
            }

    elif node_name == "execute_attack":
        responses = state_update.get("target_responses", [])
        if responses:
            last = responses[-1] if isinstance(responses, list) else responses
            data["response_preview"] = str(last.get("content", ""))[:300]

    elif node_name == "extract_findings":
        findings = state_update.get("security_findings", [])
        data["new_findings"] = [
            {
                "severity": f.get("severity", ""),
                "category": f.get("category", ""),
                "description": f.get("description", ""),
                "confidence": f.get("confidence", 0),
            }
            for f in (findings if isinstance(findings, list) else [findings])
        ]

    elif node_name == "coordinate_agents":
        consultation = state_update.get("agent_consultation", {})
        if isinstance(consultation, dict):
            data["consensus"] = consultation.get("consensus", {})

    else:
        for key in ("vulnerability_score", "test_status", "campaign_phase"):
            if key in state_update:
                data[key] = state_update[key]

    return data
