"""
CVE Data Scraper for Linux Kernel Security Vulnerabilities

This script:
- Loads CVE IDs from JSON files (train/test_kernel_cves.json)
- Implements 2-strategy approach for commit extraction:
  * Strategy 1: Direct patch ID (if available from OSIDB)
  * Strategy 2: JSON references with fix-commit filtering (if no patch_id)
- Skips CVEs that don't meet either strategy
- Uses robust git commit fetching with multi-branch support
- Saves structured data including commit patches, metadata, and file changes

Data Sources:
- train/test_kernel_cves.json (CVE lists with patch IDs)
- https://git.kernel.org/pub/scm/linux/security/vulns.git (security repo for JSON references)
- git://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git (main kernel repo)

Output: Organized directory structure with CVE data, commits, patches, and analysis reports
"""

import re
import subprocess
import time
from pathlib import Path
import logging
from typing import List, Dict, Optional
import json


# Change to 5 for processing only 5 CVEs
CVE_CNT_LIMIT = None


# Setup logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class CVEDataScraper:
    def __init__(self, base_data_dir: str = "data"):
        self.base_data_dir = Path(base_data_dir)
        self.base_data_dir.mkdir(exist_ok=True)
        self._fetched_refs = set()  # Track fetched refs to avoid duplicates

        # Git repository setup
        self.linux_repo_path = self.base_data_dir / "linux_kernel_repo"
        self.vulns_repo_path = self.base_data_dir / "linux_security_vulns"

        # Strategy statistics
        self.strategy_stats = {
            "direct_patch": 0,
            "json_references_filtered": 0,
            "skipped_no_strategy": 0,
        }

        self.cve_list = self._load_cve_data()

        # Ensure stable repository remote is available for stable tree commits
        # (Only if not in single CVE mode - predictor will handle it separately)

    def _ensure_stable_remote_available(self):
        """Automatically add and fetch stable remote if needed for CVE commits in stable tree"""
        try:
            if not self.linux_repo_path.exists():
                logger.warning(f"Linux kernel repo not found at {self.linux_repo_path}")
                return

            # Check if stable remote already exists
            result = subprocess.run(
                ["git", "remote", "get-url", "stable"],
                cwd=self.linux_repo_path,
                capture_output=True,
                text=True,
            )

            if result.returncode == 0:
                logger.info("Stable remote already configured")
                return

            logger.info("Adding stable kernel tree remote...")

            # Add stable remote
            add_result = subprocess.run(
                [
                    "git",
                    "remote",
                    "add",
                    "stable",
                    "git://git.kernel.org/pub/scm/linux/kernel/git/stable/linux.git",
                ],
                cwd=self.linux_repo_path,
                capture_output=True,
            )

            if add_result.returncode != 0:
                logger.warning(f"Failed to add stable remote: {add_result.stderr}")
                return

            logger.info(
                "Fetching stable kernel tree commits (this may take a while)..."
            )

            # Fetch stable commits
            fetch_result = subprocess.run(
                ["git", "fetch", "stable"],
                cwd=self.linux_repo_path,
                capture_output=True,
            )

            if fetch_result.returncode == 0:
                logger.info("Successfully fetched stable tree commits")
            else:
                logger.warning(f"Failed to fetch stable commits: {fetch_result.stderr}")

        except Exception as e:
            logger.warning(f"Error setting up stable remote: {e}")

    def _load_cve_data(self) -> List[Dict]:
        """Load CVE data from true labeled JSON files"""
        all_cves = []

        # Load training CVEs
        train_file = self.base_data_dir / "train_kernel_cves.json"
        if train_file.exists():
            with open(train_file, "r", encoding="utf-8") as f:
                train_data = json.load(f)
            all_cves.extend(train_data)
            logger.info(f"Loaded {len(train_data)} training CVEs")

        # Load test CVEs
        test_file = self.base_data_dir / "test_kernel_cves.json"
        if test_file.exists():
            with open(test_file, "r", encoding="utf-8") as f:
                test_data = json.load(f)
            all_cves.extend(test_data)
            logger.info(f"Loaded {len(test_data)} test CVEs")

        logger.info(f"Total loaded: {len(all_cves)} CVEs")
        return all_cves

    def setup_repositories(self):
        """Clone or update both Linux kernel and security vulnerabilities repositories"""
        logger.info("Setting up repositories...")

        # Setup security vulnerabilities repo
        self.setup_vulns_repo()

        # Setup Linux kernel repository
        self.setup_linux_repo()

        logger.info("All repositories setup complete")

    def setup_linux_repo(self):
        """Clone or update the Linux kernel git repository"""
        logger.info("Setting up Linux kernel git repository...")

        if not self.linux_repo_path.exists():
            logger.info("Cloning Linux kernel repository (this will take a while)...")
            cmd = [
                "git",
                "clone",
                "git://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git",
                str(self.linux_repo_path),
            ]
            try:
                subprocess.run(cmd, check=True)
            except subprocess.CalledProcessError:
                logger.info("Git protocol failed, trying HTTPS...")
                cmd = [
                    "git",
                    "clone",
                    "https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git",
                    str(self.linux_repo_path),
                ]
                subprocess.run(cmd, check=True)
        else:
            logger.info("Linux kernel repository exists. Updating with git pull...")
            try:
                subprocess.run(["git", "pull"], cwd=self.linux_repo_path, check=True)
                logger.info("Linux kernel repository updated successfully")
            except subprocess.CalledProcessError as e:
                logger.warning(
                    f"Git pull failed for Linux repo: {e}. Continuing with existing repo..."
                )

    def setup_vulns_repo(self):
        """Clone or update the Linux security vulnerabilities repo"""
        logger.info("Setting up Linux security vulnerabilities repository...")

        if not self.vulns_repo_path.exists():
            logger.info("Cloning Linux security vulnerabilities repo...")
            cmd = [
                "git",
                "clone",
                "https://git.kernel.org/pub/scm/linux/security/vulns.git",
                str(self.vulns_repo_path),
            ]
            subprocess.run(cmd, check=True)
        else:
            logger.info(
                "Security vulnerabilities repo exists. Updating with git pull..."
            )
            try:
                subprocess.run(["git", "pull"], cwd=self.vulns_repo_path, check=True)
                logger.info("Security vulnerabilities repository updated successfully")
            except subprocess.CalledProcessError as e:
                logger.warning(
                    f"Git pull failed for vulns repo: {e}. Continuing with existing repo..."
                )

    def get_extraction_strategy(self, cve_data: Dict) -> str:
        """Determine extraction strategy for a CVE"""
        cve_id = cve_data["cve_id"]
        patch_id = cve_data.get("patch_id", "")

        if patch_id and patch_id.strip():  # Has non-empty patch_id
            return "direct_patch"

        # Check if JSON file exists for references
        # FIXED: Use CVE year from CVE ID, not classification
        cve_year = cve_id.split("-")[1]  # Extract "2021" from "CVE-2021-46939"
        json_paths = [
            self.vulns_repo_path / f"cve/published/{cve_year}/{cve_id}.json",
            self.vulns_repo_path / f"cve/{cve_year}/{cve_id}.json",
            self.vulns_repo_path / f"cve/rejected/{cve_year}/{cve_id}.json",
        ]

        for json_path in json_paths:
            if json_path.exists():
                return "json_references_filtered"
        return "skipped_no_strategy"

    def classify_commit_context(self, text: str) -> str:
        """Classify if commit text indicates a fix or introduce commit"""
        if not text:
            return "unknown"

        text_lower = text.lower()

        # Fix indicators (strong signals)
        fix_patterns = [
            r"\bfix\w*\b",
            r"\bpatch\w*\b",
            r"\bresolv\w*\b",
            r"\baddress\w*\b",
            r"\bcorrect\w*\b",
            r"\brepair\w*\b",
            r"\bmend\w*\b",
        ]

        # Introduce indicators (avoid these)
        introduce_patterns = [
            r"\bintroduc\w*\b",
            r"\badd\w*\b",
            r"\bimplement\w*\b",
            r"\bcaus\w*\b",
        ]

        fix_score = sum(
            len(re.findall(pattern, text_lower)) for pattern in fix_patterns
        )
        introduce_score = sum(
            len(re.findall(pattern, text_lower)) for pattern in introduce_patterns
        )

        if fix_score > introduce_score and fix_score > 0:
            return "fix"
        elif introduce_score > fix_score and introduce_score > 0:
            return "introduce"
        else:
            return "unknown"

    def extract_commits_from_json_references_filtered(self, cve_id: str) -> List[str]:
        """Extract and filter fix commits from JSON references"""

        # FIXED: Use CVE year from CVE ID (e.g., "2021" from "CVE-2021-46939")
        cve_year = cve_id.split("-")[1]
        json_paths = [
            self.vulns_repo_path / f"cve/published/{cve_year}/{cve_id}.json",
            self.vulns_repo_path / f"cve/{cve_year}/{cve_id}.json",
            self.vulns_repo_path / f"cve/rejected/{cve_year}/{cve_id}.json",
        ]

        commits = []

        for json_path in json_paths:
            if not json_path.exists():
                continue
            try:
                with open(json_path, "r", encoding="utf-8") as f:
                    json_data = json.load(f)

                # FIXED: Extract commit hashes from references - handle nested JSON structure
                references = []
                if (
                    "containers" in json_data
                    and "cna" in json_data["containers"]
                    and "references" in json_data["containers"]["cna"]
                ):
                    references = json_data["containers"]["cna"]["references"]
                elif "references" in json_data:
                    references = json_data["references"]
                else:
                    references = []

                if references:
                    for ref in references:
                        if isinstance(ref, dict):
                            # Check URL field
                            url = ref.get("url", "")
                            commit_hash = self.extract_commit_hash_from_url(url)
                            if commit_hash:
                                # For now, trust JSON references (like test1 approach)
                                commits.append(commit_hash)

                # Also check other fields that might contain commit references
                for field in ["description", "problemDescription", "mitigation"]:
                    if field in json_data:
                        field_value = json_data[field]
                        if isinstance(field_value, str):
                            commit_hash = self.extract_commit_hash_from_text(
                                field_value
                            )
                            if commit_hash:
                                commits.append(commit_hash)

            except Exception as e:
                logger.error(f"Error parsing JSON for {cve_id}: {e}")
                continue

        # Remove duplicates while preserving order
        unique_commits = list(dict.fromkeys(commits))
        return unique_commits

    def extract_commit_hash_from_url(self, url: str) -> Optional[str]:
        """Extract commit hash from various URL formats"""
        if not url:
            return None

        # Use comprehensive commit hash pattern like test1 approach
        if "git.kernel.org" in url or "github.com" in url:
            commit_pattern = r"[a-f0-9]{40}"  # Full 40-character commit hash
            matches = re.findall(commit_pattern, url)
            if matches:
                return matches[0]  # Return first match

        return None

    def extract_commit_hash_from_text(self, text: str) -> Optional[str]:
        """Extract commit hash from free text"""
        if not text:
            return None

        # Look for 40-character hex strings
        commit_pattern = r"\b[0-9a-fA-F]{40}\b"
        matches = re.findall(commit_pattern, text)
        return matches[0] if matches else None

    def extract_commits_with_strategy(self, cve_data: Dict) -> List[str]:
        """Extract commits using the appropriate strategy"""
        strategy = self.get_extraction_strategy(cve_data)
        cve_id = cve_data["cve_id"]

        if strategy == "direct_patch":
            # Strategy 1: Use patch ID directly
            patch_id = cve_data["patch_id"]
            self.strategy_stats["direct_patch"] += 1
            return [patch_id]

        elif strategy == "json_references_filtered":
            # Strategy 2: Parse JSON references with filtering
            commits = self.extract_commits_from_json_references_filtered(cve_id)
            self.strategy_stats["json_references_filtered"] += 1
            return commits

        else:
            # No strategy available - skip
            logger.warning(f"No extraction strategy available for {cve_id} - skipping")
            self.strategy_stats["skipped_no_strategy"] += 1
            return []

    def process_cve_with_strategy(self, cve_data: Dict) -> Optional[Dict]:
        """Process CVE information using multi-strategy approach"""
        cve_id = cve_data["cve_id"]
        strategy = self.get_extraction_strategy(cve_data)

        logger.info(
            f"Processing {cve_id} with multi-strategy approach - Strategy: {strategy}"
        )

        cve_dir = self.base_data_dir / cve_id
        cve_dir.mkdir(exist_ok=True)

        # Extract commits using strategy
        commit_urls = []
        commits = self.extract_commits_with_strategy(cve_data)

        if not commits:
            logger.warning(
                f"No commits found for {cve_id} using strategy '{strategy}' - saving metadata anyway"
            )
            # FIXED: Don't return None, continue to save metadata
        else:
            # Convert commit hashes to URLs for compatibility with existing fetch logic
            for commit_hash in commits:
                commit_url = f"https://git.kernel.org/stable/c/{commit_hash}"
                commit_urls.append(commit_url)

            logger.info(f"Found {len(commit_urls)} commit URLs for {cve_id}")

            # Use existing robust commit fetching logic
            if commit_urls:
                self.fetch_commit_info_from_linux_repo(cve_id, commit_urls)

        # FIXED: Always save metadata regardless of whether commits were found
        metadata = {
            "cve_id": cve_id,
            "source_files": [],
            "json_data": None,
            "commit_urls": commit_urls,
            "affected_files": [],
            "scraped_at": time.time(),
            "method": "multi_strategy",
            "strategy": strategy,
            "severity": cve_data.get("severity"),
            "created_date": cve_data.get("created_date"),
            "classification": cve_data.get("classification"),
            "cvss_score": cve_data.get("cvss_score"),
            "patch_id": cve_data.get("patch_id"),
            "commit_count": len(commit_urls),
        }

        metadata_path = cve_dir / "metadata.json"
        with open(metadata_path, "w") as f:
            json.dump(metadata, f, indent=2)

        logger.info(
            f"Successfully processed {cve_id} using strategy '{strategy}' - {len(commit_urls)} commits"
        )
        return metadata

    def fetch_commit_info_from_linux_repo(self, cve_id: str, commit_urls: List[str]):
        """Fetch detailed commit information from Linux repository"""

        cve_dir = self.base_data_dir / cve_id
        commits_dir = cve_dir / "commits"
        commits_dir.mkdir(exist_ok=True)

        # Track success/failure statistics
        found_commits = []
        missing_commits = []

        for i, commit_url in enumerate(commit_urls, 1):
            # Extract commit hash from URL
            match = re.search(r"/c/([0-9a-fA-F]+)", commit_url)
            if not match:
                match = re.search(r"commit/([0-9a-fA-F]+)", commit_url)
            if not match:
                continue

            commit_hash = match.group(1)

            try:
                # First, try to fetch from remote if commit is not found locally
                self._ensure_commit_exists(commit_hash)

                # Get commit details
                commit_info = {"commit_hash": commit_hash, "commit_url": commit_url}
                # Get commit message and metadata
                show_cmd = ["git", "show", "--no-patch", "--format=fuller", commit_hash]
                result = subprocess.run(
                    show_cmd,
                    cwd=self.linux_repo_path,
                    capture_output=True,
                    text=True,
                    check=True,
                )
                commit_info["commit_details"] = result.stdout
                # Get the actual patch
                patch_cmd = ["git", "show", commit_hash]
                result = subprocess.run(
                    patch_cmd,
                    cwd=self.linux_repo_path,
                    capture_output=True,
                    text=True,
                    check=True,
                )
                commit_info["patch_content"] = result.stdout
                # Get list of changed files
                files_cmd = ["git", "show", "--name-only", "--format=", commit_hash]
                result = subprocess.run(
                    files_cmd,
                    cwd=self.linux_repo_path,
                    capture_output=True,
                    text=True,
                    check=True,
                )
                commit_info["changed_files"] = [
                    f.strip() for f in result.stdout.split("\n") if f.strip()
                ]
                # Get diff stats
                stats_cmd = ["git", "show", "--stat", "--format=", commit_hash]
                result = subprocess.run(
                    stats_cmd,
                    cwd=self.linux_repo_path,
                    capture_output=True,
                    text=True,
                    check=True,
                )
                commit_info["diff_stats"] = result.stdout

                # Save commit information
                commit_file = commits_dir / f"{commit_hash}.json"
                with open(commit_file, "w") as f:
                    json.dump(commit_info, f, indent=2)

                # Also save patch as separate file for easier processing
                patch_file = commits_dir / f"{commit_hash}.patch"
                with open(patch_file, "w") as f:
                    f.write(commit_info["patch_content"])

                found_commits.append(commit_hash)
                logger.info(f"Saved commit info for {commit_hash}")

            except subprocess.CalledProcessError as e:
                missing_commits.append(commit_hash)
                logger.warning(f"Commit {commit_hash} not found locally: {str(e)}")

                # Save commit as missing with URL for later reference
                missing_file = commits_dir / f"{commit_hash}_missing.json"
                missing_info = {
                    "commit_hash": commit_hash,
                    "commit_url": commit_url,
                    "status": "not_found_locally",
                    "error": str(e),
                }
                with open(missing_file, "w") as f:
                    json.dump(missing_info, f, indent=2)

            except Exception as e:
                missing_commits.append(commit_hash)
                logger.error(f"Unexpected error with commit {commit_hash}: {str(e)}")

        logger.info(
            f"Commit summary for {cve_id}: {len(found_commits)} found, {len(missing_commits)} missing"
        )

        return found_commits, missing_commits

    def _ensure_commit_exists(self, commit_hash: str):
        """ENHANCED: Robust multi-strategy commit fetching"""
        # First check if commit exists
        check_cmd = ["git", "cat-file", "-e", commit_hash]
        result = subprocess.run(
            check_cmd, cwd=self.linux_repo_path, capture_output=True, text=True
        )

        if result.returncode == 0:
            return  # Commit exists

        # Strategy 1: Full upstream fetch (only once per run)
        if "origin" not in self._fetched_refs:
            try:
                fetch_cmd = ["git", "fetch", "origin"]
                subprocess.run(
                    fetch_cmd,
                    cwd=self.linux_repo_path,
                    capture_output=True,
                    text=True,
                    check=True,
                    timeout=120,
                )
                self._fetched_refs.add("origin")
            except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
                pass

            # Check if commit exists now
            check_result = subprocess.run(
                check_cmd, cwd=self.linux_repo_path, capture_output=True
            )
            if check_result.returncode == 0:
                return

        # Strategy 2: Try different branches (only once each per run)
        for branch in ["master", "main", "linux-next"]:
            branch_ref = f"origin/{branch}"
            if branch_ref not in self._fetched_refs:
                try:
                    fetch_cmd = ["git", "fetch", "origin", branch]
                    subprocess.run(
                        fetch_cmd,
                        cwd=self.linux_repo_path,
                        capture_output=True,
                        text=True,
                        check=True,
                        timeout=60,
                    )
                    self._fetched_refs.add(branch_ref)

                    check_result = subprocess.run(
                        check_cmd, cwd=self.linux_repo_path, capture_output=True
                    )
                    if check_result.returncode == 0:
                        return
                except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
                    continue

        logger.warning(
            f"Could not fetch commit {commit_hash} after trying multiple strategies"
        )

    def process_single_cve(self, cve_input) -> bool:
        """Process a single CVE using multi-strategy approach"""
        try:
            # Handle both string CVE ID and dict CVE data
            if isinstance(cve_input, str):
                # Create minimal CVE data structure from CVE ID
                cve_data = {
                    "cve_id": cve_input,
                    "severity_label": "UNKNOWN",  # Placeholder
                    "patch_id": None,
                    "classification": "2024+"
                    if int(cve_input.split("-")[1]) >= 2024
                    else "2024-",
                }
            else:
                cve_data = cve_input

            result = self.process_cve_with_strategy(cve_data)
            return result is not None
        except Exception as e:
            logger.error(f"Error processing single CVE: {e}")
            return False

    def process_all_cves(self, limit: Optional[int] = None):
        """Process all CVEs using multi-strategy approach"""
        cves_to_process = self.cve_list[:limit] if limit else self.cve_list

        logger.info(
            f"Processing {len(cves_to_process)} CVEs with multi-strategy approach"
        )
        logger.info("Strategy 1: Direct patch ID (when patch_id available)")
        logger.info(
            "Strategy 2: JSON references filtered (when no patch_id but JSON exists)"
        )
        logger.info("Skipped: CVEs with no viable strategy")

        # Preview first 5 CVEs and their strategies
        logger.info("Preview of first 5 CVEs:")
        for i, cve_data in enumerate(cves_to_process[:5], 1):
            patch_id = cve_data.get("patch_id", "")
            strategy = self.get_extraction_strategy(cve_data)
            logger.info(
                f"  {i}. {cve_data['cve_id']} - Strategy: {strategy} - Patch ID: {'YES' if patch_id.strip() else 'NO'}"
            )

        success_count = 0
        processed_count = 0

        for i, cve_data in enumerate(cves_to_process, 1):
            cve_id = cve_data["cve_id"]
            strategy = self.get_extraction_strategy(cve_data)

            logger.info(
                f"Processing {i}/{len(cves_to_process)}: {cve_id} (Strategy: {strategy})"
            )

            if self.process_single_cve(cve_data):
                success_count += 1

            processed_count += 1

            # Progress update every 10 CVEs
            if i % 10 == 0:
                logger.info(
                    f"Progress: {i}/{len(cves_to_process)} CVEs processed ({success_count} successful)"
                )

            # Add small delay to avoid overwhelming the system
            time.sleep(0.1)

        logger.info(
            f"Completed processing. Success: {success_count}/{processed_count} (Total CVEs: {len(cves_to_process)})"
        )

        logger.info("Strategy Statistics:")
        for strategy, count in self.strategy_stats.items():
            percentage = (count / len(cves_to_process)) * 100 if cves_to_process else 0
            logger.info(f"  {strategy}: {count} ({percentage:.1f}%)")

    def get_summary_stats(self):
        """Get summary statistics of processed CVEs"""
        summary = {
            "total_cves": len(self.cve_list),
            "processed_cves": 0,
            "cves_with_commits": 0,
            "total_commits": 0,
            "found_commits": 0,
            "missing_commits": 0,
            "strategy_breakdown": self.strategy_stats.copy(),
        }

        for cve_data in self.cve_list:
            cve_id = cve_data["cve_id"]
            cve_dir = self.base_data_dir / cve_id
            metadata_file = cve_dir / "metadata.json"

            if metadata_file.exists():
                summary["processed_cves"] += 1

                with open(metadata_file, "r") as f:
                    metadata = json.load(f)

                if metadata.get("commit_urls"):
                    summary["cves_with_commits"] += 1
                    summary["total_commits"] += len(metadata["commit_urls"])

                # Count actual commit files found
                commits_dir = cve_dir / "commits"
                if commits_dir.exists():
                    found_files = list(commits_dir.glob("*.json"))
                    missing_files = list(commits_dir.glob("*_missing.json"))

                    # Don't count missing files in found count
                    actual_found = [
                        f for f in found_files if not f.name.endswith("_missing.json")
                    ]
                    summary["found_commits"] += len(actual_found)
                    summary["missing_commits"] += len(missing_files)

        return summary


def main():
    scraper = CVEDataScraper()

    # Check if git is available
    try:
        subprocess.run(["git", "--version"], capture_output=True, check=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        logger.error("git is not installed. Please install git first.")
        return

    # Setup both repositories
    scraper.setup_repositories()

    scraper.process_all_cves(limit=CVE_CNT_LIMIT)

    # Print summary statistics
    stats = scraper.get_summary_stats()
    logger.info(f"Summary: {stats}")

    if stats["processed_cves"] > 0:
        success_rate = (stats["found_commits"] / max(stats["total_commits"], 1)) * 100
        logger.info(
            f"Commit success rate: {success_rate:.1f}% ({stats['found_commits']}/{stats['total_commits']})"
        )


if __name__ == "__main__":
    main()
