"""
Vulnerability Pattern Database - Persistent storage for discovered vulnerabilities.

This module:
1. Records vulnerability patterns discovered by red team analysis
2. Persists patterns across missions and sessions
3. Uses accumulated patterns to improve future red team prompts
4. Tracks which patterns are most common/severe
"""

import json
import hashlib
import sys
from pathlib import Path
from dataclasses import dataclass, field, asdict
from typing import List, Dict, Optional, Any
from datetime import datetime
from collections import defaultdict
from enum import Enum

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))
from atlasforge_config import AF_DATA_DIR


class VulnerabilityCategory(Enum):
    """Categories of vulnerabilities."""
    INJECTION = "injection"
    BOUNDARY = "boundary"
    TYPE_CONFUSION = "type_confusion"
    RESOURCE_EXHAUSTION = "resource_exhaustion"
    STATE_CORRUPTION = "state_corruption"
    CONCURRENCY = "concurrency"
    ERROR_HANDLING = "error_handling"
    LOGIC_FLAW = "logic_flaw"
    AUTHENTICATION = "authentication"
    AUTHORIZATION = "authorization"
    CRYPTO = "crypto"
    CONFIG = "config"
    OTHER = "other"


@dataclass
class VulnerabilityPattern:
    """A vulnerability pattern discovered across missions."""
    pattern_id: str
    category: VulnerabilityCategory
    name: str
    description: str
    code_pattern: str  # Regex or code pattern that identifies this vuln
    example_code: str  # Example vulnerable code
    fix_guidance: str  # How to fix this type of vuln
    severity_typical: str  # Typical severity: critical/high/medium/low
    occurrences: int = 1
    missions_found_in: List[str] = field(default_factory=list)
    first_seen: str = ""
    last_seen: str = ""
    tags: List[str] = field(default_factory=list)

    def __hash__(self):
        return hash(self.pattern_id)

    def to_dict(self) -> dict:
        result = asdict(self)
        result['category'] = self.category.value
        return result

    @staticmethod
    def from_dict(data: dict) -> 'VulnerabilityPattern':
        data['category'] = VulnerabilityCategory(data['category'])
        return VulnerabilityPattern(**data)


@dataclass
class MissionVulnerabilityRecord:
    """Record of vulnerabilities found in a specific mission."""
    mission_id: str
    timestamp: str
    code_path: str
    patterns_found: List[str]  # Pattern IDs
    findings_count: int
    critical_count: int
    high_count: int
    medium_count: int
    low_count: int

    def to_dict(self) -> dict:
        return asdict(self)


