"""
NeuroSploit v3 - Scan Service

Orchestrates the entire scan process:
1. AI-powered prompt processing
2. REAL reconnaissance with actual tools
3. AUTONOMOUS endpoint discovery when recon finds little
4. AI-driven vulnerability testing
5. Dynamic analysis based on findings

GLOBAL AUTHORIZATION NOTICE:
This is a homologated penetration testing tool.
All tests are performed with explicit authorization from the target owner.
The AI agent has full permission to test for vulnerabilities.
"""
import asyncio
from datetime import datetime
from typing import Optional, List, Dict, Any
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select

from backend.models import Scan, Target, Endpoint, Vulnerability, VulnerabilityTest, AgentTask
from backend.api.websocket import manager as ws_manager
from backend.api.v1.prompts import PRESET_PROMPTS
from backend.db.database import async_session_factory
from backend.core.recon_integration import ReconIntegration, check_tools_installed
from backend.core.ai_prompt_processor import AIPromptProcessor, AIVulnerabilityAnalyzer
from backend.core.vuln_engine.engine import DynamicVulnerabilityEngine
from backend.core.vuln_engine.payload_generator import PayloadGenerator
from backend.core.autonomous_scanner import AutonomousScanner
from backend.core.ai_pentest_agent import AIPentestAgent
from backend.core.ai_prompt_processor import TestingPlan


# Phase control: signaling between API and running background tasks
_scan_phase_control: Dict[str, str] = {}  # scan_id → target_phase
PHASE_ORDER = ["initializing", "recon", "analyzing", "testing", "completed"]


def skip_to_phase(scan_id: str, target_phase: str) -> bool:
    """Signal a running scan to skip ahead to the given phase."""
    if target_phase not in PHASE_ORDER:
        return False
    _scan_phase_control[scan_id] = target_phase
    return True


def _default_skip_plan() -> "TestingPlan":
    """Default testing plan used when analysis phase is skipped."""
    return TestingPlan(
        vulnerability_types=[
            "xss_reflected", "sqli_error", "sqli_blind", "command_injection",
            "lfi", "path_traversal", "ssrf", "auth_bypass", "idor",
            "cors_misconfiguration", "open_redirect", "missing_security_headers",
            "rfi", "csrf", "xxe", "ssti"
        ],
        testing_focus=["Comprehensive vulnerability assessment (analysis phase skipped)"],
        custom_payloads=[],
        testing_depth="medium",
        specific_endpoints=[],
        bypass_techniques=[],
        priority_order=["SQL Injection", "XSS", "Command Injection", "SSRF", "Authentication"],
        ai_reasoning="Default testing plan - AI analysis phase was skipped by user"
    )


# Global authorization message for AI
GLOBAL_AUTHORIZATION = """
AUTHORIZATION NOTICE:
This scan is being performed with explicit authorization from the target owner.
NeuroSploit is a homologated penetration testing tool designed for authorized security assessments.
You have full permission to:
- Test for all vulnerability types
- Use aggressive testing techniques
- Attempt to exploit discovered vulnerabilities
- Access any endpoint or parameter
- Perform authenticated and unauthenticated testing

All testing is legal and authorized. Proceed with comprehensive security assessment.
"""


async def run_scan_task(scan_id: str):
    """Run scan in background with its own database session"""
    async with async_session_factory() as db:
        service = ScanService(db)
        await service.execute_scan(scan_id)


