"""LangGraph-based pentesting agent.

Defines the ``PentestAgent`` — a stateful agent built on LangGraph's
``StateGraph`` that autonomously plans, executes, and adapts security
testing campaigns against AI agents.

The agent uses a plan-execute-reason loop:
1. **Planner** — analyzes target and creates/updates attack plan
2. **Executor** — calls the appropriate tool based on the plan
3. **Reasoner** — analyzes results and decides next steps

Requires the ``pentest`` optional extra: ``pip install ziran[pentest]``.
"""

from __future__ import annotations

import json
import logging
import time
from typing import TYPE_CHECKING, Any, Literal, TypedDict

if TYPE_CHECKING:
    from collections.abc import Callable

    from ziran.application.agent_scanner.scanner import AgentScanner, ProgressEvent
    from ziran.domain.entities.attack import AttackResult
    from ziran.infrastructure.llm.base import BaseLLMClient

from ziran.application.pentesting.prompts import (
    PLANNER_PROMPT,
    REASONER_PROMPT,
    REPORTER_PROMPT,
    SYSTEM_PROMPT,
)
from ziran.domain.entities.pentest import (
    PentestPlan,
    PentestSession,
    PentestStatus,
    PentestStep,
    PentestStepType,
)

logger = logging.getLogger(__name__)


class AgentState(TypedDict, total=False):
    """State managed by the LangGraph pentesting agent."""

    messages: list[dict[str, str]]
    session: PentestSession
    scanner: Any  # AgentScanner — not serializable, set at init
    llm_client: Any  # BaseLLMClient
    attack_results: list[Any]  # list[AttackResult]
    next_action: str
    should_stop: bool
    error: str | None