class VulnerabilityDatabase:
    """
    Persistent database of vulnerability patterns.

    Usage:
        db = VulnerabilityDatabase()

        # Record a finding
        db.record_finding(
            mission_id="mission_123",
            category=VulnerabilityCategory.INJECTION,
            name="SQL Injection via f-string",
            description="Using f-string interpolation in SQL queries",
            code_pattern=r'f".*SELECT.*{.*}.*"',
            example_code='query = f"SELECT * FROM users WHERE id = {user_id}"',
            fix_guidance="Use parameterized queries instead",
            severity="critical",
            code_path="path/to/file.py"
        )

        # Get patterns for improving red team prompts
        patterns = db.get_common_patterns(limit=10)

        # Generate enhanced prompts
        prompt_enhancement = db.generate_prompt_enhancement()
    """

    DEFAULT_DB_PATH = AF_DATA_DIR / "vulnerability_patterns.json"

    def __init__(self, db_path: Optional[Path] = None):
        """Initialize database, loading from disk if exists."""
        self.db_path = db_path or self.DEFAULT_DB_PATH
        self.patterns: Dict[str, VulnerabilityPattern] = {}
        self.mission_records: Dict[str, MissionVulnerabilityRecord] = {}
        self.stats = {
            "total_findings": 0,
            "total_missions": 0,
            "by_category": defaultdict(int),
            "by_severity": defaultdict(int)
        }
        self._load()

    def _load(self):
        """Load database from disk."""
        if self.db_path.exists():
            try:
                with open(self.db_path, 'r') as f:
                    data = json.load(f)

                self.patterns = {
                    pid: VulnerabilityPattern.from_dict(pdata)
                    for pid, pdata in data.get("patterns", {}).items()
                }

                self.mission_records = {
                    mid: MissionVulnerabilityRecord(**mdata)
                    for mid, mdata in data.get("mission_records", {}).items()
                }

                self.stats = data.get("stats", self.stats)
                self.stats["by_category"] = defaultdict(int, self.stats.get("by_category", {}))
                self.stats["by_severity"] = defaultdict(int, self.stats.get("by_severity", {}))
            except (json.JSONDecodeError, KeyError) as e:
                print(f"Warning: Could not load vulnerability database: {e}")

    def _save(self):
        """Save database to disk."""
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

        data = {
            "patterns": {
                pid: p.to_dict()
                for pid, p in self.patterns.items()
            },
            "mission_records": {
                mid: m.to_dict()
                for mid, m in self.mission_records.items()
            },
            "stats": {
                "total_findings": self.stats["total_findings"],
                "total_missions": self.stats["total_missions"],
                "by_category": dict(self.stats["by_category"]),
                "by_severity": dict(self.stats["by_severity"])
            },
            "last_updated": datetime.now().isoformat()
        }

        with open(self.db_path, 'w') as f:
            json.dump(data, f, indent=2)

    def _generate_pattern_id(
        self,
        category: VulnerabilityCategory,
        name: str,
        code_pattern: str
    ) -> str:
        """Generate unique pattern ID from components."""
        content = f"{category.value}:{name}:{code_pattern}"
        return hashlib.sha256(content.encode()).hexdigest()[:16]

    def record_finding(
        self,
        mission_id: str,
        category: VulnerabilityCategory,
        name: str,
        description: str,
        code_pattern: str,
        example_code: str,
        fix_guidance: str,
        severity: str,
        code_path: str,
        tags: Optional[List[str]] = None
    ) -> VulnerabilityPattern:
        """
        Record a vulnerability finding.

        If the pattern already exists, updates occurrence count.
        If new, creates a new pattern.

        Returns the pattern (new or existing).
        """
        pattern_id = self._generate_pattern_id(category, name, code_pattern)
        now = datetime.now().isoformat()

        if pattern_id in self.patterns:
            # Update existing pattern
            pattern = self.patterns[pattern_id]
            pattern.occurrences += 1
            pattern.last_seen = now
            if mission_id not in pattern.missions_found_in:
                pattern.missions_found_in.append(mission_id)
            if tags:
                pattern.tags = list(set(pattern.tags + tags))
        else:
            # Create new pattern
            pattern = VulnerabilityPattern(
                pattern_id=pattern_id,
                category=category,
                name=name,
                description=description,
                code_pattern=code_pattern,
                example_code=example_code,
                fix_guidance=fix_guidance,
                severity_typical=severity,
                occurrences=1,
                missions_found_in=[mission_id],
                first_seen=now,
                last_seen=now,
                tags=tags or []
            )
            self.patterns[pattern_id] = pattern

        # Update stats
        self.stats["total_findings"] += 1
        self.stats["by_category"][category.value] += 1
        self.stats["by_severity"][severity] += 1

        self._save()
        return pattern

    def record_mission(
        self,
        mission_id: str,
        code_path: str,
        patterns_found: List[str],
        severity_counts: Dict[str, int]
    ):
        """Record a mission's vulnerability scan results."""
        record = MissionVulnerabilityRecord(
            mission_id=mission_id,
            timestamp=datetime.now().isoformat(),
            code_path=code_path,
            patterns_found=patterns_found,
            findings_count=sum(severity_counts.values()),
            critical_count=severity_counts.get("critical", 0),
            high_count=severity_counts.get("high", 0),
            medium_count=severity_counts.get("medium", 0),
            low_count=severity_counts.get("low", 0)
        )
        self.mission_records[mission_id] = record

        if mission_id not in [m.mission_id for m in self.mission_records.values()]:
            self.stats["total_missions"] += 1

        self._save()

    def get_pattern(self, pattern_id: str) -> Optional[VulnerabilityPattern]:
        """Get a specific pattern by ID."""
        return self.patterns.get(pattern_id)

    def get_patterns_by_category(
        self,
        category: VulnerabilityCategory
    ) -> List[VulnerabilityPattern]:
        """Get all patterns in a category."""
        return [p for p in self.patterns.values() if p.category == category]

    def get_common_patterns(self, limit: int = 10) -> List[VulnerabilityPattern]:
        """Get most common patterns sorted by occurrence count."""
        return sorted(
            self.patterns.values(),
            key=lambda p: p.occurrences,
            reverse=True
        )[:limit]

    def get_severe_patterns(self, limit: int = 10) -> List[VulnerabilityPattern]:
        """Get patterns sorted by severity (critical first)."""
        severity_order = {"critical": 4, "high": 3, "medium": 2, "low": 1}
        return sorted(
            self.patterns.values(),
            key=lambda p: severity_order.get(p.severity_typical, 0),
            reverse=True
        )[:limit]

    def get_recent_patterns(self, limit: int = 10) -> List[VulnerabilityPattern]:
        """Get most recently seen patterns."""
        return sorted(
            self.patterns.values(),
            key=lambda p: p.last_seen,
            reverse=True
        )[:limit]

    def generate_prompt_enhancement(self, max_patterns: int = 5) -> str:
        """
        Generate an enhancement for red team prompts based on common patterns.

        This adds context about commonly found vulnerabilities to help
        the red team agent focus on real-world issues.
        """
        common = self.get_common_patterns(max_patterns)

        if not common:
            return ""

        lines = [
            "\n## Common Vulnerability Patterns (from historical data)",
            "Based on previous analysis, pay special attention to:",
            ""
        ]

        for i, pattern in enumerate(common, 1):
            lines.append(f"{i}. **{pattern.name}** ({pattern.category.value})")
            lines.append(f"   - Found {pattern.occurrences} times across {len(pattern.missions_found_in)} missions")
            lines.append(f"   - Look for: {pattern.code_pattern}")
            lines.append(f"   - Example: `{pattern.example_code[:80]}...`" if len(pattern.example_code) > 80 else f"   - Example: `{pattern.example_code}`")
            lines.append("")

        return "\n".join(lines)

    def get_statistics(self) -> Dict[str, Any]:
        """Get database statistics."""
        return {
            "total_patterns": len(self.patterns),
            "total_findings": self.stats["total_findings"],
            "total_missions": self.stats["total_missions"],
            "by_category": dict(self.stats["by_category"]),
            "by_severity": dict(self.stats["by_severity"]),
            "most_common": [
                {"name": p.name, "count": p.occurrences}
                for p in self.get_common_patterns(5)
            ]
        }

    def export_patterns(self, filepath: Path) -> None:
        """Export all patterns to a file."""
        with open(filepath, 'w') as f:
            json.dump({
                "patterns": [p.to_dict() for p in self.patterns.values()],
                "stats": self.get_statistics(),
                "exported_at": datetime.now().isoformat()
            }, f, indent=2)

    def clear(self):
        """Clear all data (use with caution!)."""
        self.patterns = {}
        self.mission_records = {}
        self.stats = {
            "total_findings": 0,
            "total_missions": 0,
            "by_category": defaultdict(int),
            "by_severity": defaultdict(int)
        }
        self._save()


