"""Pentesting orchestrator — public Python API for the pentesting agent.

Provides ``PentestOrchestrator`` which wires up the LangGraph agent,
scanner, LLM client, and deduplicator. This is the main entry point
for both the CLI and programmatic usage.

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

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import Callable

    from ziran.application.agent_scanner.scanner import ProgressEvent
    from ziran.domain.interfaces.adapter import BaseAgentAdapter
    from ziran.infrastructure.llm.base import BaseLLMClient

from ziran.domain.entities.pentest import PentestSession

logger = logging.getLogger(__name__)


class PentestOrchestrator:
    """High-level orchestrator for autonomous pentesting sessions.

    Wires up the pentesting agent, scanner, LLM client, and
    deduplicator into a single interface. Supports both autonomous
    and interactive (directive-based) modes.

    Example::

        orchestrator = PentestOrchestrator(
            adapter=my_adapter,
            llm_client=my_llm_client,
        )
        session = await orchestrator.run_autonomous(
            goal="Find prompt injection vulnerabilities",
        )
        print(f"Found {session.total_vulnerabilities} vulnerabilities")
    """

    def __init__(
        self,
        adapter: BaseAgentAdapter,
        llm_client: BaseLLMClient,
        *,
        config: dict[str, Any] | None = None,
        max_iterations: int = 10,
        embedding_model: str = "text-embedding-3-small",
        attack_timeout: float = 120.0,
        on_progress: Callable[[ProgressEvent], None] | None = None,
        on_agent_event: Callable[[str, str], None] | None = None,
        show_progress: bool = True,
    ) -> None:
        """Initialize the orchestrator.

        Args:
            adapter: Agent adapter for the target agent.
            llm_client: LLM client for agent reasoning and embeddings.
            config: Optional scanner configuration overrides.
            max_iterations: Maximum agent iterations per session.
            embedding_model: Model for finding deduplication embeddings.
            attack_timeout: Per-attack timeout in seconds (default 120).
            on_progress: Callback for scan-level progress events.
            on_agent_event: Callback for agent-level events (plan/execute/reason).
            show_progress: Show built-in Rich progress display (default True).
                Ignored if custom on_progress or on_agent_event are provided.
        """
        from ziran.application.agent_scanner.scanner import AgentScanner
        from ziran.application.pentesting.agent import PentestAgent
        from ziran.application.pentesting.deduplication import FindingDeduplicator

        # Set up built-in progress display unless custom callbacks are provided
        self._progress_display = None
        if show_progress and on_progress is None and on_agent_event is None:
            try:
                from ziran.application.pentesting.progress import PentestProgressDisplay

                self._progress_display = PentestProgressDisplay()
                on_progress = self._progress_display.on_progress
                on_agent_event = self._progress_display.on_agent_event
            except Exception:
                pass  # Rich not available or import error — run silently

        scanner_config = dict(config or {})
        scanner_config.setdefault("llm_client", llm_client)
        scanner_config.setdefault("attack_timeout", attack_timeout)

        self.adapter = adapter
        self.llm_client = llm_client
        self.scanner = AgentScanner(adapter=adapter, config=scanner_config)
        self.agent = PentestAgent(
            scanner=self.scanner,
            llm_client=llm_client,
            max_iterations=max_iterations,
            on_progress=on_progress,
            on_agent_event=on_agent_event,
        )
        self.deduplicator = FindingDeduplicator(
            llm_client=llm_client,
            embedding_model=embedding_model,
        )
        self._sessions: dict[str, PentestSession] = {}

    async def run_autonomous(
        self,
        goal: str,
        *,
        max_iterations: int | None = None,
    ) -> PentestSession:
        """Run a fully autonomous pentesting session.

        The agent plans, executes, and adapts without human intervention.

        Args:
            goal: The pentesting objective.
            max_iterations: Override max iterations for this session.

        Returns:
            Completed PentestSession with findings and steps.
        """
        session = PentestSession(
            max_iterations=max_iterations or self.agent.max_iterations,
        )

        logger.info(
            "Starting autonomous pentesting session %s: %s",
            session.session_id,
            goal,
        )

        session = await self.agent.run(goal=goal, session=session)

        # Deduplicate findings
        try:
            attack_results = self.scanner._attack_results
            if attack_results:
                findings = await self.deduplicator.deduplicate(attack_results)
                session.findings = findings
                logger.info(
                    "Deduplicated %d attack results into %d findings",
                    len(attack_results),
                    len(findings),
                )
        except Exception as e:
            logger.warning("Finding deduplication failed: %s", e)

        self._sessions[session.session_id] = session

        # Show built-in progress summary
        if self._progress_display is not None:
            self._progress_display.show_summary(
                session_id=session.session_id,
                status=session.status.value,
                iterations=session.iteration_count,
                vulnerabilities=session.total_vulnerabilities,
                attacks=session.total_attacks_executed,
                findings=len(session.findings),
                tokens=session.token_usage.get("total_tokens", 0),
            )

        return session

    async def start_interactive(
        self,
        goal: str,
        *,
        max_iterations: int | None = None,
    ) -> PentestSession:
        """Start an interactive pentesting session.

        Creates a session that can receive human directives via
        ``send_directive()``.

        Args:
            goal: The initial pentesting objective.
            max_iterations: Override max iterations.

        Returns:
            Session ready for interactive use.
        """
        session = PentestSession(
            max_iterations=max_iterations or self.agent.max_iterations,
        )
        session.add_message("system", f"Pentesting session started with goal: {goal}")

        # Run initial planning
        session = await self.agent.run(goal=goal, session=session)
        self._sessions[session.session_id] = session
        return session

    async def send_directive(
        self,
        session_id: str,
        message: str,
    ) -> PentestSession:
        """Send a human directive to an interactive session.

        Args:
            session_id: ID of the session.
            message: Human directive.

        Returns:
            Updated session.
        """
        session = self._sessions.get(session_id)
        if session is None:
            raise ValueError(f"Session {session_id} not found")

        session = await self.agent.send_directive(session, message)
        self._sessions[session_id] = session
        return session

    def get_session(self, session_id: str) -> PentestSession | None:
        """Retrieve a session by ID."""
        return self._sessions.get(session_id)

    async def export_report(
        self,
        session: PentestSession,
        *,
        output_dir: str | Path = "ziran_pentest_results",
        formats: list[str] | None = None,
    ) -> list[Path]:
        """Export session results to files.

        Args:
            session: The session to export.
            output_dir: Directory for output files.
            formats: List of formats to export ('json', 'md'). Default: both.

        Returns:
            List of generated file paths.
        """
        out = Path(output_dir)
        out.mkdir(parents=True, exist_ok=True)
        report_formats = formats or ["json", "md", "html"]
        generated: list[Path] = []

        session_data = session.model_dump(mode="json")

        if "json" in report_formats:
            json_path = out / f"pentest_{session.session_id}.json"
            json_path.write_text(json.dumps(session_data, indent=2, default=str))
            generated.append(json_path)

        if "md" in report_formats:
            md_path = out / f"pentest_{session.session_id}.md"
            md_content = self._generate_markdown_report(session)
            md_path.write_text(md_content)
            generated.append(md_path)

        if "html" in report_formats:
            html_path = out / f"pentest_{session.session_id}.html"
            html_content = self._generate_html_report(session)
            html_path.write_text(html_content)
            generated.append(html_path)

        return generated

    @staticmethod
    def _generate_markdown_report(session: PentestSession) -> str:
        """Generate a Markdown report from a session."""
        lines: list[str] = []

        # ── Header ──────────────────────────────────────────────
        lines.append(f"# Pentesting Report — {session.session_id}")
        lines.append("")
        lines.append(
            f"> **Status:** {session.status.value} · "
            f"**Goal:** {session.plan.goal if session.plan else 'N/A'}"
        )
        lines.append("")

        # ── Summary table ───────────────────────────────────────
        tokens = session.token_usage
        lines.extend(
            [
                "## Summary",
                "",
                "| Metric | Value |",
                "|--------|-------|",
                f"| Iterations | {session.iteration_count} |",
                f"| Vulnerabilities | {session.total_vulnerabilities} |",
                f"| Unique Findings | {len(session.findings)} |",
                f"| Attacks Executed | {session.total_attacks_executed} |",
                f"| Tokens Used | {tokens.get('total_tokens', 0):,} |",
                "",
            ]
        )

        # ── Severity breakdown ──────────────────────────────────
        if session.findings:
            from collections import Counter

            sev_counts: Counter[str] = Counter()
            for f in session.findings:
                sev_counts[f.severity.upper()] += 1

            sev_order = ["CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO"]
            sev_emoji = {
                "CRITICAL": "🔴",
                "HIGH": "🟠",
                "MEDIUM": "🟡",
                "LOW": "🟢",
                "INFO": "🔵",
            }
            lines.append("### Severity Breakdown")
            lines.append("")
            lines.append("| Severity | Count |")
            lines.append("|----------|-------|")
            for sev in sev_order:
                if sev in sev_counts:
                    lines.append(f"| {sev_emoji.get(sev, '')} {sev} | {sev_counts[sev]} |")
            lines.append("")

        # ── Findings table ──────────────────────────────────────
        if session.findings:
            lines.extend(
                [
                    "## Findings",
                    "",
                    "| # | Severity | Finding | OWASP | Evidence | Confidence |",
                    "|---|----------|---------|-------|----------|------------|",
                ]
            )
            for i, finding in enumerate(session.findings, 1):
                sev = finding.severity.upper()
                owasp = ", ".join(finding.owasp_categories) if finding.owasp_categories else "—"
                evidence_n = len(finding.attack_result_ids)
                conf = f"{finding.confidence:.0%}"
                lines.append(
                    f"| {i} | {sev} | {finding.canonical_title} | {owasp} | {evidence_n} | {conf} |"
                )
            lines.append("")

            # ── Finding details (collapsible) ───────────────────
            lines.append("### Finding Details")
            lines.append("")
            for i, finding in enumerate(session.findings, 1):
                lines.append(
                    f"<details><summary><strong>{i}. [{finding.severity.upper()}] "
                    f"{finding.canonical_title}</strong></summary>"
                )
                lines.append("")
                lines.append(f"**Description:** {finding.canonical_description}")
                lines.append("")
                if finding.owasp_categories:
                    lines.append(f"**OWASP:** {', '.join(finding.owasp_categories)}")
                if finding.attack_result_ids:
                    lines.append(f"**Related Vectors:** `{'`, `'.join(finding.attack_result_ids)}`")
                if finding.remediation_suggestion:
                    lines.append("")
                    lines.append(f"**Remediation:** {finding.remediation_suggestion}")
                lines.append("")
                lines.append("</details>")
                lines.append("")

        # ── Steps timeline ──────────────────────────────────────
        if session.steps:
            lines.extend(
                [
                    "## Agent Steps",
                    "",
                    "| # | Type | Tool | Result | Duration |",
                    "|---|------|------|--------|----------|",
                ]
            )
            for i, step in enumerate(session.steps, 1):
                stype = step.step_type.value
                tool = step.tool_name or "—"
                result = step.result_summary[:60]
                dur = f"{step.duration_seconds:.1f}s"
                lines.append(f"| {i} | {stype} | {tool} | {result} | {dur} |")
            lines.append("")

        return "\n".join(lines)

    @staticmethod
    def _generate_html_report(session: PentestSession) -> str:
        """Generate a self-contained HTML report from a pentest session."""
        import html as html_mod
        from collections import Counter

        sev_order = ["CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO"]
        sev_colors = {
            "CRITICAL": "#ef4444",
            "HIGH": "#f97316",
            "MEDIUM": "#eab308",
            "LOW": "#22c55e",
            "INFO": "#3b82f6",
        }

        # Severity counts
        sev_counts: Counter[str] = Counter()
        for f in session.findings:
            sev_counts[f.severity.upper()] += 1

        total_findings = len(session.findings) or 1

        # Build severity bar chart rows
        sev_bars = ""
        for sev in sev_order:
            cnt = sev_counts.get(sev, 0)
            if cnt:
                pct = int(100 * cnt / total_findings)
                color = sev_colors.get(sev, "#6b7280")
                sev_bars += (
                    f'<div class="bar-row">'
                    f'<span class="bar-label">{sev}</span>'
                    f'<div class="bar-track">'
                    f'<div class="bar-fill" style="width:{pct}%;'
                    f'background:{color}"></div>'
                    f"</div>"
                    f'<span class="bar-count">{cnt}</span>'
                    f"</div>\n"
                )

        # Build findings table rows
        findings_rows = ""
        for i, f in enumerate(session.findings, 1):
            sev = f.severity.upper()
            owasp = ", ".join(f.owasp_categories) if f.owasp_categories else "—"
            evidence_n = len(f.attack_result_ids)
            conf = f"{f.confidence:.0%}"
            desc_escaped = html_mod.escape(f.canonical_description[:200])
            vectors = ", ".join(f.attack_result_ids[:5])
            if len(f.attack_result_ids) > 5:
                vectors += f" +{len(f.attack_result_ids) - 5} more"
            remediation = (
                html_mod.escape(f.remediation_suggestion) if f.remediation_suggestion else ""
            )
            findings_rows += (
                f'<tr class="finding-row" '
                f'onclick="this.nextElementSibling.classList.toggle('
                f"'hidden')\">"
                f"<td>{i}</td>"
                f'<td><span class="badge {sev.lower()}">{sev}</span></td>'
                f"<td>{html_mod.escape(f.canonical_title)}</td>"
                f"<td>{owasp}</td><td>{evidence_n}</td>"
                f"<td>{conf}</td></tr>\n"
                f'<tr class="detail-row hidden"><td colspan="6">'
                f"<p>{desc_escaped}</p>"
                f'<p class="meta">Vectors: <code>{vectors}</code></p>'
                + (f'<p class="remediation">\U0001f4a1 {remediation}</p>' if remediation else "")
                + "</td></tr>\n"
            )

        # Build steps table rows
        steps_rows = ""
        for i, step in enumerate(session.steps, 1):
            stype = html_mod.escape(step.step_type.value)
            tool = html_mod.escape(step.tool_name) if step.tool_name else "—"
            result = html_mod.escape(step.result_summary[:80])
            dur = f"{step.duration_seconds:.1f}s"
            steps_rows += (
                f"<tr><td>{i}</td><td>{stype}</td>"
                f"<td>{tool}</td><td>{result}</td><td>{dur}</td></tr>\n"
            )

        goal = html_mod.escape(session.plan.goal if session.plan else "N/A")
        tokens = session.token_usage

        css = (
            ":root{--bg:#0f172a;--surface:#1e293b;--border:#334155;"
            "--text:#e2e8f0;--muted:#94a3b8;--accent:#3b82f6}"
            "*{margin:0;padding:0;box-sizing:border-box}"
            "body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',"
            "Roboto,sans-serif;background:var(--bg);color:var(--text);"
            "line-height:1.6;padding:2rem}"
            "h1{font-size:1.8rem;margin-bottom:.5rem}"
            "h2{font-size:1.3rem;margin:2rem 0 1rem;"
            "border-bottom:1px solid var(--border);padding-bottom:.5rem}"
            ".subtitle{color:var(--muted);margin-bottom:2rem}"
            ".grid{display:grid;"
            "grid-template-columns:repeat(auto-fit,minmax(140px,1fr));"
            "gap:1rem;margin-bottom:2rem}"
            ".card{background:var(--surface);border-radius:8px;"
            "padding:1.2rem;text-align:center;"
            "border:1px solid var(--border)}"
            ".card .value{font-size:2rem;font-weight:700;"
            "color:var(--accent)}"
            ".card .label{font-size:.8rem;color:var(--muted);"
            "text-transform:uppercase;letter-spacing:.05em}"
            ".bar-row{display:flex;align-items:center;"
            "gap:.5rem;margin-bottom:.4rem}"
            ".bar-label{width:70px;font-size:.85rem;text-align:right}"
            ".bar-track{flex:1;height:20px;background:var(--surface);"
            "border-radius:4px;overflow:hidden}"
            ".bar-fill{height:100%;border-radius:4px}"
            ".bar-count{width:30px;font-size:.85rem}"
            "table{width:100%;border-collapse:collapse}"
            "th,td{padding:.6rem .8rem;text-align:left;"
            "border-bottom:1px solid var(--border);font-size:.9rem}"
            "th{background:var(--surface);color:var(--muted);"
            "font-size:.8rem;text-transform:uppercase;"
            "letter-spacing:.04em;position:sticky;top:0}"
            ".finding-row{cursor:pointer}"
            ".finding-row:hover{background:rgba(59,130,246,.08)}"
            ".detail-row td{background:var(--surface);"
            "padding:1rem 1.5rem}"
            ".detail-row.hidden{display:none}"
            ".detail-row .meta{color:var(--muted);"
            "font-size:.85rem;margin-top:.5rem}"
            ".detail-row .remediation{color:#34d399;margin-top:.5rem}"
            "code{background:rgba(59,130,246,.15);"
            "padding:.15em .4em;border-radius:3px;font-size:.85em}"
            ".badge{padding:.2em .6em;border-radius:4px;"
            "font-size:.75rem;font-weight:600}"
            ".badge.critical{background:rgba(239,68,68,.2);"
            "color:#f87171}"
            ".badge.high{background:rgba(249,115,22,.2);color:#fb923c}"
            ".badge.medium{background:rgba(234,179,8,.2);color:#facc15}"
            ".badge.low{background:rgba(34,197,94,.2);color:#4ade80}"
            ".badge.info{background:rgba(59,130,246,.2);color:#60a5fa}"
        )

        return (
            "<!DOCTYPE html>\n"
            '<html lang="en">\n<head>\n'
            '<meta charset="utf-8">\n'
            '<meta name="viewport" '
            'content="width=device-width,initial-scale=1">\n'
            f"<title>Pentest Report — {session.session_id}</title>\n"
            f"<style>{css}</style>\n"
            "</head>\n<body>\n"
            "<h1>\U0001f6e1\ufe0f Pentest Report</h1>\n"
            f'<p class="subtitle">Session '
            f"<code>{session.session_id}</code> · {goal}</p>\n\n"
            '<div class="grid">\n'
            f'  <div class="card"><div class="value">'
            f"{session.total_vulnerabilities}</div>"
            f'<div class="label">Vulnerabilities</div></div>\n'
            f'  <div class="card"><div class="value">'
            f"{len(session.findings)}</div>"
            f'<div class="label">Unique Findings</div></div>\n'
            f'  <div class="card"><div class="value">'
            f"{session.iteration_count}</div>"
            f'<div class="label">Iterations</div></div>\n'
            f'  <div class="card"><div class="value">'
            f"{session.total_attacks_executed}</div>"
            f'<div class="label">Attacks</div></div>\n'
            f'  <div class="card"><div class="value">'
            f"{tokens.get('total_tokens', 0):,}</div>"
            f'<div class="label">Tokens</div></div>\n'
            "</div>\n\n"
            "<h2>Severity Breakdown</h2>\n"
            f"{sev_bars}\n"
            "<h2>Findings</h2>\n"
            '<p style="color:var(--muted);font-size:.85rem;'
            'margin-bottom:.5rem;">'
            "Click a row to expand details</p>\n"
            "<table>\n<thead><tr>"
            "<th>#</th><th>Severity</th><th>Finding</th>"
            "<th>OWASP</th><th>Evidence</th><th>Confidence</th>"
            "</tr></thead>\n<tbody>\n"
            f"{findings_rows}"
            "</tbody>\n</table>\n\n"
            "<h2>Agent Steps</h2>\n"
            "<table>\n<thead><tr>"
            "<th>#</th><th>Type</th><th>Tool</th>"
            "<th>Result</th><th>Duration</th>"
            "</tr></thead>\n<tbody>\n"
            f"{steps_rows}"
            "</tbody>\n</table>\n\n"
            '<p style="margin-top:3rem;color:var(--muted);'
            'font-size:.8rem;text-align:center;">\n'
            "  Generated by Ziran Pentesting Agent"
            f" · {session.status.value}\n"
            "</p>\n</body>\n</html>"
        )
