"""
AIPTX AI Strategy Generator - Intelligent Scan Planning

Uses AI to generate optimal scan strategies based on:
- Target profile analysis
- Attack surface assessment
- Risk prioritization
- Resource constraints
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Optional

from aipt_v2.strategy.target_analyzer import TargetProfile, TargetType, Technology

logger = logging.getLogger(__name__)


class ScanIntensity(str, Enum):
    """Scan intensity levels."""
    QUICK = "quick"          # Fast scan, basic checks
    STANDARD = "standard"    # Balanced scan
    THOROUGH = "thorough"    # Deep scan, all checks
    STEALTH = "stealth"      # Low-noise scan


class ScanPhaseType(str, Enum):
    """Types of scan phases."""
    RECON = "reconnaissance"
    SAST = "static_analysis"
    DAST = "dynamic_analysis"
    AUTH = "authentication"
    BUSINESS_LOGIC = "business_logic"
    API = "api_testing"
    WEBSOCKET = "websocket_testing"
    VALIDATION = "poc_validation"
    REPORTING = "reporting"


@dataclass
class ScanPhase:
    """A phase in the scan plan."""
    phase_type: ScanPhaseType
    name: str
    description: str
    agents: list[str]
    tools: list[str] = field(default_factory=list)
    timeout_minutes: int = 30
    priority: int = 1  # Lower = higher priority
    depends_on: list[str] = field(default_factory=list)
    enabled: bool = True
    config: dict = field(default_factory=dict)


@dataclass
class ScanPlan:
    """Complete scan plan generated by AI."""
    target: str
    profile: TargetProfile
    intensity: ScanIntensity
    phases: list[ScanPhase] = field(default_factory=list)
    total_estimated_minutes: int = 0
    created_at: datetime = field(default_factory=datetime.now)
    rationale: str = ""
    risk_assessment: dict = field(default_factory=dict)
    custom_config: dict = field(default_factory=dict)

    def get_phase_order(self) -> list[ScanPhase]:
        """Get phases in execution order based on dependencies."""
        # Simple topological sort
        executed = set()
        ordered = []

        while len(ordered) < len(self.phases):
            for phase in self.phases:
                if phase.name in executed:
                    continue
                if all(dep in executed for dep in phase.depends_on):
                    ordered.append(phase)
                    executed.add(phase.name)

        return ordered


class AIStrategyGenerator:
    """
    AI-powered scan strategy generator.

    Generates optimal scan plans based on:
    - Target profile (type, technologies, features)
    - Risk assessment
    - Available resources
    - Time constraints

    Usage:
        generator = AIStrategyGenerator()
        plan = await generator.generate(profile, intensity=ScanIntensity.STANDARD)
    """

    def __init__(self, llm_client: Optional[Any] = None):
        """
        Initialize strategy generator.

        Args:
            llm_client: Optional LLM client for AI-powered planning
        """
        self._llm = llm_client

    async def generate(
        self,
        profile: TargetProfile,
        intensity: ScanIntensity = ScanIntensity.STANDARD,
        constraints: Optional[dict] = None,
    ) -> ScanPlan:
        """
        Generate a scan plan for the target.

        Args:
            profile: Target profile from analyzer
            intensity: Scan intensity level
            constraints: Optional constraints (time, scope, etc.)

        Returns:
            ScanPlan with phases and configuration
        """
        constraints = constraints or {}

        plan = ScanPlan(
            target=profile.target,
            profile=profile,
            intensity=intensity,
        )

        # Generate risk assessment
        plan.risk_assessment = self._assess_risk(profile)

        # Generate phases based on target type
        plan.phases = self._generate_phases(profile, intensity, constraints)

        # Calculate estimated time
        plan.total_estimated_minutes = sum(p.timeout_minutes for p in plan.phases)

        # Generate rationale
        plan.rationale = self._generate_rationale(profile, plan)

        logger.info(
            f"[AI Planner] Generated {len(plan.phases)} phase plan "
            f"(~{plan.total_estimated_minutes} minutes)"
        )

        return plan

    def _assess_risk(self, profile: TargetProfile) -> dict:
        """Assess risk level of the target."""
        risk = {
            "overall": "medium",
            "factors": [],
            "score": 50,
        }

        # High-risk indicators
        if profile.has_auth:
            risk["factors"].append("Authentication present - credential attacks possible")
            risk["score"] += 10

        if profile.has_file_upload:
            risk["factors"].append("File upload - potential RCE vector")
            risk["score"] += 15

        if profile.has_api:
            risk["factors"].append("API endpoints - injection and auth bypass risks")
            risk["score"] += 10

        if profile.has_graphql:
            risk["factors"].append("GraphQL - introspection and complexity attacks")
            risk["score"] += 10

        # Technology-specific risks
        risky_techs = [
            Technology.WORDPRESS, Technology.DRUPAL, Technology.PHP, Technology.LARAVEL
        ]
        if any(t in profile.technologies for t in risky_techs):
            risk["factors"].append("Legacy/common framework with known vulnerabilities")
            risk["score"] += 10

        # Determine overall risk level
        if risk["score"] >= 70:
            risk["overall"] = "high"
        elif risk["score"] >= 40:
            risk["overall"] = "medium"
        else:
            risk["overall"] = "low"

        return risk

    def _generate_phases(
        self,
        profile: TargetProfile,
        intensity: ScanIntensity,
        constraints: dict,
    ) -> list[ScanPhase]:
        """Generate scan phases based on profile."""
        phases = []

        # Time multipliers based on intensity
        time_mult = {
            ScanIntensity.QUICK: 0.5,
            ScanIntensity.STANDARD: 1.0,
            ScanIntensity.THOROUGH: 2.0,
            ScanIntensity.STEALTH: 1.5,
        }[intensity]

        # Phase 1: Reconnaissance (for web targets)
        if profile.target_type not in [TargetType.LOCAL_DIRECTORY, TargetType.GITHUB_REPO]:
            phases.append(ScanPhase(
                phase_type=ScanPhaseType.RECON,
                name="reconnaissance",
                description="Discover attack surface: subdomains, directories, technologies",
                agents=["ReconAgent"],
                tools=["subfinder", "httpx", "nuclei-tech-detect"],
                timeout_minutes=int(15 * time_mult),
                priority=1,
            ))

        # Phase 2: SAST (for source code)
        if profile.source_path or profile.target_type in [
            TargetType.LOCAL_DIRECTORY, TargetType.GITHUB_REPO, TargetType.GITLAB_REPO
        ]:
            phases.append(ScanPhase(
                phase_type=ScanPhaseType.SAST,
                name="static_analysis",
                description="Analyze source code for vulnerabilities and secrets",
                agents=["SASTAgent"],
                tools=["semgrep", "bandit", "trufflehog"],
                timeout_minutes=int(20 * time_mult),
                priority=1,
                config={
                    "languages": list(profile.languages.keys()) if profile.languages else [],
                    "enable_secret_detection": True,
                },
            ))

        # Phase 3: DAST (for web targets)
        if profile.target_type in [
            TargetType.WEB_APP, TargetType.SPA, TargetType.REST_API
        ]:
            phases.append(ScanPhase(
                phase_type=ScanPhaseType.DAST,
                name="dynamic_analysis",
                description="Test for runtime vulnerabilities: XSS, SQLi, SSRF",
                agents=["DASTAgent"],
                tools=["nuclei", "dalfox", "sqlmap"],
                timeout_minutes=int(30 * time_mult),
                priority=2,
                depends_on=["reconnaissance"] if "reconnaissance" in [p.name for p in phases] else [],
                config={
                    "test_xss": True,
                    "test_sqli": True,
                    "test_ssrf": True,
                    "test_lfi": True,
                },
            ))

        # Phase 4: API Testing (for APIs)
        if profile.has_api or profile.target_type in [TargetType.REST_API, TargetType.GRAPHQL_API]:
            phases.append(ScanPhase(
                phase_type=ScanPhaseType.API,
                name="api_testing",
                description="Test API endpoints for security issues",
                agents=["DASTAgent"],
                tools=["nuclei-api", "arjun"],
                timeout_minutes=int(20 * time_mult),
                priority=2,
                depends_on=["reconnaissance"] if "reconnaissance" in [p.name for p in phases] else [],
                config={
                    "test_auth_bypass": True,
                    "test_mass_assignment": True,
                    "test_rate_limiting": True,
                },
            ))

        # Phase 5: GraphQL Testing
        if profile.has_graphql or profile.target_type == TargetType.GRAPHQL_API:
            phases.append(ScanPhase(
                phase_type=ScanPhaseType.API,
                name="graphql_testing",
                description="Test GraphQL for introspection, injection, DoS",
                agents=["DASTAgent"],
                tools=["graphql-voyager", "clairvoyance"],
                timeout_minutes=int(15 * time_mult),
                priority=2,
                config={
                    "test_introspection": True,
                    "test_injection": True,
                    "test_complexity_dos": True,
                },
            ))

        # Phase 6: WebSocket Testing
        if profile.has_websocket or profile.target_type == TargetType.WEBSOCKET:
            phases.append(ScanPhase(
                phase_type=ScanPhaseType.WEBSOCKET,
                name="websocket_testing",
                description="Test WebSocket connections for security issues",
                agents=["WebSocketAgent"],
                tools=["websocat"],
                timeout_minutes=int(15 * time_mult),
                priority=3,
                config={
                    "test_injection": True,
                    "test_cswsh": True,
                    "test_replay": True,
                },
            ))

        # Phase 7: Authentication Testing
        if profile.has_auth:
            phases.append(ScanPhase(
                phase_type=ScanPhaseType.AUTH,
                name="auth_testing",
                description="Test authentication and session management",
                agents=["DASTAgent", "BusinessLogicAgent"],
                tools=["hydra", "nuclei-auth"],
                timeout_minutes=int(20 * time_mult),
                priority=3,
                config={
                    "test_default_creds": True,
                    "test_session_management": True,
                    "test_password_policy": True,
                },
            ))

        # Phase 8: Business Logic Testing
        if profile.has_auth or profile.has_forms:
            phases.append(ScanPhase(
                phase_type=ScanPhaseType.BUSINESS_LOGIC,
                name="business_logic",
                description="Test for business logic flaws: IDOR, race conditions",
                agents=["BusinessLogicAgent"],
                tools=[],
                timeout_minutes=int(20 * time_mult),
                priority=4,
                depends_on=["auth_testing"] if "auth_testing" in [p.name for p in phases] else [],
                config={
                    "test_idor": True,
                    "test_race_conditions": True,
                    "test_workflow_bypass": True,
                },
            ))

        # Phase 9: PoC Validation (always last before reporting)
        phases.append(ScanPhase(
            phase_type=ScanPhaseType.VALIDATION,
            name="poc_validation",
            description="Validate findings with working exploits",
            agents=["PoCValidator"],
            tools=[],
            timeout_minutes=int(15 * time_mult),
            priority=8,
            depends_on=[p.name for p in phases if p.phase_type != ScanPhaseType.VALIDATION],
            config={
                "min_severity": "medium" if intensity == ScanIntensity.QUICK else "low",
                "generate_poc": True,
                "collect_evidence": True,
            },
        ))

        # Phase 10: Reporting (always last)
        phases.append(ScanPhase(
            phase_type=ScanPhaseType.REPORTING,
            name="reporting",
            description="Generate final report with findings and evidence",
            agents=[],
            tools=["sarif-generator"],
            timeout_minutes=5,
            priority=9,
            depends_on=["poc_validation"],
            config={
                "format": ["sarif", "json", "html"],
                "include_evidence": True,
            },
        ))

        return phases

    def _generate_rationale(self, profile: TargetProfile, plan: ScanPlan) -> str:
        """Generate human-readable rationale for the plan."""
        rationale_parts = []

        rationale_parts.append(
            f"Target type detected: {profile.target_type.value}. "
        )

        if profile.technologies:
            techs = ", ".join(t.value for t in profile.technologies[:5])
            rationale_parts.append(f"Technologies: {techs}. ")

        if plan.risk_assessment.get("overall") == "high":
            rationale_parts.append(
                "High-risk profile detected - comprehensive scanning recommended. "
            )

        phase_summary = ", ".join(p.name for p in plan.phases[:5])
        rationale_parts.append(f"Phases: {phase_summary}.")

        return "".join(rationale_parts)


# Convenience function
async def generate_scan_plan(
    profile: TargetProfile,
    intensity: ScanIntensity = ScanIntensity.STANDARD,
    **kwargs,
) -> ScanPlan:
    """
    Convenience function to generate a scan plan.

    Args:
        profile: Target profile
        intensity: Scan intensity
        **kwargs: Additional constraints

    Returns:
        ScanPlan
    """
    generator = AIStrategyGenerator()
    return await generator.generate(profile, intensity, constraints=kwargs)