def record_red_team_findings(
    db: VulnerabilityDatabase,
    mission_id: str,
    code_path: str,
    findings: List[Dict[str, Any]]
):
    """
    Helper function to record findings from a red team result.

    Args:
        db: VulnerabilityDatabase instance
        mission_id: Current mission ID
        code_path: Path to code being analyzed
        findings: List of finding dicts from RedTeamResult
    """
    patterns_found = []
    severity_counts = defaultdict(int)

    for finding in findings:
        # Map red team categories to vulnerability categories
        category_map = {
            "injection": VulnerabilityCategory.INJECTION,
            "boundary": VulnerabilityCategory.BOUNDARY,
            "type_confusion": VulnerabilityCategory.TYPE_CONFUSION,
            "resource": VulnerabilityCategory.RESOURCE_EXHAUSTION,
            "state_corruption": VulnerabilityCategory.STATE_CORRUPTION,
            "concurrency": VulnerabilityCategory.CONCURRENCY,
            "error_handling": VulnerabilityCategory.ERROR_HANDLING,
            "logic": VulnerabilityCategory.LOGIC_FLAW,
        }

        category = category_map.get(
            finding.get("category", "other"),
            VulnerabilityCategory.OTHER
        )

        pattern = db.record_finding(
            mission_id=mission_id,
            category=category,
            name=finding.get("title", "Unknown"),
            description=finding.get("description", ""),
            code_pattern=finding.get("affected_code", ""),
            example_code=finding.get("reproduction_steps", [""])[0] if finding.get("reproduction_steps") else "",
            fix_guidance=finding.get("suggested_fix", ""),
            severity=finding.get("severity", "medium"),
            code_path=code_path,
            tags=[category.value]
        )
        patterns_found.append(pattern.pattern_id)
        severity_counts[finding.get("severity", "medium")] += 1

    # Record mission summary
    db.record_mission(
        mission_id=mission_id,
        code_path=code_path,
        patterns_found=patterns_found,
        severity_counts=dict(severity_counts)
    )