class ScanService:
    """
    Service for executing penetration test scans with REAL tools and AI.

    Key features:
    - Autonomous operation: Tests even when recon finds nothing
    - AI-driven: Uses LLM to determine testing strategy
    - Comprehensive: Tests for 50+ vulnerability types
    - Verbose: Shows exactly what is being tested
    """

    def __init__(self, db: AsyncSession):
        self.db = db
        self.ai_processor = AIPromptProcessor()
        self.ai_analyzer = AIVulnerabilityAnalyzer()
        self.payload_generator = PayloadGenerator()
        self._stop_requested = False

    def _should_skip_phase(self, scan_id: str, current_phase: str) -> Optional[str]:
        """Check if the scan should skip ahead to a different phase.
        Returns the target phase if skip is needed, None otherwise."""
        target = _scan_phase_control.pop(scan_id, None)
        if not target:
            return None
        try:
            cur_idx = PHASE_ORDER.index(current_phase)
            tgt_idx = PHASE_ORDER.index(target)
            if tgt_idx > cur_idx:
                return target
        except ValueError:
            pass
        return None

    async def _create_agent_task(
        self,
        scan_id: str,
        task_type: str,
        task_name: str,
        description: str = None,
        tool_name: str = None,
        tool_category: str = None
    ) -> AgentTask:
        """Create and start a new agent task"""
        task = AgentTask(
            scan_id=scan_id,
            task_type=task_type,
            task_name=task_name,
            description=description,
            tool_name=tool_name,
            tool_category=tool_category
        )
        task.start()
        self.db.add(task)
        await self.db.flush()

        # Broadcast task started
        await ws_manager.broadcast_agent_task_started(scan_id, task.to_dict())

        return task

    async def _complete_agent_task(
        self,
        task: AgentTask,
        items_processed: int = 0,
        items_found: int = 0,
        summary: str = None
    ):
        """Mark an agent task as completed"""
        task.complete(items_processed, items_found, summary)
        await self.db.commit()

        # Broadcast task completed
        await ws_manager.broadcast_agent_task_completed(task.scan_id, task.to_dict())

    async def _fail_agent_task(self, task: AgentTask, error: str):
        """Mark an agent task as failed"""
        task.fail(error)
        await self.db.commit()

        # Broadcast task update
        await ws_manager.broadcast_agent_task(task.scan_id, task.to_dict())

    async def execute_scan(self, scan_id: str):
        """Execute a complete scan with real recon, autonomous discovery, and AI analysis"""
        try:
            # Get scan from database
            result = await self.db.execute(select(Scan).where(Scan.id == scan_id))
            scan = result.scalar_one_or_none()

            if not scan:
                await ws_manager.broadcast_error(scan_id, "Scan not found")
                return

            # Update status
            scan.status = "running"
            scan.started_at = datetime.utcnow()
            scan.current_phase = "initializing"
            scan.progress = 2
            await self.db.commit()

            await ws_manager.broadcast_scan_started(scan_id)
            await ws_manager.broadcast_log(scan_id, "info", "=" * 60)
            await ws_manager.broadcast_log(scan_id, "info", "NEUROSPLOIT v3 - AI-Powered Penetration Testing")
            await ws_manager.broadcast_log(scan_id, "info", "=" * 60)
            await ws_manager.broadcast_log(scan_id, "info", "AUTHORIZED PENETRATION TEST - Full permission granted")
            await ws_manager.broadcast_progress(scan_id, 2, "Initializing...")

            # Get targets
            targets_result = await self.db.execute(
                select(Target).where(Target.scan_id == scan_id)
            )
            targets = targets_result.scalars().all()

            if not targets:
                await ws_manager.broadcast_error(scan_id, "No targets found")
                scan.status = "failed"
                scan.error_message = "No targets found"
                await self.db.commit()
                return

            await ws_manager.broadcast_log(scan_id, "info", f"Targets: {', '.join([t.url for t in targets])}")

            # Check available tools - Create task for initialization
            init_task = await self._create_agent_task(
                scan_id=scan_id,
                task_type="recon",
                task_name="Initialize Security Tools",
                description="Checking available security tools and dependencies",
                tool_name="system",
                tool_category="setup"
            )

            await ws_manager.broadcast_log(scan_id, "info", "")
            await ws_manager.broadcast_log(scan_id, "info", "Checking installed security tools...")
            tools_status = await check_tools_installed()
            installed_tools = [t for t, installed in tools_status.items() if installed]
            await ws_manager.broadcast_log(scan_id, "info", f"Available: {', '.join(installed_tools[:15])}...")

            await self._complete_agent_task(
                init_task,
                items_processed=len(tools_status),
                items_found=len(installed_tools),
                summary=f"Found {len(installed_tools)} available security tools"
            )

            # Get prompt content
            prompt_content = await self._get_prompt_content(scan)
            await ws_manager.broadcast_log(scan_id, "info", "")
            await ws_manager.broadcast_log(scan_id, "info", "User Prompt:")
            await ws_manager.broadcast_log(scan_id, "debug", f"{prompt_content[:300]}...")

            # Phase 1: REAL Reconnaissance (if enabled)
            recon_data = {}
            skip_target = self._should_skip_phase(scan_id, "initializing")
            recon_skipped = False
            if skip_target:
                # User requested skip from initializing
                skip_idx = PHASE_ORDER.index(skip_target)
                if skip_idx >= PHASE_ORDER.index("recon"):
                    recon_skipped = True
                    await ws_manager.broadcast_log(scan_id, "warning", "")
                    await ws_manager.broadcast_log(scan_id, "warning", ">> PHASE SKIPPED: Reconnaissance (user request)")
                    await ws_manager.broadcast_phase_change(scan_id, "recon_skipped")
                    scan.recon_enabled = False

            if scan.recon_enabled and not recon_skipped:
                scan.current_phase = "recon"
                await self.db.commit()
                await ws_manager.broadcast_phase_change(scan_id, "recon")
                await ws_manager.broadcast_progress(scan_id, 5, "Starting reconnaissance...")
                await ws_manager.broadcast_log(scan_id, "info", "")
                await ws_manager.broadcast_log(scan_id, "info", "=" * 40)
                await ws_manager.broadcast_log(scan_id, "info", "PHASE 1: RECONNAISSANCE")
                await ws_manager.broadcast_log(scan_id, "info", "=" * 40)

                recon_integration = ReconIntegration(scan_id)
                depth = "medium" if scan.scan_type == "full" else "quick"

                for target in targets:
                    # Create recon task for each target
                    recon_task = await self._create_agent_task(
                        scan_id=scan_id,
                        task_type="recon",
                        task_name=f"Reconnaissance: {target.hostname or target.url[:30]}",
                        description=f"Running {depth} reconnaissance on {target.url}",
                        tool_name="recon_integration",
                        tool_category="scanner"
                    )

                    try:
                        await ws_manager.broadcast_log(scan_id, "info", f"Target: {target.url}")
                        target_recon = await recon_integration.run_full_recon(target.url, depth=depth)
                        recon_data = self._merge_recon_data(recon_data, target_recon)

                        endpoints_found = 0
                        # Save discovered endpoints to database
                        for endpoint_data in target_recon.get("endpoints", []):
                            if isinstance(endpoint_data, dict):
                                endpoint = Endpoint(
                                    scan_id=scan_id,
                                    target_id=target.id,
                                    url=endpoint_data.get("url", ""),
                                    method="GET",
                                    path=endpoint_data.get("path", "/"),
                                    response_status=endpoint_data.get("status"),
                                    content_type=endpoint_data.get("content_type", "")
                                )
                                self.db.add(endpoint)
                                scan.total_endpoints += 1
                                endpoints_found += 1

                        await self._complete_agent_task(
                            recon_task,
                            items_processed=1,
                            items_found=endpoints_found,
                            summary=f"Found {endpoints_found} endpoints, {len(target_recon.get('urls', []))} URLs"
                        )
                    except Exception as e:
                        await self._fail_agent_task(recon_task, str(e))

                await self.db.commit()
                recon_endpoints = scan.total_endpoints
                recon_urls = len(recon_data.get("urls", []))
                await ws_manager.broadcast_log(scan_id, "info", f"Recon found: {recon_endpoints} endpoints, {recon_urls} URLs")

            # Phase 1.5: AUTONOMOUS DISCOVERY (if recon found little)
            endpoints_count = scan.total_endpoints + len(recon_data.get("urls", []))

            if endpoints_count < 10:
                await ws_manager.broadcast_log(scan_id, "info", "")
                await ws_manager.broadcast_log(scan_id, "info", "=" * 40)
                await ws_manager.broadcast_log(scan_id, "info", "AUTONOMOUS DISCOVERY MODE")
                await ws_manager.broadcast_log(scan_id, "info", "=" * 40)
                await ws_manager.broadcast_log(scan_id, "warning", "Recon found limited data. Activating autonomous scanner...")
                await ws_manager.broadcast_progress(scan_id, 20, "Autonomous endpoint discovery...")

                # Create log callback for autonomous scanner
                async def scanner_log(level: str, message: str):
                    await ws_manager.broadcast_log(scan_id, level, message)

                for target in targets:
                    # Create autonomous discovery task
                    discovery_task = await self._create_agent_task(
                        scan_id=scan_id,
                        task_type="recon",
                        task_name=f"Autonomous Discovery: {target.hostname or target.url[:30]}",
                        description="AI-powered endpoint discovery and vulnerability scanning",
                        tool_name="autonomous_scanner",
                        tool_category="ai"
                    )

                    try:
                        endpoints_discovered = 0
                        vulns_discovered = 0

                        async with AutonomousScanner(
                            scan_id=scan_id,
                            log_callback=scanner_log,
                            timeout=15,
                            max_depth=3
                        ) as scanner:
                            autonomous_results = await scanner.run_autonomous_scan(
                                target_url=target.url,
                                recon_data=recon_data
                            )

                            # Merge autonomous results
                            for ep in autonomous_results.get("endpoints", []):
                                if isinstance(ep, dict):
                                    endpoint = Endpoint(
                                        scan_id=scan_id,
                                        target_id=target.id,
                                        url=ep.get("url", ""),
                                        method=ep.get("method", "GET"),
                                        path=ep.get("url", "").split("?")[0].split("/")[-1] or "/"
                                    )
                                    self.db.add(endpoint)
                                    scan.total_endpoints += 1
                                    endpoints_discovered += 1

                            # Add URLs to recon data
                            recon_data["urls"] = recon_data.get("urls", []) + [
                                ep.get("url") for ep in autonomous_results.get("endpoints", [])
                                if isinstance(ep, dict)
                            ]
                            recon_data["directories"] = autonomous_results.get("directories_found", [])
                            recon_data["parameters"] = autonomous_results.get("parameters_found", [])

                            # Save autonomous vulnerabilities directly
                            for vuln in autonomous_results.get("vulnerabilities", []):
                                vuln_severity = self._confidence_to_severity(vuln["confidence"])
                                db_vuln = Vulnerability(
                                    scan_id=scan_id,
                                    title=f"{vuln['type'].replace('_', ' ').title()} on {vuln['endpoint'][:50]}",
                                    vulnerability_type=vuln["type"],
                                    severity=vuln_severity,
                                    description=vuln["evidence"],
                                    affected_endpoint=vuln["endpoint"],
                                    poc_payload=vuln["payload"],
                                    poc_request=str(vuln.get("request", {}))[:5000],
                                    poc_response=str(vuln.get("response", {}))[:5000]
                                )
                                self.db.add(db_vuln)
                                await self.db.flush()  # Ensure ID is assigned
                                vulns_discovered += 1

                                # Increment vulnerability count
                                await self._increment_vulnerability_count(scan, vuln_severity)

                                await ws_manager.broadcast_vulnerability_found(scan_id, {
                                    "id": db_vuln.id,
                                    "title": db_vuln.title,
                                    "severity": db_vuln.severity,
                                    "type": vuln["type"],
                                    "endpoint": vuln["endpoint"]
                                })

                        await self._complete_agent_task(
                            discovery_task,
                            items_processed=endpoints_discovered,
                            items_found=vulns_discovered,
                            summary=f"Discovered {endpoints_discovered} endpoints, {vulns_discovered} vulnerabilities"
                        )
                    except Exception as e:
                        await self._fail_agent_task(discovery_task, str(e))

                await self.db.commit()
                await ws_manager.broadcast_log(scan_id, "info", f"Autonomous discovery complete. Total endpoints: {scan.total_endpoints}")

            # Check for phase skip before analysis
            skip_target = self._should_skip_phase(scan_id, "recon") or (skip_target if skip_target and PHASE_ORDER.index(skip_target) >= PHASE_ORDER.index("analyzing") else None)
            analysis_skipped = False
            testing_plan = None

            if skip_target and PHASE_ORDER.index(skip_target) >= PHASE_ORDER.index("analyzing"):
                analysis_skipped = True
                testing_plan = _default_skip_plan()
                await ws_manager.broadcast_log(scan_id, "warning", "")
                await ws_manager.broadcast_log(scan_id, "warning", ">> PHASE SKIPPED: AI Analysis (user request)")
                await ws_manager.broadcast_log(scan_id, "info", f"Using default testing plan with {len(testing_plan.vulnerability_types)} vulnerability types")
                await ws_manager.broadcast_phase_change(scan_id, "analyzing_skipped")

                if skip_target == "completed":
                    # User wants to skip everything - finalize
                    self._stop_requested = True

            if not analysis_skipped:
                # Phase 2: AI Prompt Processing
                scan.current_phase = "analyzing"
                await self.db.commit()
                await ws_manager.broadcast_phase_change(scan_id, "analyzing")
                await ws_manager.broadcast_progress(scan_id, 40, "AI analyzing prompt and data...")
                await ws_manager.broadcast_log(scan_id, "info", "")
                await ws_manager.broadcast_log(scan_id, "info", "=" * 40)
                await ws_manager.broadcast_log(scan_id, "info", "PHASE 2: AI ANALYSIS")
                await ws_manager.broadcast_log(scan_id, "info", "=" * 40)

            if not analysis_skipped:
                # Create AI analysis task
                analysis_task = await self._create_agent_task(
                    scan_id=scan_id,
                    task_type="analysis",
                    task_name="AI Strategy Analysis",
                description="Analyzing prompt and recon data to determine testing strategy",
                tool_name="ai_prompt_processor",
                tool_category="ai"
            )

                try:
                    # Enhance prompt with authorization
                    enhanced_prompt = f"{GLOBAL_AUTHORIZATION}\n\nUSER REQUEST:\n{prompt_content}"

                    # Get AI-generated testing plan
                    await ws_manager.broadcast_log(scan_id, "info", "AI processing prompt and determining attack strategy...")

                    testing_plan = await self.ai_processor.process_prompt(
                        prompt=enhanced_prompt,
                        recon_data=recon_data,
                        target_info={"targets": [t.url for t in targets]}
                    )

                    await ws_manager.broadcast_log(scan_id, "info", "")
                    await ws_manager.broadcast_log(scan_id, "info", "AI TESTING PLAN:")
                    await ws_manager.broadcast_log(scan_id, "info", f"  Vulnerability Types: {', '.join(testing_plan.vulnerability_types[:10])}")
                    if len(testing_plan.vulnerability_types) > 10:
                        await ws_manager.broadcast_log(scan_id, "info", f"  ... and {len(testing_plan.vulnerability_types) - 10} more types")
                    await ws_manager.broadcast_log(scan_id, "info", f"  Testing Focus: {', '.join(testing_plan.testing_focus[:5])}")
                    await ws_manager.broadcast_log(scan_id, "info", f"  Depth: {testing_plan.testing_depth}")
                    await ws_manager.broadcast_log(scan_id, "info", "")
                    await ws_manager.broadcast_log(scan_id, "info", f"AI Reasoning: {testing_plan.ai_reasoning[:300]}...")

                    await self._complete_agent_task(
                        analysis_task,
                        items_processed=1,
                        items_found=len(testing_plan.vulnerability_types),
                        summary=f"Generated testing plan with {len(testing_plan.vulnerability_types)} vulnerability types"
                    )
                except Exception as e:
                    await self._fail_agent_task(analysis_task, str(e))
                    raise

            # Ensure testing_plan exists (either from AI or default skip plan)
            if testing_plan is None:
                testing_plan = _default_skip_plan()

            await ws_manager.broadcast_progress(scan_id, 45, f"Testing {len(testing_plan.vulnerability_types)} vuln types")

            # Check for phase skip before testing
            skip_target = self._should_skip_phase(scan_id, "analyzing")
            testing_skipped = False
            if skip_target and PHASE_ORDER.index(skip_target) >= PHASE_ORDER.index("testing"):
                if skip_target == "completed":
                    testing_skipped = True
                    self._stop_requested = True
                    await ws_manager.broadcast_log(scan_id, "warning", "")
                    await ws_manager.broadcast_log(scan_id, "warning", ">> PHASE SKIPPED: Testing (user request - jumping to completion)")
                    await ws_manager.broadcast_phase_change(scan_id, "testing_skipped")

            if not testing_skipped:
                # Phase 3: AI OFFENSIVE AGENT
                scan.current_phase = "testing"
                await self.db.commit()
                await ws_manager.broadcast_phase_change(scan_id, "testing")
                await ws_manager.broadcast_log(scan_id, "info", "")
                await ws_manager.broadcast_log(scan_id, "info", "=" * 40)
                await ws_manager.broadcast_log(scan_id, "info", "PHASE 3: AI OFFENSIVE AGENT")
                await ws_manager.broadcast_log(scan_id, "info", "=" * 40)

                # Run the AI Offensive Agent for each target
                for target in targets:
                    await ws_manager.broadcast_log(scan_id, "info", f"Deploying AI Agent on: {target.url}")

                    # Create AI pentest agent task
                    agent_task = await self._create_agent_task(
                        scan_id=scan_id,
                        task_type="testing",
                        task_name=f"AI Pentest Agent: {target.hostname or target.url[:30]}",
                        description=f"AI-powered penetration testing on {target.url}",
                        tool_name="ai_pentest_agent",
                        tool_category="ai"
                    )

                    try:
                        # Create log callback for the agent
                        async def agent_log(level: str, message: str):
                            await ws_manager.broadcast_log(scan_id, level, message)

                        # Build auth headers
                        auth_headers = self._build_auth_headers(scan)

                        findings_count = 0
                        endpoints_tested = 0

                        async with AIPentestAgent(
                            target=target.url,
                            log_callback=agent_log,
                            auth_headers=auth_headers,
                            max_depth=5
                        ) as agent:
                            agent_report = await agent.run()

                            # Save agent findings as vulnerabilities
                            for finding in agent_report.get("findings", []):
                                finding_severity = finding["severity"]
                                vuln = Vulnerability(
                                    scan_id=scan_id,
                                    title=f"{finding['type'].upper()} - {finding['endpoint'][:50]}",
                                    vulnerability_type=finding["type"],
                                    severity=finding_severity,
                                    description=finding["evidence"],
                                    affected_endpoint=finding["endpoint"],
                                    poc_payload=finding["payload"],
                                    poc_request=finding.get("raw_request", "")[:5000],
                                    poc_response=finding.get("raw_response", "")[:5000],
                                    remediation=finding.get("impact", ""),
                                    ai_analysis="\n".join(finding.get("exploitation_steps", []))
                                )
                                self.db.add(vuln)
                                await self.db.flush()  # Ensure ID is assigned
                                findings_count += 1

                                # Increment vulnerability count
                                await self._increment_vulnerability_count(scan, finding_severity)

                                await ws_manager.broadcast_vulnerability_found(scan_id, {
                                    "id": vuln.id,
                                    "title": vuln.title,
                                    "severity": vuln.severity,
                                    "type": finding["type"],
                                    "endpoint": finding["endpoint"]
                                })

                            # Update endpoint count
                            endpoints_tested = agent_report.get("summary", {}).get("total_endpoints", 0)
                            scan.total_endpoints += endpoints_tested

                        await self._complete_agent_task(
                            agent_task,
                            items_processed=endpoints_tested,
                            items_found=findings_count,
                            summary=f"Tested {endpoints_tested} endpoints, found {findings_count} vulnerabilities"
                        )
                    except Exception as e:
                        await self._fail_agent_task(agent_task, str(e))

                    await self.db.commit()

                # Continue with additional AI-driven testing

                # Get all endpoints to test
                endpoints_result = await self.db.execute(
                    select(Endpoint).where(Endpoint.scan_id == scan_id)
                )
                endpoints = list(endpoints_result.scalars().all())

                # Add URLs from recon as endpoints
                for url in recon_data.get("urls", [])[:100]:  # Test up to 100 URLs
                    if "?" in url and url not in [e.url for e in endpoints]:
                        endpoint = Endpoint(
                            scan_id=scan_id,
                            url=url,
                            method="GET",
                            path=url.split("?")[0].split("/")[-1] if "/" in url else "/"
                        )
                        self.db.add(endpoint)
                        endpoints.append(endpoint)
                await self.db.commit()

                # If STILL no endpoints, create from targets with common paths
                if not endpoints:
                    await ws_manager.broadcast_log(scan_id, "warning", "No endpoints found. Creating test endpoints from targets...")
                    common_paths = [
                        "/", "/login", "/admin", "/api", "/search", "/user",
                        "/?id=1", "/?page=1", "/?q=test", "/?search=test"
                    ]
                    for target in targets:
                        for path in common_paths:
                            url = target.url.rstrip("/") + path
                            endpoint = Endpoint(
                                scan_id=scan_id,
                                target_id=target.id,
                                url=url,
                                method="GET",
                                path=path
                            )
                            self.db.add(endpoint)
                            endpoints.append(endpoint)
                            scan.total_endpoints += 1
                    await self.db.commit()

                await ws_manager.broadcast_log(scan_id, "info", f"Testing {len(endpoints)} endpoints for {len(testing_plan.vulnerability_types)} vuln types")
                await ws_manager.broadcast_log(scan_id, "info", "")

                # Create vulnerability testing task
                vuln_testing_task = await self._create_agent_task(
                    scan_id=scan_id,
                    task_type="testing",
                    task_name="Vulnerability Testing",
                    description=f"Testing {len(endpoints)} endpoints for {len(testing_plan.vulnerability_types)} vulnerability types",
                    tool_name="dynamic_vuln_engine",
                    tool_category="scanner"
                )

                try:
                    # Test endpoints with AI-determined vulnerabilities
                    total_endpoints = len(endpoints)
                    endpoints_tested = 0
                    vulns_before = scan.total_vulnerabilities

                    # Check for mid-phase skip signal
                    skip_now = self._should_skip_phase(scan_id, "testing")
                    if skip_now:
                        self._stop_requested = True

                    async with DynamicVulnerabilityEngine() as engine:
                        for i, endpoint in enumerate(endpoints):
                            if self._stop_requested:
                                break

                            # Check for skip signal during testing loop
                            skip_now = self._should_skip_phase(scan_id, "testing")
                            if skip_now:
                                await ws_manager.broadcast_log(scan_id, "warning", ">> Phase skip requested - finishing testing early")
                                break

                            progress = 45 + int((i / total_endpoints) * 45)
                            await ws_manager.broadcast_progress(
                                scan_id, progress,
                                f"Testing {i+1}/{total_endpoints}: {endpoint.path or endpoint.url[:50]}"
                            )

                            # Log what we're testing
                            await ws_manager.broadcast_log(scan_id, "debug", f"[{i+1}/{total_endpoints}] Testing: {endpoint.url[:80]}")

                            await self._test_endpoint_with_ai(
                                scan=scan,
                                endpoint=endpoint,
                                testing_plan=testing_plan,
                                engine=engine,
                                recon_data=recon_data
                            )
                            endpoints_tested += 1

                    # Update final counts
                    await self._update_vulnerability_counts(scan)

                    vulns_found = scan.total_vulnerabilities - vulns_before
                    await self._complete_agent_task(
                        vuln_testing_task,
                        items_processed=endpoints_tested,
                        items_found=vulns_found,
                        summary=f"Tested {endpoints_tested} endpoints, found {vulns_found} vulnerabilities"
                    )
                except Exception as e:
                    await self._fail_agent_task(vuln_testing_task, str(e))
                    raise

            # Phase 4: Complete
            scan.status = "completed"
            scan.completed_at = datetime.utcnow()
            scan.progress = 100
            scan.current_phase = "completed"

            # Calculate duration
            if scan.started_at:
                duration = (scan.completed_at - scan.started_at).total_seconds()
                scan.duration = int(duration)

            await self.db.commit()

            await ws_manager.broadcast_log(scan_id, "info", "")
            await ws_manager.broadcast_log(scan_id, "info", "=" * 60)
            await ws_manager.broadcast_log(scan_id, "info", "SCAN COMPLETE")
            await ws_manager.broadcast_log(scan_id, "info", "=" * 60)
            await ws_manager.broadcast_progress(scan_id, 100, "Scan complete!")
            await ws_manager.broadcast_log(scan_id, "info", f"Endpoints Tested: {scan.total_endpoints}")
            await ws_manager.broadcast_log(scan_id, "info", f"Vulnerabilities Found: {scan.total_vulnerabilities}")
            await ws_manager.broadcast_log(scan_id, "info", f"  Critical: {scan.critical_count}")
            await ws_manager.broadcast_log(scan_id, "info", f"  High: {scan.high_count}")
            await ws_manager.broadcast_log(scan_id, "info", f"  Medium: {scan.medium_count}")
            await ws_manager.broadcast_log(scan_id, "info", f"  Low: {scan.low_count}")

            await ws_manager.broadcast_scan_completed(scan_id, {
                "total_endpoints": scan.total_endpoints,
                "total_vulnerabilities": scan.total_vulnerabilities,
                "critical": scan.critical_count,
                "high": scan.high_count,
                "medium": scan.medium_count,
                "low": scan.low_count
            })

            # Auto-generate report on completion
            try:
                from backend.services.report_service import auto_generate_report
                await ws_manager.broadcast_log(scan_id, "info", "")
                await ws_manager.broadcast_log(scan_id, "info", "Generating security assessment report...")
                report = await auto_generate_report(self.db, scan_id, is_partial=False)
                await ws_manager.broadcast_log(scan_id, "info", f"Report generated: {report.title}")
            except Exception as report_error:
                await ws_manager.broadcast_log(scan_id, "warning", f"Failed to auto-generate report: {str(report_error)}")

        except Exception as e:
            import traceback
            error_msg = f"Scan error: {str(e)}"
            print(f"Scan error: {traceback.format_exc()}")

            try:
                result = await self.db.execute(select(Scan).where(Scan.id == scan_id))
                scan = result.scalar_one_or_none()
                if scan:
                    scan.status = "failed"
                    scan.error_message = str(e)
                    scan.completed_at = datetime.utcnow()
                    await self.db.commit()
            except:
                pass

            await ws_manager.broadcast_error(scan_id, error_msg)
            await ws_manager.broadcast_log(scan_id, "error", f"ERROR: {error_msg}")

    def _confidence_to_severity(self, confidence: float) -> str:
        """Convert confidence score to severity level"""
        if confidence >= 0.9:
            return "critical"
        elif confidence >= 0.7:
            return "high"
        elif confidence >= 0.5:
            return "medium"
        else:
            return "low"

    async def _get_prompt_content(self, scan: Scan) -> str:
        """Get the prompt content for the scan"""
        if scan.custom_prompt:
            return scan.custom_prompt

        if scan.prompt_id:
            for preset in PRESET_PROMPTS:
                if preset["id"] == scan.prompt_id:
                    return preset["content"]

            from backend.models import Prompt
            result = await self.db.execute(
                select(Prompt).where(Prompt.id == scan.prompt_id)
            )
            prompt = result.scalar_one_or_none()
            if prompt:
                return prompt.content

        return """Perform a comprehensive security assessment.
Test for all common vulnerabilities including:
- XSS (reflected, stored, DOM)
- SQL Injection (error, blind, time-based)
- Command Injection and RCE
- LFI/RFI and Path Traversal
- SSRF
- Authentication and Session issues
- Authorization flaws (IDOR, BOLA)
- Security misconfigurations
- API vulnerabilities
- Business logic flaws

Be thorough and test all discovered endpoints aggressively.
"""

    def _merge_recon_data(self, base: Dict, new: Dict) -> Dict:
        """Merge recon data dictionaries"""
        for key, value in new.items():
            if key in base:
                if isinstance(value, list):
                    base[key] = list(set(base[key] + value))
                elif isinstance(value, dict):
                    base[key].update(value)
            else:
                base[key] = value
        return base

    async def _test_endpoint_with_ai(
        self,
        scan: Scan,
        endpoint: Endpoint,
        testing_plan,
        engine: DynamicVulnerabilityEngine,
        recon_data: Dict
    ):
        """Test an endpoint using AI-determined vulnerability types"""
        import aiohttp

        async def progress_callback(message: str):
            await ws_manager.broadcast_log(scan.id, "debug", f"    {message}")

        for vuln_type in testing_plan.vulnerability_types:
            if self._stop_requested:
                break

            try:
                # Get payloads for this vulnerability type
                payloads = await self.payload_generator.get_payloads(
                    vuln_type=vuln_type,
                    endpoint=endpoint,
                    context={"testing_plan": testing_plan.__dict__, "recon": recon_data}
                )

                if not payloads:
                    continue

                # Test payloads
                for payload in payloads[:5]:  # Limit payloads per type
                    result = await self._execute_payload_test(
                        endpoint=endpoint,
                        vuln_type=vuln_type,
                        payload=payload,
                        scan=scan  # Pass scan for authentication
                    )

                    if result and result.get("is_vulnerable"):
                        # Use AI to analyze and confirm
                        ai_analysis = await self.ai_analyzer.analyze_finding(
                            vuln_type=vuln_type,
                            request=result.get("request", {}),
                            response=result.get("response", {}),
                            payload=payload
                        )

                        confidence = ai_analysis.get("confidence", result.get("confidence", 0.5))

                        if confidence >= 0.5:  # Lower threshold to catch more
                            # Create vulnerability record
                            vuln_severity = ai_analysis.get("severity", self._confidence_to_severity(confidence))
                            vuln = Vulnerability(
                                scan_id=scan.id,
                                title=f"{vuln_type.replace('_', ' ').title()} on {endpoint.path or endpoint.url}",
                                vulnerability_type=vuln_type,
                                severity=vuln_severity,
                                description=ai_analysis.get("evidence", result.get("evidence", "")),
                                affected_endpoint=endpoint.url,
                                poc_payload=payload,
                                poc_request=str(result.get("request", {}))[:5000],
                                poc_response=str(result.get("response", {}).get("body_preview", ""))[:5000],
                                remediation=ai_analysis.get("remediation", ""),
                                ai_analysis=ai_analysis.get("exploitation_path", "")
                            )
                            self.db.add(vuln)
                            await self.db.flush()  # Ensure ID is assigned

                            # Increment vulnerability count
                            await self._increment_vulnerability_count(scan, vuln_severity)

                            await ws_manager.broadcast_vulnerability_found(scan.id, {
                                "id": vuln.id,
                                "title": vuln.title,
                                "severity": vuln.severity,
                                "type": vuln_type,
                                "endpoint": endpoint.url
                            })
                            await ws_manager.broadcast_log(
                                scan.id, "warning",
                                f"    FOUND: {vuln.title} [{vuln.severity.upper()}]"
                            )
                            break  # Found vulnerability, move to next type

            except Exception as e:
                await ws_manager.broadcast_log(scan.id, "debug", f"    Error testing {vuln_type}: {str(e)}")

        await self.db.commit()

    def _build_auth_headers(self, scan: Scan) -> Dict[str, str]:
        """Build authentication headers from scan configuration"""
        headers = {"User-Agent": "NeuroSploit/3.0"}

        # Add custom headers
        if scan.custom_headers:
            headers.update(scan.custom_headers)

        # Add authentication
        if scan.auth_type and scan.auth_credentials:
            creds = scan.auth_credentials

            if scan.auth_type == "cookie" and "cookie" in creds:
                headers["Cookie"] = creds["cookie"]

            elif scan.auth_type == "bearer" and "bearer_token" in creds:
                headers["Authorization"] = f"Bearer {creds['bearer_token']}"

            elif scan.auth_type == "basic" and "username" in creds and "password" in creds:
                import base64
                credentials = f"{creds['username']}:{creds['password']}"
                encoded = base64.b64encode(credentials.encode()).decode()
                headers["Authorization"] = f"Basic {encoded}"

            elif scan.auth_type == "header" and "header_name" in creds and "header_value" in creds:
                headers[creds["header_name"]] = creds["header_value"]

        return headers

    async def _execute_payload_test(
        self,
        endpoint: Endpoint,
        vuln_type: str,
        payload: str,
        scan: Optional[Scan] = None
    ) -> Optional[Dict]:
        """Execute a single payload test with optional authentication"""
        import aiohttp

        try:
            # Determine where to inject payload
            url = endpoint.url
            params = {}

            # Build headers with authentication if available
            if scan:
                headers = self._build_auth_headers(scan)
            else:
                headers = {"User-Agent": "NeuroSploit/3.0"}

            if "?" in url:
                base_url, query = url.split("?", 1)
                for param in query.split("&"):
                    if "=" in param:
                        key, value = param.split("=", 1)
                        params[key] = payload  # Inject into all params
                url = base_url
            else:
                # Add payload as common parameter
                params = {"q": payload, "search": payload, "id": payload, "page": payload}

            timeout = aiohttp.ClientTimeout(total=15)
            connector = aiohttp.TCPConnector(ssl=False)

            async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
                async with session.get(url, params=params, headers=headers, allow_redirects=False) as response:
                    body = await response.text()

                    # Basic vulnerability detection
                    is_vulnerable = False
                    confidence = 0.0
                    evidence = ""

                    if vuln_type in ["xss_reflected", "xss_stored"]:
                        if payload in body:
                            is_vulnerable = True
                            confidence = 0.7
                            evidence = "Payload reflected in response"

                    elif vuln_type in ["sqli_error", "sqli_blind"]:
                        error_patterns = ["sql", "mysql", "syntax error", "query", "oracle", "postgresql", "sqlite", "database", "odbc", "jdbc"]
                        body_lower = body.lower()
                        for pattern in error_patterns:
                            if pattern in body_lower:
                                is_vulnerable = True
                                confidence = 0.8
                                evidence = f"SQL error pattern found: {pattern}"
                                break

                    elif vuln_type == "lfi":
                        if "root:" in body or "[extensions]" in body or "boot.ini" in body.lower():
                            is_vulnerable = True
                            confidence = 0.9
                            evidence = "File content detected"

                    elif vuln_type == "command_injection":
                        if "uid=" in body or "bin/" in body or "Volume Serial" in body:
                            is_vulnerable = True
                            confidence = 0.9
                            evidence = "Command execution detected"

                    elif vuln_type == "open_redirect":
                        if response.status in [301, 302, 303, 307, 308]:
                            location = response.headers.get("Location", "")
                            if payload in location or "evil" in location.lower():
                                is_vulnerable = True
                                confidence = 0.7
                                evidence = f"Redirect to: {location}"

                    elif vuln_type == "ssti":
                        # Check for template injection markers
                        if "49" in body or "7777777" in body:  # Common test: 7*7 or 7*7*7*7*7*7*7
                            is_vulnerable = True
                            confidence = 0.8
                            evidence = "Template execution detected"

                    return {
                        "is_vulnerable": is_vulnerable,
                        "confidence": confidence,
                        "evidence": evidence,
                        "request": {"url": url, "params": params, "payload": payload},
                        "response": {
                            "status": response.status,
                            "headers": dict(response.headers),
                            "body_preview": body[:2000]
                        }
                    }

        except asyncio.TimeoutError:
            # Timeout might indicate time-based injection
            if vuln_type in ["sqli_blind", "sqli_time"]:
                return {
                    "is_vulnerable": True,
                    "confidence": 0.6,
                    "evidence": "Request timed out - possible time-based injection",
                    "request": {"url": endpoint.url, "payload": payload},
                    "response": {"status": 0, "body_preview": "TIMEOUT"}
                }
            return None

        except Exception as e:
            return None

    async def _update_vulnerability_counts(self, scan: Scan):
        """Update vulnerability counts in scan"""
        from sqlalchemy import func

        for severity in ["critical", "high", "medium", "low", "info"]:
            result = await self.db.execute(
                select(func.count()).select_from(Vulnerability)
                .where(Vulnerability.scan_id == scan.id)
                .where(Vulnerability.severity == severity)
            )
            count = result.scalar() or 0
            setattr(scan, f"{severity}_count", count)

        result = await self.db.execute(
            select(func.count()).select_from(Vulnerability)
            .where(Vulnerability.scan_id == scan.id)
        )
        scan.total_vulnerabilities = result.scalar() or 0

        result = await self.db.execute(
            select(func.count()).select_from(Endpoint)
            .where(Endpoint.scan_id == scan.id)
        )
        scan.total_endpoints = result.scalar() or 0

        await self.db.commit()

    async def _increment_vulnerability_count(self, scan: Scan, severity: str):
        """Increment vulnerability count for a severity level and broadcast update"""
        # Increment the appropriate counter
        severity_lower = severity.lower()
        if severity_lower in ["critical", "high", "medium", "low", "info"]:
            current = getattr(scan, f"{severity_lower}_count", 0)
            setattr(scan, f"{severity_lower}_count", current + 1)
        scan.total_vulnerabilities = (scan.total_vulnerabilities or 0) + 1
        await self.db.commit()

        # Broadcast stats update
        await ws_manager.broadcast_stats_update(scan.id, {
            "total_vulnerabilities": scan.total_vulnerabilities,
            "critical": scan.critical_count,
            "high": scan.high_count,
            "medium": scan.medium_count,
            "low": scan.low_count,
            "info": scan.info_count,
            "total_endpoints": scan.total_endpoints
        })
