"""AI-powered vulnerability analysis"""

from dataclasses import dataclass, asdict
from typing import List, Dict, Any, Optional, Callable
import json
import logging

from threat_radar.ai.llm_client import LLMClient, get_llm_client
from threat_radar.ai.prompt_templates import (
    create_analysis_prompt,
    create_batch_analysis_prompt,
    create_summary_consolidation_prompt,
)
from threat_radar.core.grype_integration import GrypeScanResult, GrypeVulnerability

logger = logging.getLogger(__name__)


@dataclass
class VulnerabilityInsight:
    """AI-generated insight for a single vulnerability"""

    cve_id: str
    package_name: str
    exploitability: str  # HIGH, MEDIUM, LOW
    exploitability_details: str
    attack_vectors: List[str]
    business_impact: str  # HIGH, MEDIUM, LOW
    business_impact_details: str
    recommendations: List[str]


@dataclass
class VulnerabilityAnalysis:
    """Complete AI analysis of vulnerability scan results"""

    vulnerabilities: List[VulnerabilityInsight]
    summary: str
    metadata: Dict[str, Any]

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            "vulnerabilities": [asdict(v) for v in self.vulnerabilities],
            "summary": self.summary,
            "metadata": self.metadata,
        }


class VulnerabilityAnalyzer:
    """Analyzes vulnerabilities using AI to provide context and insights"""

    def __init__(
        self,
        llm_client: Optional[LLMClient] = None,
        provider: Optional[str] = None,
        model: Optional[str] = None,
        batch_size: int = 25,
        auto_batch_threshold: int = 30,
    ):
        """
        Initialize vulnerability analyzer.

        Args:
            llm_client: Pre-configured LLM client (optional)
            provider: AI provider if not using pre-configured client
            model: Model name if not using pre-configured client
            batch_size: Number of vulnerabilities per batch (default: 25)
            auto_batch_threshold: Trigger batching when count exceeds this (default: 30)
        """
        self.llm_client = llm_client or get_llm_client(provider=provider, model=model)
        self.batch_size = batch_size
        self.auto_batch_threshold = auto_batch_threshold

    def analyze_scan_result(
        self,
        scan_result: GrypeScanResult,
        temperature: float = 0.3,
        batch_mode: Optional[str] = "auto",
        progress_callback: Optional[Callable[[int, int, int], None]] = None,
    ) -> VulnerabilityAnalysis:
        """
        Analyze a complete Grype scan result with smart batching.

        Args:
            scan_result: GrypeScanResult from Grype scan
            temperature: LLM temperature for analysis (lower = more consistent)
            batch_mode: Batching mode - "auto", "enabled", or "disabled"
            progress_callback: Optional callback(batch_num, total_batches, analyzed_count)

        Returns:
            VulnerabilityAnalysis with AI-generated insights
        """
        vuln_count = len(scan_result.vulnerabilities)

        # Determine whether to use batch processing
        use_batching = False
        if batch_mode == "enabled":
            use_batching = True
        elif batch_mode == "auto" and vuln_count > self.auto_batch_threshold:
            use_batching = True
            logger.info(
                f"Large scan detected ({vuln_count} CVEs > {self.auto_batch_threshold} threshold). "
                f"Using batch processing..."
            )
        elif batch_mode == "disabled":
            use_batching = False

        # Route to appropriate analysis method
        if use_batching:
            return self._analyze_with_batches(
                scan_result, temperature, progress_callback
            )
        else:
            return self._analyze_standard(scan_result, temperature)

    def _analyze_standard(
        self, scan_result: GrypeScanResult, temperature: float = 0.3
    ) -> VulnerabilityAnalysis:
        """
        Standard single-pass analysis (original behavior).

        Args:
            scan_result: GrypeScanResult from Grype scan
            temperature: LLM temperature for analysis

        Returns:
            VulnerabilityAnalysis with AI-generated insights
        """
        # Convert Grype vulnerabilities to dict format
        vulnerabilities_data = [
            self._vuln_to_dict(v) for v in scan_result.vulnerabilities
        ]

        # Generate analysis using AI
        prompt = create_analysis_prompt(vulnerabilities_data)

        try:
            response = self.llm_client.generate_json(prompt, temperature=temperature)

            # Parse response into structured format
            insights = [
                VulnerabilityInsight(**vuln_data)
                for vuln_data in response.get("vulnerabilities", [])
            ]

            metadata = {
                "target": scan_result.target,
                "total_vulnerabilities": scan_result.total_count,
                "severity_counts": scan_result.severity_counts,
                "scan_metadata": scan_result.scan_metadata,
            }

            return VulnerabilityAnalysis(
                vulnerabilities=insights,
                summary=response.get("summary", ""),
                metadata=metadata,
            )

        except Exception as e:
            raise RuntimeError(f"Failed to analyze vulnerabilities: {str(e)}")

    def _analyze_with_batches(
        self,
        scan_result: GrypeScanResult,
        temperature: float = 0.3,
        progress_callback: Optional[Callable[[int, int, int], None]] = None,
    ) -> VulnerabilityAnalysis:
        """
        Analyze vulnerabilities in batches for large scans.

        Args:
            scan_result: GrypeScanResult from Grype scan
            temperature: LLM temperature for analysis
            progress_callback: Optional callback(batch_num, total_batches, analyzed_count)

        Returns:
            VulnerabilityAnalysis with consolidated results
        """
        all_insights = []
        batch_summaries = []
        total_vulns = len(scan_result.vulnerabilities)
        total_batches = (total_vulns + self.batch_size - 1) // self.batch_size

        logger.info(
            f"Processing {total_vulns} vulnerabilities in {total_batches} batches of {self.batch_size}"
        )

        # Process each batch
        for batch_idx in range(total_batches):
            start_idx = batch_idx * self.batch_size
            end_idx = min(start_idx + self.batch_size, total_vulns)
            batch = scan_result.vulnerabilities[start_idx:end_idx]
            batch_num = batch_idx + 1

            logger.debug(
                f"Processing batch {batch_num}/{total_batches} ({len(batch)} vulnerabilities)"
            )

            # Convert to dict format
            batch_data = [self._vuln_to_dict(v) for v in batch]

            # Create batch-specific prompt
            prompt = create_batch_analysis_prompt(batch_data, batch_num, total_batches)

            try:
                response = self.llm_client.generate_json(
                    prompt, temperature=temperature
                )

                # Parse batch insights
                batch_insights = [
                    VulnerabilityInsight(**vuln_data)
                    for vuln_data in response.get("vulnerabilities", [])
                ]
                all_insights.extend(batch_insights)

                # Store batch summary
                batch_summary = response.get("summary", f"Batch {batch_num} analyzed")
                batch_summaries.append(batch_summary)

                logger.debug(
                    f"Batch {batch_num} complete: {len(batch_insights)} insights generated"
                )

                # Call progress callback if provided
                if progress_callback:
                    progress_callback(batch_num, total_batches, len(all_insights))

            except Exception as e:
                logger.warning(
                    f"Batch {batch_num} failed: {str(e)}. Continuing with remaining batches..."
                )
                batch_summaries.append(f"Batch {batch_num}: Analysis failed")
                continue

        # Generate consolidated summary
        high_priority_count = sum(
            1
            for insight in all_insights
            if insight.exploitability == "HIGH" or insight.business_impact == "HIGH"
        )

        try:
            consolidation_prompt = create_summary_consolidation_prompt(
                target=scan_result.target,
                total_vulnerabilities=total_vulns,
                severity_counts=scan_result.severity_counts,
                batch_summaries=batch_summaries,
                high_priority_count=high_priority_count,
            )
            consolidated_summary = self.llm_client.generate(
                consolidation_prompt, temperature=temperature
            )
        except Exception as e:
            logger.warning(f"Failed to generate consolidated summary: {e}")
            consolidated_summary = (
                f"Analyzed {len(all_insights)} of {total_vulns} vulnerabilities across {total_batches} batches. "
                f"{high_priority_count} high-priority vulnerabilities identified."
            )

        # Build metadata with batch information
        metadata = {
            "target": scan_result.target,
            "total_vulnerabilities": total_vulns,
            "severity_counts": scan_result.severity_counts,
            "scan_metadata": scan_result.scan_metadata,
            "batch_processing": True,
            "batches_processed": total_batches,
            "batch_size": self.batch_size,
            "insights_generated": len(all_insights),
        }

        return VulnerabilityAnalysis(
            vulnerabilities=all_insights,
            summary=consolidated_summary,
            metadata=metadata,
        )

    def analyze_vulnerabilities(
        self, vulnerabilities: List[GrypeVulnerability], temperature: float = 0.3
    ) -> VulnerabilityAnalysis:
        """
        Analyze a list of vulnerabilities.

        Args:
            vulnerabilities: List of GrypeVulnerability objects
            temperature: LLM temperature for analysis

        Returns:
            VulnerabilityAnalysis with AI-generated insights
        """
        vulnerabilities_data = [self._vuln_to_dict(v) for v in vulnerabilities]

        prompt = create_analysis_prompt(vulnerabilities_data)

        try:
            response = self.llm_client.generate_json(prompt, temperature=temperature)

            insights = [
                VulnerabilityInsight(**vuln_data)
                for vuln_data in response.get("vulnerabilities", [])
            ]

            # Calculate severity counts
            severity_counts = {}
            for vuln in vulnerabilities:
                severity = vuln.severity.lower()
                severity_counts[severity] = severity_counts.get(severity, 0) + 1

            metadata = {
                "total_vulnerabilities": len(vulnerabilities),
                "severity_counts": severity_counts,
            }

            return VulnerabilityAnalysis(
                vulnerabilities=insights,
                summary=response.get("summary", ""),
                metadata=metadata,
            )

        except Exception as e:
            raise RuntimeError(f"Failed to analyze vulnerabilities: {str(e)}")

    def _vuln_to_dict(self, vuln: GrypeVulnerability) -> Dict[str, Any]:
        """Convert GrypeVulnerability to dictionary"""
        return {
            "id": vuln.id,
            "severity": vuln.severity,
            "package_name": vuln.package_name,
            "package_version": vuln.package_version,
            "package_type": vuln.package_type,
            "fixed_in_version": vuln.fixed_in_version,
            "description": vuln.description,
            "cvss_score": vuln.cvss_score,
            "urls": vuln.urls,
            "data_source": vuln.data_source,
        }

    def filter_by_severity(
        self,
        scan_result: GrypeScanResult,
        min_severity: str,
    ) -> GrypeScanResult:
        """
        Filter scan results by minimum severity level.

        Args:
            scan_result: Original scan result
            min_severity: Minimum severity (critical, high, medium, low, negligible)

        Returns:
            Filtered GrypeScanResult with only vulnerabilities >= min_severity
        """
        severity_order = {
            "critical": 4,
            "high": 3,
            "medium": 2,
            "low": 1,
            "negligible": 0,
        }

        min_severity_lower = min_severity.lower()
        if min_severity_lower not in severity_order:
            raise ValueError(
                f"Invalid severity: {min_severity}. "
                f"Valid values: critical, high, medium, low, negligible"
            )

        threshold = severity_order[min_severity_lower]

        # Filter vulnerabilities
        filtered_vulns = [
            v
            for v in scan_result.vulnerabilities
            if severity_order.get(v.severity.lower(), 0) >= threshold
        ]

        # Recalculate severity counts
        filtered_severity_counts = {}
        for vuln in filtered_vulns:
            severity = vuln.severity.lower()
            filtered_severity_counts[severity] = (
                filtered_severity_counts.get(severity, 0) + 1
            )

        logger.info(
            f"Severity filter: {len(filtered_vulns)}/{scan_result.total_count} "
            f"vulnerabilities >= {min_severity.upper()}"
        )

        return GrypeScanResult(
            target=scan_result.target,
            vulnerabilities=filtered_vulns,
            total_count=len(filtered_vulns),
            severity_counts=filtered_severity_counts,
            scan_metadata={
                **(scan_result.scan_metadata or {}),
                "severity_filter": min_severity_lower,
                "original_count": scan_result.total_count,
            },
        )

    def get_high_priority_vulnerabilities(
        self, analysis: VulnerabilityAnalysis
    ) -> List[VulnerabilityInsight]:
        """
        Filter vulnerabilities to high priority ones.

        Args:
            analysis: VulnerabilityAnalysis result

        Returns:
            List of high priority vulnerabilities
        """
        return [
            v
            for v in analysis.vulnerabilities
            if v.exploitability == "HIGH" or v.business_impact == "HIGH"
        ]

    def get_critical_paths(self, analysis: VulnerabilityAnalysis) -> List[str]:
        """
        Extract critical attack paths from analysis.

        Args:
            analysis: VulnerabilityAnalysis result

        Returns:
            List of critical attack vector descriptions
        """
        critical_vectors = set()

        for vuln in analysis.vulnerabilities:
            if vuln.exploitability == "HIGH":
                critical_vectors.update(vuln.attack_vectors)

        return list(critical_vectors)
