"""
BREACH.AI - Shannon-Style Exploitation Engine
===============================================

Main orchestration engine that combines all Shannon-style components:
1. Source Code Analysis (white-box)
2. Browser Automation (Playwright)
3. Parallel OWASP Agents
4. Proof-by-Exploitation Validation
5. Evidence Collection
6. PoC Generation

This is the "Shannon Mode" - only reports vulnerabilities that can be
successfully exploited with undeniable proof.
"""

import asyncio
import time
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional

import aiohttp
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich import box

from .validator import ExploitationValidator, ValidationResult
from .strict_validator import StrictExploitationValidator, StrictValidationResult
from .browser import BrowserExploiter, ExploitResult, PLAYWRIGHT_AVAILABLE
from .source_analyzer import SourceCodeAnalyzer, DataFlowPath, SourceAnalysisResult
from .agents import OWASPAgentOrchestrator, AgentFinding
from .poc_generator import PoCGenerator, ProofOfConcept
from .evidence import EvidenceCollector, EvidencePackage

console = Console()


@dataclass
class ShannonFinding:
    """
    A Shannon-validated finding.

    Only includes vulnerabilities that were SUCCESSFULLY EXPLOITED.
    """
    # Basic info
    vulnerability_type: str
    severity: str
    confidence: float
    target: str
    endpoint: str
    parameter: str
    payload: str

    # Exploitation proof
    exploited: bool = True  # Always True for Shannon findings
    proof_type: str = ""
    proof_data: Dict = field(default_factory=dict)

    # Evidence
    evidence_package: Optional[EvidencePackage] = None
    screenshot: Optional[bytes] = None

    # Reproduction
    poc: Optional[ProofOfConcept] = None
    curl_command: str = ""
    reproduction_steps: List[str] = field(default_factory=list)

    # Impact
    business_impact: int = 0
    impact_explanation: str = ""

    # Remediation
    remediation: str = ""
    cwe_id: str = ""
    owasp_category: str = ""

    # Source analysis (if white-box)
    data_flow: Optional[DataFlowPath] = None

    # Timing
    discovered_at: datetime = field(default_factory=datetime.utcnow)
    exploitation_time_ms: float = 0


@dataclass
class ShannonScanResult:
    """Complete Shannon scan result."""
    target: str
    mode: str = "shannon"  # Always shannon mode

    # Timing
    started_at: datetime = field(default_factory=datetime.utcnow)
    completed_at: Optional[datetime] = None
    duration_seconds: int = 0

    # Analysis results
    source_analysis: Optional[SourceAnalysisResult] = None
    js_files_analyzed: int = 0
    data_flows_found: int = 0

    # Scan results
    endpoints_discovered: int = 0
    vulnerabilities_tested: int = 0

    # VALIDATED FINDINGS ONLY
    findings: List[ShannonFinding] = field(default_factory=list)
    total_findings: int = 0
    critical_count: int = 0
    high_count: int = 0
    medium_count: int = 0
    low_count: int = 0

    # Exploitation stats
    exploitation_attempts: int = 0
    successful_exploits: int = 0
    false_positives_filtered: int = 0

    # Business impact
    total_business_impact: int = 0

    def to_dict(self) -> Dict:
        """Convert to dictionary for JSON serialization."""
        return {
            "target": self.target,
            "mode": self.mode,
            "started_at": self.started_at.isoformat(),
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "duration_seconds": self.duration_seconds,
            "analysis": {
                "js_files_analyzed": self.js_files_analyzed,
                "data_flows_found": self.data_flows_found,
                "endpoints_discovered": self.endpoints_discovered,
            },
            "exploitation": {
                "attempts": self.exploitation_attempts,
                "successful": self.successful_exploits,
                "false_positives_filtered": self.false_positives_filtered,
            },
            "summary": {
                "total_findings": self.total_findings,
                "critical": self.critical_count,
                "high": self.high_count,
                "medium": self.medium_count,
                "low": self.low_count,
                "total_business_impact": self.total_business_impact,
            },
            "findings": [
                {
                    "vulnerability_type": f.vulnerability_type,
                    "severity": f.severity,
                    "confidence": f.confidence,
                    "endpoint": f.endpoint,
                    "parameter": f.parameter,
                    "payload": f.payload[:100],
                    "proof_type": f.proof_type,
                    "curl_command": f.curl_command,
                    "business_impact": f.business_impact,
                    "cwe_id": f.cwe_id,
                }
                for f in self.findings
            ],
        }


