"""
PentestAgent — autonomous pentesting agent.

Main agent loop: Think → Act → Observe → Think → Act → ...
Uses LLM to decide what to test next based on all context gathered so far.
"""

from __future__ import annotations

import asyncio
import json
import signal
import sys
import time
from typing import Any

from rich.panel import Panel
from rich.table import Table

from senshi.agent.actions import Action, ActionExecutor, ActionResult
from senshi.agent.context import PentestContext
from senshi.agent.planner import AgentPlanner
from senshi.ai.brain import Brain
from senshi.ai.prompts.agent import EXPLORE_ENDPOINT_PROMPT
from senshi.core.session import Session
from senshi.dast.crawler import Crawler, DiscoveredEndpoint
from senshi.dast.tech_detector import TechDetector
from senshi.reporters.models import (
    Confidence, Finding, ScanMode, ScanResult, ScanState, Severity,
)
from senshi.utils.logger import (
    console, get_logger, print_error, print_finding, print_status, print_success,
)

logger = get_logger("senshi.agent.pentest_agent")


class PentestAgent:
    """
    Autonomous pentesting agent.

    Uses LLM to decide what to test next based on accumulated context.
    Supports browser-based exploitation, WebSocket testing, multi-account
    IDOR, and automatic PoC generation.
    """

    def __init__(
        self,
        target: str,
        brain: Brain,
        session: Session,
        *,
        max_iterations: int = 50,
        strict_mode: bool = False,
        budget: int = 0,
        browser_enabled: bool = False,
        ws_enabled: bool = False,
        stealth: bool = False,
        auth2: str = "",
        target_profile: str = "",
        output: str = "",
    ) -> None:
        self.target = target
        self.brain = brain
        self.session = session
        self.max_iterations = max_iterations
        self.strict_mode = strict_mode
        self.browser_enabled = browser_enabled
        self.ws_enabled = ws_enabled
        self.stealth = stealth
        self.auth2 = auth2
        self.target_profile_name = target_profile
        self.output = output

        self.context = PentestContext(target)
        self.planner = AgentPlanner(brain, budget=budget)
        self.executor = ActionExecutor()
        self.scan_state: ScanState | None = None

        self._browser_exploiter = None
        self._ws_tester = None
        self._auth_manager = None

    async def run(self) -> ScanResult:
        """
        Main agent loop.

        Phase 1: Recon (automated)
        Phase 2: Autonomous testing (LLM-driven)
        Phase 3: Exploitation confirmation (browser)
        Phase 4: PoC + Report generation
        """
        start_time = time.time()

        result = ScanResult(
            target=self.target,
            mode=ScanMode.DAST,
            provider=self.brain.provider,
            model=self.brain.model,
        )

        # Progressive save
        if self.output:
            output_path = self.output
        else:
            from senshi.core.engine import _generate_output_path
            output_path = _generate_output_path(self.target)
        self.scan_state = ScanState(output_path)
        console.print(f"[dim]Results saving to: {output_path}[/dim]")

        # Ctrl+C handler
        original_handler = signal.getsignal(signal.SIGINT)
        signal.signal(signal.SIGINT, self._handle_interrupt)

        try:
            # ── Phase 1: Recon ──────────────────────────────────
            console.print("\n[bold cyan]Phase 1:[/bold cyan] Reconnaissance...")
            await self._recon()
            result.endpoints_discovered = len(self.context.endpoints)

            if not self.context.endpoints:
                print_error("No endpoints discovered, cannot pentest")
                result.completed_at = ""
                return result

            # ── Phase 2: Register action handlers ───────────────
            self._register_handlers()

            # ── Phase 3: Agent loop ─────────────────────────────
            console.print(f"\n[bold cyan]Phase 2:[/bold cyan] Autonomous testing "
                          f"(max {self.max_iterations} iterations)...\n")

            for i in range(self.max_iterations):
                self.context.iteration = i + 1

                # Ask LLM what to do next
                console.print(f"  [bold]Iteration {i + 1}[/bold]")
                action = await self.planner.next_action(self.context)

                if action.type == "done":
                    console.print(f"    [green]✓ Agent finished: {action.reasoning}[/green]")
                    break

                console.print(f"    [yellow]▶[/yellow] {action.type}: {action.reasoning}")

                # Check finding cap
                if action.type in ("scan_endpoint", "fuzz_parameter", "test_ssrf", "test_idor", "test_auth", "test_graphql"):
                    ep = action.params.get("endpoint", action.params.get("ws_url", ""))
                    cat = action.params.get("vuln_type", action.params.get("payload_type", ""))
                    if not cat:
                        if action.type == "test_ssrf": cat = "ssrf"
                        elif action.type == "test_idor": cat = "idor"
                        elif action.type == "test_auth": cat = "auth"
                        elif action.type == "test_graphql": cat = "graphql"
                    
                    if ep and cat:
                        from senshi.agent.actions import VULN_TYPE_ALIASES
                        cat_alias = VULN_TYPE_ALIASES.get(cat.lower(), cat.lower())
                        count = sum(1 for f in self.context.findings if f.endpoint == ep and f.category == cat_alias)
                        if count >= 2:
                            msg = f"Skipping {action.type} on {ep} for {cat_alias} — already found max {count} issues (cap reached)"
                            console.print(f"    [yellow]⚠ {msg}[/yellow]")
                            self.context.add_observation(action.type, ep, f"Skipped: already found {count} issues")
                            continue

                # Execute the action
                action_result = await self.executor.execute(action, self.context)

                # Record observation
                self.context.add_observation(
                    action_type=action.type,
                    target=action.params.get("endpoint", action.params.get("url", "")),
                    result_summary=action_result.summary,
                    is_interesting=action_result.is_interesting,
                    raw_data=action_result.raw_data,
                )

                # Record findings
                for finding in action_result.findings:
                    self.context.add_finding(finding)
                    if self.scan_state:
                        self.scan_state.add_finding(finding)
                    print_finding(
                        finding.severity.value,
                        finding.title,
                        finding.endpoint or finding.file_path,
                    )

                # Record failures
                if not action_result.findings and not action_result.is_interesting:
                    ep = action.params.get("endpoint", action.params.get("url", ""))
                    vt = action.params.get("vuln_type", action.params.get("payload_type", ""))
                    self.context.failed_tests.append({
                        "action": action.type,
                        "endpoint": ep,
                        "vuln_type": vt,
                        "reason": action_result.error or "no issues found"
                    })

                # Mark tested
                if action.type in ("scan_endpoint", "fuzz_parameter"):
                    ep = action.params.get("endpoint", "")
                    vt = action.params.get("vuln_type", action.params.get("payload_type", ""))
                    res = "confirmed" if action_result.findings else (
                        "interesting" if action_result.is_interesting else "negative"
                    )
                    self.context.mark_tested(ep, vt, result=res)

                console.print(f"    [dim]{action_result.summary}[/dim]\n")

            # ── Phase 4: PoC generation ─────────────────────────
            if self.context.findings:
                console.print(f"\n[bold cyan]Phase 3:[/bold cyan] Generating PoCs...")
                await self._generate_pocs()

            # ── Phase 5: Chain analysis ─────────────────────────
            if len(self.context.findings) >= 2:
                console.print(f"\n[bold cyan]Phase 4:[/bold cyan] Analyzing chains...")
                from senshi.ai.chain_builder import ChainBuilder
                chain_builder = ChainBuilder(self.brain)
                chains = chain_builder.build_chains(
                    self.context.findings, target_description=self.target
                )
                result.chains = chains
                if chains:
                    print_success(f"Found {len(chains)} exploit chains")

            result.findings = self.context.findings
            result.completed_at = time.strftime("%Y-%m-%dT%H:%M:%S")

            if self.scan_state:
                self.scan_state.findings = self.context.findings
                self.scan_state.llm_calls = self.context.llm_calls
                self.scan_state.complete()

        finally:
            signal.signal(signal.SIGINT, original_handler)

        elapsed = time.time() - start_time
        self._print_dashboard(result, elapsed, output_path)
        return result

    # ── Recon ─────────────────────────────────────────────────

    async def _recon(self) -> None:
        """Automated reconnaissance — tech detect, crawl, browser recon."""
        # Tech detection
        print_status("Detecting technology stack...")
        try:
            tech_detector = TechDetector(self.session)
            tech_info = tech_detector.detect()
            self.context.tech_stack = tech_info
            print_success(f"Tech: {tech_detector.get_summary(tech_info)}")
        except Exception as e:
            logger.warning(f"Tech detection failed: {e}")

        # Standard crawl
        print_status("Crawling endpoints...")
        crawler = Crawler(self.session, brain=self.brain)
        endpoints = crawler.crawl()
        self.context.add_endpoints([ep.to_dict() for ep in endpoints])
        print_success(f"Crawled {len(endpoints)} endpoints")

        # Browser recon
        if self.browser_enabled:
            try:
                from senshi.dast.browser_recon import BrowserRecon
                print_status("Browser recon (headless Chromium)...")
                recon = BrowserRecon(timeout=30)
                browser_eps = recon.discover(self.target, auth=self.session._auth_raw if hasattr(self.session, '_auth_raw') else "")
                new_eps = [ep.to_dict() for ep in browser_eps]
                self.context.add_endpoints(new_eps)
                print_success(f"Browser discovered {len(browser_eps)} additional endpoints")
            except ImportError:
                print_error("Playwright not installed — skipping browser recon")
            except Exception as e:
                logger.warning(f"Browser recon failed: {e}")

        # Target profile endpoints
        if self.target_profile_name:
            self._load_target_profile()

        # Show endpoints
        console.print(f"\n  [bold]{len(self.context.endpoints)} endpoints discovered:[/bold]")
        for ep in self.context.endpoints[:20]:
            params = ", ".join(ep.get("params", []))
            console.print(f"    {ep.get('method', 'GET'):4s} {ep['url']}"
                          f"{' (' + params + ')' if params else ''}")
        if len(self.context.endpoints) > 20:
            console.print(f"    [dim]... and {len(self.context.endpoints) - 20} more[/dim]")

    def _load_target_profile(self) -> None:
        """Load pre-built target profile endpoints."""
        try:
            from senshi.targets import get_profile
            profile = get_profile(self.target_profile_name)
            if profile and profile.get("endpoints"):
                self.context.add_endpoints(profile["endpoints"])
                if profile.get("scope"):
                    self.context.scope_rules = profile["scope"]
                print_success(f"Loaded target profile: {self.target_profile_name}")
        except Exception as e:
            logger.debug(f"Target profile not found: {e}")

    # ── Handler Registration ──────────────────────────────────

    def _register_handlers(self) -> None:
        """Register action handlers based on available capabilities."""
        self.executor.register("scan_endpoint", self._handle_scan)
        self.executor.register("fuzz_parameter", self._handle_fuzz)
        self.executor.register("test_auth", self._handle_test_auth)
        self.executor.register("test_idor", self._handle_test_idor)
        self.executor.register("test_ssrf", self._handle_test_ssrf)
        self.executor.register("explore_endpoint", self._handle_explore)
        self.executor.register("escalate", self._handle_escalate)

        # Browser tests (optional)
        if self.browser_enabled:
            self.executor.register("browser_test", self._handle_browser_test)

        # WebSocket tests (optional)
        if self.ws_enabled:
            self.executor.register("websocket_test", self._handle_ws_test)

    # ── Action Handlers ───────────────────────────────────────

    async def _handle_scan(self, params: dict, ctx: PentestContext) -> ActionResult:
        """Run a specific scanner against an endpoint."""
        from senshi.core.engine import DAST_SCANNERS
        from senshi.dast.crawler import DiscoveredEndpoint

        from urllib.parse import urlparse, parse_qs
        
        endpoint_url = params.get("endpoint", "")
        vuln_type = params.get("vuln_type", "xss")

        from senshi.agent.actions import VULN_TYPE_ALIASES
        vuln_type = VULN_TYPE_ALIASES.get(vuln_type.lower(), vuln_type.lower())

        scanner_class = DAST_SCANNERS.get(vuln_type)
        if not scanner_class:
            return ActionResult(success=False, error=f"Unknown vuln type: {vuln_type}")

        # Build endpoint object
        url_params = list(parse_qs(urlparse(endpoint_url).query).keys())
        base_url = endpoint_url.split("?")[0]
        
        ep_data = next((e for e in ctx.endpoints if e["url"].split("?")[0] == base_url), None)
        
        if ep_data:
            merged_params = list(set(ep_data.get("params", []) + url_params))
        else:
            merged_params = url_params

        ep = DiscoveredEndpoint(
            url=endpoint_url,
            method=ep_data.get("method", "GET") if ep_data else "GET",
            params=merged_params,
        )

        def on_finding(f: Finding) -> None:
            if getattr(self, "scan_state", None):
                self.scan_state.add_finding(f)

        scanner = scanner_class(
            session=self.session,
            brain=self.brain,
            endpoints=[ep],
            tech_summary=ctx.tech_summary,
            max_payloads=15,
            rate_limit=self.session.rate_limiter.min_interval if hasattr(self.session, 'rate_limiter') else 1.0,
            on_finding=on_finding,
        )

        try:
            findings = scanner.scan()
            return ActionResult(
                success=True,
                findings=findings,
                is_interesting=len(findings) > 0,
                observations=[f"Scanned {endpoint_url} for {vuln_type}: {'found ' + str(len(findings)) + ' issues' if findings else 'clean'}"],
            )
        except Exception as e:
            return ActionResult(success=False, error=str(e))

    async def _handle_fuzz(self, params: dict, ctx: PentestContext) -> ActionResult:
        """Fuzz a specific parameter."""
        endpoint = params.get("endpoint", "")
        param = params.get("param", "")
        payload_type = params.get("payload_type", "injection")

        if not endpoint or not param:
            return ActionResult(success=False, error="Missing endpoint or param")

        from senshi.dast.fuzzer import DeterministicFuzzer
        fuzzer = DeterministicFuzzer(self.session)
        
        vuln_types = None
        # Map generic LLM payload types to fuzzer payload sets
        if payload_type.lower() in ["sqli", "injection", "sql"]:
            vuln_types = ["sqli"]
        elif payload_type.lower() in ["xss", "cross_site_scripting"]:
            vuln_types = ["xss"]
        elif payload_type.lower() in ["ssrf"]:
            vuln_types = ["ssrf"]
        elif payload_type.lower() in ["cmdi", "command_injection"]:
            vuln_types = ["cmdi"]
            
        findings = fuzzer.fuzz_endpoint(endpoint, method="GET", params=[param], vuln_types=vuln_types)
        
        if not findings:
            findings = fuzzer.fuzz_endpoint(endpoint, method="POST", params=[param], vuln_types=vuln_types)

        return ActionResult(
            success=True,
            findings=findings,
            is_interesting=len(findings) > 0,
            observations=[f"Fuzzed {param} on {endpoint}: {'found issues' if findings else 'no anomalies'}"],
        )

    async def _handle_test_auth(self, params: dict, ctx: PentestContext) -> ActionResult:
        """Test authentication on an endpoint using AuthScanner."""
        params["vuln_type"] = "auth"
        return await self._handle_scan(params, ctx)

    async def _handle_test_idor(self, params: dict, ctx: PentestContext) -> ActionResult:
        """Test for IDOR using IdorScanner."""
        params["vuln_type"] = "idor"
        return await self._handle_scan(params, ctx)

    async def _handle_test_ssrf(self, params: dict, ctx: PentestContext) -> ActionResult:
        """Test for SSRF using SsrfScanner."""
        params["vuln_type"] = "ssrf"
        return await self._handle_scan(params, ctx)

    async def _handle_explore(self, params: dict, ctx: PentestContext) -> ActionResult:
        """Explore an unknown endpoint."""
        url = params.get("url", "")
        try:
            r = self.session.get(url)
            # Ask LLM to analyze the response
            prompt = EXPLORE_ENDPOINT_PROMPT.format(
                endpoint=url,
                method="GET",
                status=r.status_code,
                headers=json.dumps(dict(r.headers), indent=2)[:1000],
                body=r.body[:3000],
            )
            ctx.llm_calls += 1
            analysis = await self.brain.async_think(
                system_prompt="You are a security researcher analyzing an endpoint.",
                user_prompt=prompt,
                json_schema={"type": "object"},
            )

            observations = []
            if isinstance(analysis, dict):
                if analysis.get("discovered_params"):
                    # Update endpoint in context
                    for ep in ctx.endpoints:
                        if ep["url"] == url:
                            ep["params"] = list(set(ep.get("params", []) + analysis["discovered_params"]))
                    observations.append(f"Discovered params: {analysis['discovered_params']}")
                if analysis.get("vuln_types_to_test"):
                    observations.append(f"Suggested tests: {analysis['vuln_types_to_test']}")
                if analysis.get("notes"):
                    observations.append(f"Notes: {analysis['notes']}")

            return ActionResult(
                success=True,
                is_interesting=bool(observations),
                observations=observations or [f"Explored {url}: status {r.status_code}"],
                raw_data=analysis if isinstance(analysis, dict) else {},
            )
        except Exception as e:
            return ActionResult(success=False, error=str(e))

    async def _handle_escalate(self, params: dict, ctx: PentestContext) -> ActionResult:
        """Attempt to escalate a finding."""
        idx = params.get("finding_index", 0)
        if idx >= len(ctx.findings):
            return ActionResult(success=False, error="Invalid finding index")

        finding = ctx.findings[idx]
        return ActionResult(
            success=True,
            observations=[f"Escalation analysis for: {finding.title}"],
        )

    async def _handle_browser_test(self, params: dict, ctx: PentestContext) -> ActionResult:
        """Run a Playwright browser test."""
        try:
            from senshi.browser.exploiter import BrowserExploiter
        except ImportError:
            return ActionResult(success=False, error="Playwright not installed")

        url = params.get("url", "")
        test_type = params.get("test_type", "xss_confirm")

        try:
            exploiter = BrowserExploiter()
            await exploiter.start()

            if test_type == "xss_confirm":
                payload = params.get("payload", "<script>alert(1)</script>")
                result = await exploiter.confirm_xss(url, payload)
                findings = []
                if result.get("confirmed"):
                    findings.append(Finding(
                        title=f"Confirmed XSS in {url}",
                        severity=Severity.HIGH,
                        confidence=Confidence.CONFIRMED,
                        category="xss",
                        mode=ScanMode.DAST,
                        endpoint=url,
                        payload=payload,
                        confirmed=True,
                    ))
                return ActionResult(
                    success=True,
                    findings=findings,
                    is_interesting=bool(findings),
                    observations=[f"Browser XSS test: {'CONFIRMED' if findings else 'not executable'}"],
                )
            return ActionResult(success=True, observations=[f"Browser test {test_type} completed"])
        except Exception as e:
            return ActionResult(success=False, error=str(e))

    async def _handle_ws_test(self, params: dict, ctx: PentestContext) -> ActionResult:
        """Test a WebSocket endpoint."""
        try:
            from senshi.websocket.tester import WebSocketTester
        except ImportError:
            return ActionResult(success=False, error="websockets not installed")

        ws_url = params.get("ws_url", "")
        try:
            tester = WebSocketTester(self.brain)
            findings = await tester.test_endpoint(ws_url)
            return ActionResult(
                success=True,
                findings=findings,
                is_interesting=len(findings) > 0,
                observations=[f"WebSocket test: {len(findings)} issues found"],
            )
        except Exception as e:
            return ActionResult(success=False, error=str(e))

    # ── PoC Generation ────────────────────────────────────────

    async def _generate_pocs(self) -> None:
        """Generate PoCs for all findings."""
        try:
            from senshi.poc.generator import PoCGenerator
            gen = PoCGenerator(self.brain)
            for finding in self.context.findings:
                try:
                    poc = await gen.generate(finding)
                    if poc:
                        finding.poc_curl = poc.get("curl", "")
                        finding.poc_python = poc.get("python_script", "")
                        finding.poc_steps = poc.get("browser_steps", [])
                except Exception as e:
                    logger.debug(f"PoC generation failed for {finding.title}: {e}")
            print_success(f"Generated PoCs for {len(self.context.findings)} findings")
        except ImportError:
            logger.debug("PoC generator not available")

    # ── Interrupt ─────────────────────────────────────────────

    def _handle_interrupt(self, signum: Any, frame: Any) -> None:
        """Save partial results on Ctrl+C."""
        console.print("\n[yellow]⚡ Pentest interrupted! Saving partial results...[/yellow]")
        if self.scan_state:
            self.scan_state.findings = self.context.findings
            self.scan_state.llm_calls = self.context.llm_calls
            self.scan_state.interrupt()
            console.print(f"[green]✓ {len(self.context.findings)} findings saved to "
                          f"{self.scan_state.output_path}[/green]")
        sys.exit(0)

    # ── Dashboard ─────────────────────────────────────────────

    def _print_dashboard(self, result: ScanResult, elapsed: float, output_path: str) -> None:
        """Print scan summary dashboard."""
        console.print("\n")

        mins = int(elapsed // 60)
        secs = int(elapsed % 60)
        duration = f"{mins}m {secs}s" if mins else f"{secs}s"

        header = Panel(
            f"[bold]Target:[/bold] {result.target}\n"
            f"[bold]Mode:[/bold] Autonomous Pentest\n"
            f"[bold]Duration:[/bold] {duration}\n"
            f"[bold]Iterations:[/bold] {self.context.iteration}\n"
            f"[bold]Endpoints:[/bold] {result.endpoints_discovered}\n"
            f"[bold]LLM Calls:[/bold] {self.context.llm_calls}\n"
            f"[bold]Findings:[/bold] {len(result.findings)}\n"
            f"[bold]Chains:[/bold] {len(result.chains)}",
            title="[bold]Senshi Autonomous Pentest Report[/bold]",
            border_style="cyan",
        )
        console.print(header)

        if result.findings:
            table = Table(show_header=True, header_style="bold")
            table.add_column("Sev", width=10)
            table.add_column("Conf", width=10)
            table.add_column("Finding")
            table.add_column("PoC", width=4)

            for f in sorted(result.findings, key=lambda x: x.severity.rank, reverse=True):
                sev = f.severity.value
                color = {"critical": "red bold", "high": "red", "medium": "yellow"}.get(sev, "blue")
                has_poc = "✓" if getattr(f, "poc_curl", "") else ""
                table.add_row(f"[{color}]{sev.upper()}[/{color}]", f.confidence.value, f.title, has_poc)

            console.print(table)

        console.print(f"\n[bold green]Results saved to:[/bold green] {output_path}")
        if self.strict_mode:
            confirmed = [f for f in result.findings if getattr(f, "confirmed", False)]
            console.print(f"[dim]Strict mode: {len(confirmed)}/{len(result.findings)} confirmed[/dim]")
        console.print()