if __name__ == "__main__":
    # Self-test
    print("Vulnerability Database - Self Test")
    print("=" * 50)

    # Use temp location for test
    test_db_path = Path("/tmp/test_vuln_db.json")
    db = VulnerabilityDatabase(db_path=test_db_path)

    # Clear any existing test data
    db.clear()

    # Record some test findings
    print("\n1. Recording test findings...")
    db.record_finding(
        mission_id="test_mission_1",
        category=VulnerabilityCategory.INJECTION,
        name="SQL Injection via f-string",
        description="Using f-string interpolation in SQL queries allows injection",
        code_pattern=r'f".*SELECT.*{.*}.*"',
        example_code='query = f"SELECT * FROM users WHERE id = {user_id}"',
        fix_guidance="Use parameterized queries with ? placeholders",
        severity="critical",
        code_path="test/file.py"
    )

    db.record_finding(
        mission_id="test_mission_1",
        category=VulnerabilityCategory.BOUNDARY,
        name="Missing array bounds check",
        description="Accessing array without checking length",
        code_pattern=r'\[.*\](?!.*if.*len)',
        example_code='return items[index]',
        fix_guidance="Check array length before accessing",
        severity="high",
        code_path="test/file.py"
    )

    # Record same pattern again (different mission)
    db.record_finding(
        mission_id="test_mission_2",
        category=VulnerabilityCategory.INJECTION,
        name="SQL Injection via f-string",
        description="Using f-string interpolation in SQL queries allows injection",
        code_pattern=r'f".*SELECT.*{.*}.*"',
        example_code='query = f"SELECT * FROM users WHERE id = {user_id}"',
        fix_guidance="Use parameterized queries with ? placeholders",
        severity="critical",
        code_path="test/other_file.py"
    )

    print("\n2. Database Statistics:")
    stats = db.get_statistics()
    print(f"   Total patterns: {stats['total_patterns']}")
    print(f"   Total findings: {stats['total_findings']}")
    print(f"   By category: {stats['by_category']}")
    print(f"   By severity: {stats['by_severity']}")

    print("\n3. Common Patterns:")
    for pattern in db.get_common_patterns(5):
        print(f"   - {pattern.name}: {pattern.occurrences} occurrences")

    print("\n4. Prompt Enhancement:")
    enhancement = db.generate_prompt_enhancement()
    print(enhancement[:500] + "..." if len(enhancement) > 500 else enhancement)

    print("\n5. Persistence Test:")
    # Reload from disk
    db2 = VulnerabilityDatabase(db_path=test_db_path)
    print(f"   Loaded {len(db2.patterns)} patterns from disk")
    assert len(db2.patterns) == 2, "Persistence failed!"

    # Cleanup
    test_db_path.unlink(missing_ok=True)

    print("\nVulnerability database self-test complete!")