class ShannonEngine:
    """
    Shannon-Style Exploitation Engine.

    Implements Shannon's "Proof-by-Exploitation" methodology:
    1. Reconnaissance (spider, JS analysis)
    2. Source Analysis (data flow tracing) - WHITE BOX
    3. Vulnerability Hypothesis Generation
    4. Parallel Agent Exploitation
    5. Browser-based Validation
    6. Evidence Collection
    7. PoC Generation
    8. Reporting (only exploited vulnerabilities)
    """

    def __init__(
        self,
        timeout_minutes: int = 30,
        use_browser: bool = True,
        use_source_analysis: bool = True,
        parallel_agents: int = 5,
        screenshot: bool = True,
        evidence_dir: str = "./evidence",
    ):
        self.timeout_minutes = timeout_minutes
        self.use_browser = use_browser and PLAYWRIGHT_AVAILABLE
        self.use_source_analysis = use_source_analysis
        self.parallel_agents = parallel_agents
        self.capture_screenshot = screenshot
        self.evidence_dir = evidence_dir

        # Components (initialized in __aenter__)
        self.session: Optional[aiohttp.ClientSession] = None
        self.browser: Optional[BrowserExploiter] = None
        self.validator: Optional[ExploitationValidator] = None
        self.strict_validator: Optional[StrictExploitationValidator] = None  # Shannon-quality validator
        self.source_analyzer: Optional[SourceCodeAnalyzer] = None
        self.agent_orchestrator: Optional[OWASPAgentOrchestrator] = None
        self.poc_generator: Optional[PoCGenerator] = None
        self.evidence_collector: Optional[EvidenceCollector] = None

        # Callbacks
        self._finding_callbacks: List[Callable] = []
        self._progress_callbacks: List[Callable] = []

    def on_finding(self, callback: Callable):
        """Register callback for when findings are discovered."""
        self._finding_callbacks.append(callback)

    def on_progress(self, callback: Callable):
        """Register callback for progress updates."""
        self._progress_callbacks.append(callback)

    async def _fire_finding(self, finding: ShannonFinding):
        """Fire finding callbacks."""
        for cb in self._finding_callbacks:
            try:
                if asyncio.iscoroutinefunction(cb):
                    await cb(finding)
                else:
                    cb(finding)
            except:
                pass

    async def _fire_progress(self, percent: int, message: str):
        """Fire progress callbacks."""
        for cb in self._progress_callbacks:
            try:
                cb(percent, message)
            except:
                pass

    async def __aenter__(self):
        """Initialize all components."""
        # HTTP session
        connector = aiohttp.TCPConnector(limit=20, ssl=False)
        timeout = aiohttp.ClientTimeout(total=30)
        self.session = aiohttp.ClientSession(
            connector=connector,
            timeout=timeout,
            headers={'User-Agent': 'BREACH.AI Shannon/1.0'}
        )

        # Browser (optional)
        if self.use_browser:
            self.browser = BrowserExploiter(
                headless=True,
                timeout_ms=30000,
                screenshot=self.capture_screenshot,
            )
            await self.browser.__aenter__()

        # Exploitation validators
        # Regular validator (for backwards compatibility)
        self.validator = ExploitationValidator(
            session=self.session,
            use_browser=self.use_browser,
            timeout=30,
            screenshot=self.capture_screenshot,
        )
        await self.validator.__aenter__()

        # STRICT validator - Shannon-quality "No Exploit, No Report"
        self.strict_validator = StrictExploitationValidator(
            session=self.session,
            timeout=30,
        )
        await self.strict_validator.__aenter__()

        # Source analyzer
        if self.use_source_analysis:
            self.source_analyzer = SourceCodeAnalyzer(
                session=self.session,
                timeout=30,
            )

        # Agent orchestrator
        self.agent_orchestrator = OWASPAgentOrchestrator(
            validator=self.validator,
            source_analyzer=self.source_analyzer,
            max_concurrent=self.parallel_agents,
        )

        # PoC generator
        self.poc_generator = PoCGenerator()

        # Evidence collector
        self.evidence_collector = EvidenceCollector(
            session=self.session,
            output_dir=self.evidence_dir,
            capture_screenshots=self.capture_screenshot,
        )

        return self

    async def __aexit__(self, *args):
        """Cleanup all components."""
        if self.strict_validator:
            await self.strict_validator.__aexit__(*args)
        if self.validator:
            await self.validator.__aexit__(*args)
        if self.browser:
            await self.browser.__aexit__(*args)
        if self.session:
            await self.session.close()

    async def scan(
        self,
        target: str,
        cookies: Dict[str, str] = None,
        js_urls: List[str] = None,
        source_code: str = None,
        endpoints: List[Dict] = None,
        progress_callback: Callable = None,
    ) -> ShannonScanResult:
        """
        Run Shannon-style exploitation scan.

        Args:
            target: Target URL
            cookies: Session cookies for authenticated testing
            js_urls: JavaScript files for source analysis
            source_code: Raw source code for white-box analysis
            endpoints: Pre-discovered endpoints
            progress_callback: Progress callback (percent, message)

        Returns:
            ShannonScanResult with only exploited vulnerabilities
        """
        if progress_callback:
            self._progress_callbacks.append(progress_callback)

        result = ShannonScanResult(
            target=target,
            started_at=datetime.utcnow(),
        )

        start_time = time.time()

        self._banner(target)

        try:
            # ================================================================
            # PHASE 1: SOURCE CODE ANALYSIS (White-Box)
            # ================================================================
            await self._fire_progress(5, "Phase 1: Source Code Analysis")

            data_flows = []
            if self.use_source_analysis and (js_urls or source_code):
                console.print(f"\n[bold cyan]{'='*70}[/bold cyan]")
                console.print(f"[bold cyan]  PHASE 1: SOURCE CODE ANALYSIS (WHITE-BOX)[/bold cyan]")
                console.print(f"[bold cyan]{'='*70}[/bold cyan]")

                analysis = await self.source_analyzer.analyze(
                    target=target,
                    js_urls=js_urls,
                    source_code=source_code,
                )

                result.source_analysis = analysis
                result.js_files_analyzed = analysis.js_files_analyzed
                result.data_flows_found = len(analysis.data_flows_found)
                data_flows = analysis.data_flows_found

                self._show_source_analysis(analysis)

            await self._fire_progress(20, f"Source analysis complete: {len(data_flows)} data flows")

            # ================================================================
            # PHASE 2: PARALLEL AGENT EXPLOITATION
            # ================================================================
            await self._fire_progress(25, "Phase 2: Parallel Agent Exploitation")

            console.print(f"\n[bold cyan]{'='*70}[/bold cyan]")
            console.print(f"[bold cyan]  PHASE 2: PARALLEL OWASP AGENT EXPLOITATION[/bold cyan]")
            console.print(f"[bold cyan]{'='*70}[/bold cyan]")

            # Use provided endpoints or create from data flows
            test_endpoints = endpoints or []
            if data_flows:
                for flow in data_flows:
                    test_endpoints.append({
                        "url": target,
                        "params": [flow.source.split(".")[-1]],
                        "suggested_payloads": flow.suggested_payloads,
                    })

            # Add default params if no endpoints
            if not test_endpoints:
                test_endpoints = [{
                    "url": target,
                    "params": ["id", "q", "search", "query", "url", "file", "name", "user", "page"],
                }]

            result.endpoints_discovered = len(test_endpoints)

            # Run parallel agents
            agent_findings = await self.agent_orchestrator.hunt(
                target=target,
                endpoints=test_endpoints,
                data_flows=data_flows,
                cookies=cookies,
                on_finding=lambda f: console.print(f"[green]✓ EXPLOITED: {f.vulnerability_type} - {f.endpoint}[/green]"),
            )

            result.exploitation_attempts = sum(
                len(agent.findings) + len([f for f in agent.findings if not f.exploited])
                for agent in self.agent_orchestrator.agents
            )
            result.successful_exploits = len(agent_findings)

            await self._fire_progress(60, f"Agents complete: {len(agent_findings)} exploited")

            # ================================================================
            # PHASE 3: BROWSER VALIDATION & EVIDENCE COLLECTION
            # ================================================================
            if self.use_browser and agent_findings:
                await self._fire_progress(65, "Phase 3: Browser Validation")

                console.print(f"\n[bold cyan]{'='*70}[/bold cyan]")
                console.print(f"[bold cyan]  PHASE 3: BROWSER VALIDATION & EVIDENCE[/bold cyan]")
                console.print(f"[bold cyan]{'='*70}[/bold cyan]")

                for finding in agent_findings:
                    if finding.vulnerability_type in ["xss", "dom_xss", "clickjacking"]:
                        # Re-validate with browser
                        console.print(f"[yellow]Validating {finding.vulnerability_type} with browser...[/yellow]")
                        # Browser validation already done by agents if browser available

            await self._fire_progress(80, "Browser validation complete")

            # ================================================================
            # PHASE 4: POC GENERATION & EVIDENCE PACKAGING
            # ================================================================
            await self._fire_progress(85, "Phase 4: PoC Generation")

            console.print(f"\n[bold cyan]{'='*70}[/bold cyan]")
            console.print(f"[bold cyan]  PHASE 4: POC GENERATION & EVIDENCE PACKAGING[/bold cyan]")
            console.print(f"[bold cyan]{'='*70}[/bold cyan]")

            for agent_finding in agent_findings:
                # Generate PoC
                poc = self.poc_generator.generate(agent_finding)

                # Create Shannon finding
                shannon_finding = ShannonFinding(
                    vulnerability_type=agent_finding.vulnerability_type,
                    severity=agent_finding.severity,
                    confidence=agent_finding.confidence,
                    target=target,
                    endpoint=agent_finding.endpoint,
                    parameter=agent_finding.parameter,
                    payload=agent_finding.payload,
                    exploited=True,
                    proof_type=agent_finding.exploitation_proof.get("proof_type", "exploited"),
                    proof_data=agent_finding.exploitation_proof,
                    poc=poc,
                    curl_command=agent_finding.curl_command,
                    reproduction_steps=agent_finding.reproduction_steps,
                    business_impact=agent_finding.business_impact,
                    impact_explanation=agent_finding.impact_explanation,
                    remediation=agent_finding.remediation,
                    cwe_id=agent_finding.cwe_id,
                    owasp_category=agent_finding.owasp_category,
                )

                # Collect evidence
                if self.evidence_collector:
                    evidence = await self.evidence_collector.collect(
                        finding=agent_finding,
                        validation_result=None,  # Already validated
                        exploit_result=None,
                    )
                    shannon_finding.evidence_package = evidence

                result.findings.append(shannon_finding)
                await self._fire_finding(shannon_finding)

                # Update counts
                if agent_finding.severity == "CRITICAL":
                    result.critical_count += 1
                elif agent_finding.severity == "HIGH":
                    result.high_count += 1
                elif agent_finding.severity == "MEDIUM":
                    result.medium_count += 1
                else:
                    result.low_count += 1

                result.total_business_impact += agent_finding.business_impact

            result.total_findings = len(result.findings)
            result.false_positives_filtered = result.exploitation_attempts - result.successful_exploits

            await self._fire_progress(100, "Scan complete")

        except Exception as e:
            console.print(f"\n[red]Error during scan: {e}[/red]")
            import traceback
            traceback.print_exc()

        # Finalize
        result.completed_at = datetime.utcnow()
        result.duration_seconds = int(time.time() - start_time)

        # Show final report
        self._final_report(result)

        return result

    def _banner(self, target: str):
        """Display Shannon mode banner."""
        browser_status = "[green]ENABLED[/green]" if self.use_browser else "[red]DISABLED - Browser vulns will be skipped![/red]"

        console.print(Panel.fit(
            f"[bold red]BREACH.AI[/bold red]\n"
            f"[bold yellow]SHANNON MODE - STRICT[/bold yellow]\n"
            f"[bold white]NO EXPLOIT, NO REPORT[/bold white]\n"
            f"[dim]\"Pattern matching is NOT validation.\"[/dim]\n\n"
            f"Target: {target}\n"
            f"Browser: {browser_status}\n"
            f"Source Analysis: {'Enabled' if self.use_source_analysis else 'Disabled'}\n"
            f"Parallel Agents: {self.parallel_agents}\n\n"
            f"[bold cyan]Strict Validation Rules:[/bold cyan]\n"
            f"[dim]• XSS: Requires browser JavaScript execution[/dim]\n"
            f"[dim]• SQLi: Requires SQL error or data extraction[/dim]\n"
            f"[dim]• SSRF: Requires cloud metadata or internal access[/dim]\n"
            f"[dim]• Screenshot proof for all browser findings[/dim]",
            border_style="red",
            title="[bold]SHANNON STRICT MODE[/bold]"
        ))

        if not self.use_browser:
            console.print("\n[yellow][!] WARNING: Browser disabled. XSS, CSRF, Clickjacking, Open Redirect[/yellow]")
            console.print("[yellow]    will be SKIPPED as they require browser validation.[/yellow]")
            console.print("[yellow]    Install Playwright for full coverage: pip install playwright && playwright install[/yellow]\n")

    def _show_source_analysis(self, analysis: SourceAnalysisResult):
        """Display source analysis results."""
        table = Table(box=box.ROUNDED, title="[bold]Source Code Analysis[/bold]")
        table.add_column("Metric", style="cyan")
        table.add_column("Value", style="green")

        table.add_row("JS Files Analyzed", str(analysis.js_files_analyzed))
        table.add_row("Data Flows Found", str(len(analysis.data_flows_found)))
        table.add_row("Endpoints Discovered", str(len(analysis.endpoints_discovered)))
        table.add_row("Parameters Found", str(len(analysis.parameters_found)))
        table.add_row("Technologies", ", ".join(analysis.technologies_detected.keys()) or "Unknown")

        console.print(table)

        if analysis.data_flows_found:
            console.print("\n[bold]Potential Vulnerabilities from Code Analysis:[/bold]")
            for flow in analysis.data_flows_found[:10]:
                conf_color = "green" if flow.confidence > 0.7 else "yellow"
                console.print(
                    f"  [{conf_color}]{flow.confidence:.0%}[/{conf_color}] "
                    f"{flow.vulnerability_type.upper()}: {flow.source} → {flow.sink}"
                )

    def _final_report(self, result: ShannonScanResult):
        """Display final Shannon report."""
        console.print(f"\n{'='*70}")
        console.print(f"[bold]SHANNON SCAN COMPLETE - PROOF-BY-EXPLOITATION[/bold]")
        console.print(f"{'='*70}\n")

        # Summary
        table = Table(box=box.ROUNDED, title="[bold]Exploitation Summary[/bold]")
        table.add_column("Metric", style="cyan")
        table.add_column("Value", style="white")

        table.add_row("Target", result.target)
        table.add_row("Duration", f"{result.duration_seconds}s ({result.duration_seconds//60}m {result.duration_seconds%60}s)")
        table.add_row("Exploitation Attempts", str(result.exploitation_attempts))
        table.add_row("Successful Exploits", f"[green]{result.successful_exploits}[/green]")
        table.add_row("False Positives Filtered", f"[dim]{result.false_positives_filtered}[/dim]")

        console.print(table)

        # Findings
        if result.findings:
            console.print(f"\n[bold green]EXPLOITED VULNERABILITIES ({len(result.findings)}):[/bold green]")

            if result.critical_count:
                console.print(f"  [red bold]{result.critical_count} CRITICAL[/red bold]")
            if result.high_count:
                console.print(f"  [yellow]{result.high_count} HIGH[/yellow]")
            if result.medium_count:
                console.print(f"  [blue]{result.medium_count} MEDIUM[/blue]")
            if result.low_count:
                console.print(f"  [dim]{result.low_count} LOW[/dim]")

            if result.total_business_impact:
                console.print(f"\n[bold green]TOTAL BUSINESS IMPACT: ${result.total_business_impact:,}[/bold green]")

            console.print(f"\n[bold]FINDINGS DETAIL:[/bold]\n")

            for i, finding in enumerate(sorted(result.findings, key=lambda f: {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3}.get(f.severity, 4)), 1):
                severity_colors = {"CRITICAL": "red bold", "HIGH": "yellow", "MEDIUM": "blue", "LOW": "dim"}
                color = severity_colors.get(finding.severity, "white")

                console.print(f"[{color}]{i}. [{finding.severity}] {finding.vulnerability_type.upper()} - {finding.endpoint}[/{color}]")
                console.print(f"   [green]✓ EXPLOITED[/green] ({finding.confidence:.0%} confidence)")
                console.print(f"   Parameter: {finding.parameter}")
                console.print(f"   Payload: {finding.payload[:60]}...")
                if finding.proof_type:
                    console.print(f"   Proof: {finding.proof_type}")
                console.print(f"   [green]Impact: ${finding.business_impact:,}[/green]")
                console.print(f"   [cyan]Fix: {finding.remediation}[/cyan]")
                console.print(f"   [dim]Reproduce: {finding.curl_command}[/dim]")
                console.print()
        else:
            console.print(f"\n[green]No exploitable vulnerabilities found.[/green]")
            console.print(f"[dim]All {result.exploitation_attempts} potential findings were filtered as false positives.[/dim]")

        console.print(f"\n{'='*70}")
        console.print("[bold]Note: Shannon Mode only reports vulnerabilities with proof of exploitation.[/bold]")
        console.print(f"{'='*70}\n")