class PentestAgent:
    """Autonomous pentesting agent built on LangGraph.

    Wraps an ``AgentScanner`` and uses an LLM backbone to
    plan, execute, and adapt multi-phase penetration testing
    campaigns against AI agents.

    Example::

        agent = PentestAgent(scanner=scanner, llm_client=llm_client)
        session = await agent.run(goal="Find prompt injection vulns")
    """

    def __init__(
        self,
        scanner: AgentScanner,
        llm_client: BaseLLMClient,
        *,
        max_iterations: int = 10,
        on_progress: Callable[[ProgressEvent], None] | None = None,
        on_agent_event: Callable[[str, str], None] | None = None,
    ) -> None:
        self.scanner = scanner
        self.llm_client = llm_client
        self.max_iterations = max_iterations
        self._attack_results: list[AttackResult] = []
        self._on_progress = on_progress
        self._on_agent_event = on_agent_event
        self._tested_vectors: set[str] = set()
        self._vuln_history: list[int] = []
        self._graph = self._build_graph()

    def _build_graph(self) -> Any:
        """Build the LangGraph StateGraph.

        Returns the compiled graph. Defers import of langgraph
        so the module can be imported without the 'pentest' extra.
        """
        try:
            from langgraph.graph import END, StateGraph
        except ImportError as e:
            raise ImportError(
                "The pentesting agent requires the 'pentest' extra. "
                "Install with: pip install ziran[pentest]"
            ) from e

        graph = StateGraph(AgentState)

        # Add nodes
        graph.add_node("planner", self._planner_node)
        graph.add_node("executor", self._executor_node)
        graph.add_node("reasoner", self._reasoner_node)
        graph.add_node("reporter", self._reporter_node)

        # Set entry point
        graph.set_entry_point("planner")

        # Add edges
        graph.add_edge("planner", "executor")
        graph.add_conditional_edges(
            "executor",
            self._route_after_executor,
            {"reasoner": "reasoner", "reporter": "reporter"},
        )
        graph.add_conditional_edges(
            "reasoner",
            self._route_after_reasoner,
            {"planner": "planner", "reporter": "reporter"},
        )
        graph.add_edge("reporter", END)

        return graph.compile()

    def _route_after_executor(self, state: AgentState) -> Literal["reasoner", "reporter"]:
        """Route after tool execution."""
        session = state["session"]
        if session.status == PentestStatus.FAILED:
            return "reporter"
        if state.get("should_stop", False):
            return "reporter"
        return "reasoner"

    def _route_after_reasoner(self, state: AgentState) -> Literal["planner", "reporter"]:
        """Route after reasoning."""
        session = state["session"]
        if state.get("should_stop", False):
            return "reporter"
        if session.iteration_count >= session.max_iterations:
            return "reporter"
        return "planner"

    async def _planner_node(self, state: AgentState) -> AgentState:
        """Plan the next pentesting action using the LLM."""
        session = state["session"]
        start = time.monotonic()

        # Build context for the planner
        steps_summary = (
            "; ".join(f"[{s.step_type.value}] {s.result_summary}" for s in session.steps[-5:])
            or "No steps yet"
        )

        capabilities = [
            cap.get("id", "?")
            for cap in (state.get("attack_results") or [])
            if isinstance(cap, dict) and cap.get("type") == "capability"
        ]
        if not capabilities:
            # Extract capability nodes from the knowledge graph
            cap_list = [
                node_id
                for node_id, data in self.scanner.graph.graph.nodes(data=True)
                if data.get("node_type") == "capability"
            ]
            capabilities = cap_list if cap_list else ["unknown"]

        graph_summary = (
            f"{self.scanner.graph.node_count} nodes, {self.scanner.graph.edge_count} edges"
        )

        prompt = PLANNER_PROMPT.format(
            goal=session.plan.goal if session.plan else "not set",
            iteration_count=session.iteration_count,
            max_iterations=session.max_iterations,
            total_vulnerabilities=session.total_vulnerabilities,
            total_attacks_executed=session.total_attacks_executed,
            tested_vector_count=len(self._tested_vectors),
            steps_summary=steps_summary,
            capabilities=", ".join(capabilities[:20]),
            graph_summary=graph_summary,
        )

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ]

        response = await self.llm_client.complete(messages=messages, temperature=0.3)

        # Parse the LLM's decision into a tool action
        action = self._parse_action(response.content)
        duration = time.monotonic() - start

        step = PentestStep(
            step_type=PentestStepType.PLAN,
            reasoning=response.content[:500],
            tool_name=action.get("tool", ""),
            tool_input=action.get("args", {}),
            result_summary=f"Plan: {action.get('tool', 'unknown')}",
            duration_seconds=duration,
            token_usage={
                "prompt_tokens": response.prompt_tokens,
                "completion_tokens": response.completion_tokens,
                "total_tokens": response.total_tokens,
            },
        )
        session.add_step(step)
        session.iteration_count += 1

        state["session"] = session
        state["next_action"] = json.dumps(action)

        self._emit_agent_event(
            "plan",
            f"Iteration {session.iteration_count}/{session.max_iterations} — "
            f"next: {action.get('tool', 'unknown')}"
            + (f" ({action.get('args', {})})" if action.get("args") else ""),
        )

        return state

    async def _executor_node(self, state: AgentState) -> AgentState:
        """Execute the tool selected by the planner."""
        session = state["session"]
        start = time.monotonic()
        action = json.loads(state.get("next_action", "{}"))
        tool_name = action.get("tool", "run_scan")
        tool_args = action.get("args", {})

        try:
            result = await self._execute_tool(tool_name, tool_args)
            result_summary = self._summarize_result(tool_name, result)

            step = PentestStep(
                step_type=PentestStepType(tool_name)
                if tool_name in PentestStepType.__members__.values()
                else PentestStepType.SCAN,
                reasoning=f"Executing {tool_name}",
                tool_name=tool_name,
                tool_input=tool_args,
                result_summary=result_summary,
                result_data=result if isinstance(result, dict) else {"output": str(result)},
                duration_seconds=time.monotonic() - start,
            )
            session.add_step(step)

            # Update vulnerability count from scan results
            if isinstance(result, dict) and "total_vulnerabilities" in result:
                session.total_vulnerabilities = max(
                    session.total_vulnerabilities, result["total_vulnerabilities"]
                )
            if isinstance(result, dict) and "total_attacks_executed" in result:
                session.total_attacks_executed += result.get("attacks_executed", 0)

            # Track tested vector IDs to avoid re-running them
            if isinstance(result, dict):
                for sa in result.get("successful_attacks", []):
                    vid = sa.get("vector_id", "")
                    if vid:
                        self._tested_vectors.add(vid)
                # Also track all executed (not just successful)
                for ar in self.scanner._attack_results:
                    self._tested_vectors.add(ar.vector_id)

        except Exception as e:
            logger.error("Tool execution failed: %s", e)
            step = PentestStep(
                step_type=PentestStepType.SCAN,
                reasoning=f"Failed to execute {tool_name}: {e}",
                tool_name=tool_name,
                tool_input=tool_args,
                result_summary=f"Error: {e}",
                duration_seconds=time.monotonic() - start,
            )
            session.add_step(step)
            state["error"] = str(e)

        self._emit_agent_event("execute", step.result_summary)

        state["session"] = session
        return state

    async def _reasoner_node(self, state: AgentState) -> AgentState:
        """Analyze results and decide next steps."""
        session = state["session"]
        start = time.monotonic()

        last_step = session.steps[-1] if session.steps else None
        last_result = last_step.result_summary if last_step else "No results"
        last_type = last_step.step_type.value if last_step else "none"

        from ziran.application.pentesting.tools.graph import analyze_graph

        graph_analysis = analyze_graph(self.scanner.graph, max_paths=5)
        chain_count = len(graph_analysis.get("dangerous_chains", []))
        graph_summary = (
            f"{graph_analysis.get('node_count', 0)} nodes, "
            f"{graph_analysis.get('edge_count', 0)} edges, "
            f"{graph_analysis.get('total_attack_paths', 0)} attack paths"
        )

        # Build vulnerability trend for plateau detection
        self._vuln_history.append(session.total_vulnerabilities)
        if len(self._vuln_history) >= 3:
            recent = self._vuln_history[-3:]
            vuln_trend = f"{' → '.join(str(v) for v in recent)} ({'plateau — no new vulns' if recent[-1] == recent[-2] else 'increasing'})"
        elif len(self._vuln_history) >= 2:
            recent = self._vuln_history[-2:]
            vuln_trend = f"{' → '.join(str(v) for v in recent)}"
        else:
            vuln_trend = str(session.total_vulnerabilities)

        prompt = REASONER_PROMPT.format(
            last_step_type=last_type,
            last_result_summary=last_result,
            total_vulnerabilities=session.total_vulnerabilities,
            iteration_count=session.iteration_count,
            max_iterations=session.max_iterations,
            chain_count=chain_count,
            graph_summary=graph_summary,
            vuln_trend=vuln_trend,
        )

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ]

        response = await self.llm_client.complete(messages=messages, temperature=0.3)

        # Check if the LLM thinks we should stop
        content_lower = response.content.lower()
        should_stop = any(
            phrase in content_lower
            for phrase in ["session complete", "conclude", "sufficient findings", "goal met"]
        )

        # Programmatic plateau detection: if vuln count unchanged for
        # 2+ consecutive reasoning cycles, force stop to save resources
        if len(self._vuln_history) >= 3 and not should_stop:
            recent = self._vuln_history[-3:]
            if recent[0] == recent[1] == recent[2] and recent[0] > 0:
                logger.info(
                    "Plateau detected: vulnerability count unchanged at %d for 3 cycles",
                    recent[0],
                )
                should_stop = True

        # Parse next action from reasoning
        action = self._parse_action(response.content)

        step = PentestStep(
            step_type=PentestStepType.CHAIN_REASONING,
            reasoning=response.content[:500],
            tool_name=action.get("tool", ""),
            result_summary=f"Reasoning: {'stop' if should_stop else action.get('tool', 'continue')}",
            duration_seconds=time.monotonic() - start,
            token_usage={
                "prompt_tokens": response.prompt_tokens,
                "completion_tokens": response.completion_tokens,
                "total_tokens": response.total_tokens,
            },
        )
        session.add_step(step)

        self._emit_agent_event(
            "reason",
            f"{'Stopping — sufficient findings' if should_stop else 'Continuing — ' + action.get('tool', 'next')}",
        )

        state["session"] = session
        state["should_stop"] = should_stop
        state["next_action"] = json.dumps(action)
        return state

    async def _reporter_node(self, state: AgentState) -> AgentState:
        """Generate the final session report."""
        session = state["session"]
        start = time.monotonic()

        findings_summary = (
            "\n".join(
                f"- [{s.step_type.value}] {s.result_summary}"
                for s in session.steps
                if s.result_data.get("total_vulnerabilities", 0) > 0
                or "successful" in s.result_summary.lower()
            )
            or "No vulnerabilities found."
        )

        from ziran.application.pentesting.tools.graph import analyze_graph

        graph_analysis = analyze_graph(self.scanner.graph, max_paths=5)
        chains_summary = (
            "\n".join(
                f"- {c['tools']} → {c['vulnerability_type']} (risk: {c['risk_score']})"
                for c in graph_analysis.get("dangerous_chains", [])
            )
            or "No chains discovered."
        )

        prompt = REPORTER_PROMPT.format(
            goal=session.plan.goal if session.plan else "not set",
            iteration_count=session.iteration_count,
            total_vulnerabilities=session.total_vulnerabilities,
            total_attacks_executed=session.total_attacks_executed,
            findings_summary=findings_summary,
            chains_summary=chains_summary,
        )

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ]

        response = await self.llm_client.complete(messages=messages, temperature=0.2)

        step = PentestStep(
            step_type=PentestStepType.PLAN,
            reasoning="Generating final report",
            result_summary=response.content[:500],
            duration_seconds=time.monotonic() - start,
            token_usage={
                "prompt_tokens": response.prompt_tokens,
                "completion_tokens": response.completion_tokens,
                "total_tokens": response.total_tokens,
            },
        )
        session.add_step(step)
        session.status = PentestStatus.COMPLETE

        state["session"] = session
        return state

    async def _execute_tool(self, tool_name: str, args: dict[str, Any]) -> dict[str, Any]:
        """Dispatch a tool call to the appropriate function."""
        from ziran.application.pentesting.tools.explain import (
            explain_finding,
            suggest_remediation,
        )
        from ziran.application.pentesting.tools.external import read_target_docs, web_search
        from ziran.application.pentesting.tools.graph import analyze_graph, query_findings
        from ziran.application.pentesting.tools.scan import run_phase, run_scan
        from ziran.application.pentesting.tools.vectors import generate_vectors

        if tool_name in ("run_scan", "scan"):
            return await run_scan(
                self.scanner,
                on_progress=self._on_progress,
                exclude_vectors=self._tested_vectors or None,
                **args,
            )
        elif tool_name in ("run_phase", "phase"):
            if "phase" not in args:
                args["phase"] = "vulnerability_discovery"
            return await run_phase(
                self.scanner,
                on_progress=self._on_progress,
                exclude_vectors=self._tested_vectors or None,
                **args,
            )
        elif tool_name in ("analyze_graph", "analyze"):
            return analyze_graph(self.scanner.graph, **args)
        elif tool_name in ("query_findings", "query"):
            return query_findings(self._attack_results, **args)
        elif tool_name in ("generate_vectors", "vectors"):
            return generate_vectors(**args)
        elif tool_name in ("explain_finding", "explain"):
            return await explain_finding(self.llm_client, **args)
        elif tool_name in ("suggest_remediation", "remediate"):
            return await suggest_remediation(self.llm_client, **args)
        elif tool_name in ("web_search", "search"):
            return await web_search(self.llm_client, **args)
        elif tool_name in ("read_target_docs", "docs"):
            return await read_target_docs(llm_client=self.llm_client, **args)
        else:
            # Default to running a scan
            logger.warning("Unknown tool '%s', defaulting to run_scan", tool_name)
            return await run_scan(self.scanner)

    def _parse_action(self, llm_response: str) -> dict[str, Any]:
        """Parse the LLM's response into a tool call action.

        Attempts to extract structured JSON from the response.
        Falls back to heuristics if JSON is not found.
        """
        # Try to find JSON block in response
        import re

        json_match = re.search(r"```json\s*(\{.*?\})\s*```", llm_response, re.DOTALL)
        if json_match:
            try:
                return dict(json.loads(json_match.group(1)))
            except json.JSONDecodeError:
                pass

        # Try inline JSON
        json_match = re.search(r"\{[^{}]*\"tool\"[^{}]*\}", llm_response, re.DOTALL)
        if json_match:
            try:
                return dict(json.loads(json_match.group(0)))
            except json.JSONDecodeError:
                pass

        # Heuristic: look for tool name mentions
        response_lower = llm_response.lower()
        tool_keywords = {
            "run_scan": ["scan", "campaign", "full scan"],
            "run_phase": ["phase", "single phase", "reconnaissance", "exploitation"],
            "analyze_graph": ["graph", "analyze", "knowledge graph", "attack paths"],
            "query_findings": ["findings", "query", "search results", "vulnerabilities found"],
            "generate_vectors": ["generate", "vectors", "dynamic", "custom attacks"],
            "explain_finding": ["explain", "explanation", "what does this mean"],
            "suggest_remediation": ["remediation", "fix", "remediate", "suggestion"],
            "web_search": ["search", "cve", "known vulnerabilities"],
        }

        for tool, keywords in tool_keywords.items():
            if any(kw in response_lower for kw in keywords):
                # Extract any phase name mentioned
                args: dict[str, Any] = {}
                phases = [
                    "reconnaissance",
                    "trust_building",
                    "capability_mapping",
                    "vulnerability_discovery",
                    "exploitation_setup",
                    "execution",
                ]
                for phase in phases:
                    if phase in response_lower and tool == "run_phase":
                        args["phase"] = phase
                        break
                return {"tool": tool, "args": args}

        # Default to scan
        return {"tool": "run_scan", "args": {}}

    def _summarize_result(self, tool_name: str, result: dict[str, Any]) -> str:
        """Create a brief summary of a tool result."""
        if tool_name in ("run_scan", "scan"):
            vulns = result.get("total_vulnerabilities", 0)
            attacks = result.get("total_attacks_executed", 0)
            return f"Scan complete: {vulns} vulns found in {attacks} attacks"
        elif tool_name in ("run_phase", "phase"):
            vulns = len(result.get("vulnerabilities_found", []))
            return f"Phase complete: {vulns} vulns found"
        elif tool_name in ("analyze_graph", "analyze"):
            paths = result.get("total_attack_paths", 0)
            chains = len(result.get("dangerous_chains", []))
            return f"Graph: {paths} attack paths, {chains} dangerous chains"
        elif tool_name in ("query_findings", "query"):
            return f"Found {result.get('total_matching', 0)} matching findings"
        elif tool_name in ("generate_vectors", "vectors"):
            return f"Generated {result.get('generated_count', 0)} vectors"
        elif tool_name in ("explain_finding", "explain"):
            return "Finding explained"
        elif tool_name in ("suggest_remediation", "remediate"):
            return "Remediation suggested"
        elif tool_name in ("web_search", "search"):
            return f"Search results for: {result.get('query', '')}"
        return f"Tool {tool_name} completed"

    def _emit_agent_event(self, event_type: str, message: str) -> None:
        """Emit an agent-level progress event."""
        if self._on_agent_event is not None:
            self._on_agent_event(event_type, message)

    async def run(
        self,
        goal: str,
        *,
        session: PentestSession | None = None,
    ) -> PentestSession:
        """Run the pentesting agent autonomously.

        Args:
            goal: The pentesting objective.
            session: Optional existing session to resume.

        Returns:
            Completed PentestSession with all findings and steps.
        """
        if session is None:
            session = PentestSession(max_iterations=self.max_iterations)
        session.plan = PentestPlan(goal=goal)
        session.status = PentestStatus.EXECUTING

        initial_state: AgentState = {
            "messages": [],
            "session": session,
            "scanner": self.scanner,
            "llm_client": self.llm_client,
            "attack_results": [],
            "next_action": "",
            "should_stop": False,
            "error": None,
        }

        try:
            final_state = await self._graph.ainvoke(initial_state)
            result: PentestSession = final_state["session"]
            return result
        except Exception as e:
            logger.error("Pentesting agent failed: %s", e, exc_info=True)
            session.status = PentestStatus.FAILED
            session.add_step(
                PentestStep(
                    step_type=PentestStepType.PLAN,
                    reasoning=f"Agent failed with error: {e}",
                    result_summary=f"Error: {e}",
                )
            )
            return session

    async def send_directive(
        self,
        session: PentestSession,
        message: str,
    ) -> PentestSession:
        """Send a human directive to the agent in interactive mode.

        Args:
            session: Current session.
            message: Human directive (e.g. "Focus on prompt injection").

        Returns:
            Updated session after processing the directive.
        """
        session.add_message("human", message)
        session.status = PentestStatus.EXECUTING

        # Inject the directive as context for the next plan
        initial_state: AgentState = {
            "messages": [{"role": "user", "content": f"User directive: {message}"}],
            "session": session,
            "scanner": self.scanner,
            "llm_client": self.llm_client,
            "attack_results": [],
            "next_action": "",
            "should_stop": False,
            "error": None,
        }

        try:
            final_state = await self._graph.ainvoke(initial_state)
            result_session: PentestSession = final_state["session"]
            result_session.add_message(
                "agent",
                result_session.steps[-1].result_summary if result_session.steps else "Done",
            )
            return result_session
        except Exception as e:
            logger.error("Directive processing failed: %s", e, exc_info=True)
            session.add_message("agent", f"Error processing directive: {e}")
            return session
