"""
CVE Predictor that supports direct commit hash input
Usage:
  python cve-predict.py CVE-2022-27666  # Original mode
  python cve-predict.py CVE-2022-27666 --commit ebe48d368e97d007bfeb76fcb065d6cfc4c96645  # Direct commit mode
  python cve-predict.py CVE-2022-27666 --commits ebe48d368e97d007bfeb76fcb065d6cfc4c96645,abc123def456  # Multiple commits
"""

from pathlib import Path
import numpy as np
import pickle
import json
import subprocess
import time
from typing import Dict, Optional, List
import logging
import argparse
import re
from cve_data_scraper import CVEDataScraper  # type: ignore
from cve_feature_extraction import CVEFeatureExtractor  # type: ignore


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


class CVEPredictor:
    def __init__(self, cve_id: str, commit_hashes: Optional[List[str]] = None):
        self.cve_id = cve_id
        self.commit_hashes = commit_hashes or []
        self.predict_data_dir = Path(f"data/predict_data/{cve_id}")
        self.model = None
        self.feature_columns = None
        self.label_mapping = None
        self.timing_info = {}

    def load_model(self):
        """Load trained model and its metadata."""
        start_time = time.time()
        logger.info("Loading trained model...")

        try:
            with open("models/cve_severity_model.pkl", "rb") as f:
                self.model = pickle.load(f)

            with open("models/model_metadata.json", "r") as f:
                metadata = json.load(f)

            self.feature_columns = metadata["feature_columns"]
            self.label_mapping = {
                int(k): v for k, v in metadata["label_mapping"].items()
            }

            elapsed_time = time.time() - start_time
            self.timing_info["model_loading"] = elapsed_time
            logger.info(f"Model loaded successfully! ({elapsed_time:.2f}s)")
        except FileNotFoundError as e:
            logger.error(f"Error loading model or metadata: {e}")
            raise

    def check_repository_setup(self):
        logger.info("Checking repository setup...")

        source_vulns_repo = Path("data/linux_security_vulns")
        source_linux_repo = Path("data/linux_kernel_repo")

        missing_repos = []

        if not source_vulns_repo.exists():
            missing_repos.append(
                f"Security vulnerabilities repo: {source_vulns_repo.resolve()}"
            )
        elif not (source_vulns_repo / ".git").exists():
            logger.warning(
                f"Security repo exists but is not a git repository: {source_vulns_repo.resolve()}"
            )

        if not source_linux_repo.exists():
            missing_repos.append(f"Linux kernel repo: {source_linux_repo.resolve()}")
        elif not (source_linux_repo / ".git").exists():
            logger.warning(
                f"Linux repo exists but is not a git repository: {source_linux_repo.resolve()}"
            )

        if missing_repos:
            logger.info("Setting up missing repositories automatically...")

            # Setup missing repositories
            if not source_vulns_repo.exists():
                self._setup_vulns_repo(source_vulns_repo)
            elif not (source_vulns_repo / ".git").exists():
                logger.warning("Removing invalid vulns repo directory...")
                import shutil

                shutil.rmtree(source_vulns_repo)
                self._setup_vulns_repo(source_vulns_repo)

            if not source_linux_repo.exists():
                self._setup_linux_repo(source_linux_repo)
            elif not (source_linux_repo / ".git").exists():
                logger.warning("Removing invalid linux repo directory...")
                import shutil

                shutil.rmtree(source_linux_repo)
                self._setup_linux_repo(source_linux_repo)

            logger.info("Repository setup completed!")

        return True

    def _setup_vulns_repo(self, repo_path):
        """Clone the Linux security vulnerabilities repo"""
        logger.info("Cloning Linux security vulnerabilities repo...")
        try:
            cmd = [
                "git",
                "clone",
                "https://git.kernel.org/pub/scm/linux/security/vulns.git",
                str(repo_path),
            ]
            result = subprocess.run(cmd, capture_output=True, text=True)
            if result.returncode != 0:
                raise RuntimeError(f"Failed to clone vulns repo: {result.stderr}")
            logger.info("Successfully cloned security vulnerabilities repository")
        except Exception as e:
            raise RuntimeError(f"Failed to setup vulnerabilities repository: {e}")

    def _setup_linux_repo(self, repo_path):
        """Clone the Linux kernel git repository"""
        logger.info("Cloning Linux kernel repository (this will take a while)...")
        try:
            cmd = [
                "git",
                "clone",
                "git://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git",
                str(repo_path),
            ]
            result = subprocess.run(cmd, capture_output=True, text=True)
            if result.returncode != 0:
                # Try HTTPS fallback
                logger.warning("Git protocol failed, trying HTTPS...")
                cmd[2] = (
                    "https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git"
                )
                result = subprocess.run(cmd, capture_output=True, text=True)
                if result.returncode != 0:
                    raise RuntimeError(f"Failed to clone Linux repo: {result.stderr}")
            logger.info("Successfully cloned Linux kernel repository")
        except Exception as e:
            raise RuntimeError(f"Failed to setup Linux kernel repository: {e}")

    def setup_predict_environment(self):
        """Setup prediction data directory and link necessary repositories."""
        start_time = time.time()
        logger.info(f"Setting up prediction environment for {self.cve_id}...")

        # check if repositories exist
        if not self.check_repository_setup():
            raise FileNotFoundError(
                "Required repositories are missing. Please run the data scraper first."
            )

        self.predict_data_dir.mkdir(parents=True, exist_ok=True)

        source_vulns_repo = Path("data/linux_security_vulns")
        source_linux_repo = Path("data/linux_kernel_repo")

        predict_vulns_repo_link = self.predict_data_dir / "linux_security_vulns"
        predict_linux_repo_link = self.predict_data_dir / "linux_kernel_repo"

        # Create symlinks
        for link, source in [
            (predict_vulns_repo_link, source_vulns_repo),
            (predict_linux_repo_link, source_linux_repo),
        ]:
            try:
                # Remove existing symlink/directory if it exists
                if link.exists() or link.is_symlink():
                    logger.info(f"Removing existing symlink/directory at {link}")
                    if link.is_symlink():
                        link.unlink()
                    elif link.is_dir():
                        import shutil

                        shutil.rmtree(link)
                    else:
                        link.unlink()

                # Check if source exists
                if not source.exists():
                    logger.error(f"Source path does not exist: {source.resolve()}")
                    raise FileNotFoundError(
                        f"Source repository not found: {source.resolve()}"
                    )

                # Create the symlink
                link.symlink_to(source.resolve())
                logger.info(f"Symlinked {source.resolve()} to {link}")

            except Exception as e:
                logger.error(f"Could not create symlink for {link}: {e}")
                logger.error(f"Source: {source.resolve()}")
                logger.error(f"Target: {link.resolve()}")
                raise

        elapsed_time = time.time() - start_time
        self.timing_info["environment_setup"] = elapsed_time
        logger.info(f"Prediction environment ready! ({elapsed_time:.2f}s)")

    def _ensure_commit_exists(self, commit_hash: str, repo_path: Path):
        """Try to fetch a commit if it doesn't exist locally"""
        # check if commit exists
        check_cmd = ["git", "cat-file", "-e", commit_hash]
        result = subprocess.run(
            check_cmd, cwd=repo_path, capture_output=True, text=True
        )

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

        logger.info(f"Attempting to fetch commit {commit_hash[:8]} from remote")
        try:
            # Try to fetch the specific commit
            fetch_cmd = ["git", "fetch", "origin", commit_hash]
            subprocess.run(
                fetch_cmd, cwd=repo_path, capture_output=True, text=True, timeout=30
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
            # If specific fetch fails, trying a broader fetch
            try:
                fetch_cmd = ["git", "fetch", "origin", "--depth=100"]
                subprocess.run(
                    fetch_cmd, cwd=repo_path, capture_output=True, text=True, timeout=60
                )
            except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
                logger.warning(f"Could not fetch commit {commit_hash[:8]}")

    def scrape_from_commits(self) -> bool:
        """Directly fetch commit data using provided commit hashes."""
        if not self.commit_hashes:
            return False

        start_time = time.time()
        logger.info(
            f"Fetching commits directly for {self.cve_id}: {[h[:8] + '...' for h in self.commit_hashes]}"
        )

        commits_dir = self.predict_data_dir / "commits"
        commits_dir.mkdir(exist_ok=True)

        linux_repo_path = self.predict_data_dir / "linux_kernel_repo"
        found_commits = []

        for commit_hash in self.commit_hashes:
            try:
                # Ensure we have the commit (try to fetch if needed)
                self._ensure_commit_exists(commit_hash, linux_repo_path)

                commit_info = {"commit_hash": commit_hash, "cve_id": self.cve_id}

                # Get commit details
                show_cmd = ["git", "show", "--no-patch", "--format=fuller", commit_hash]
                result = subprocess.run(
                    show_cmd,
                    cwd=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=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=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=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)

                # Save patch as separate file
                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"Successfully fetched commit {commit_hash[:8]}")

            except subprocess.CalledProcessError as e:
                logger.error(f"Failed to fetch commit {commit_hash}: {e}")
                continue
            except Exception as e:
                logger.error(f"Unexpected error with commit {commit_hash}: {e}")
                continue

        # Creating metadata in the format expected by CVEFeatureExtractor
        metadata = {
            "cve_id": self.cve_id,
            "source_method": "direct_commits",
            "source_files": [f"commits/{commit}.json" for commit in found_commits],
            "json_data": None,  # We don't have JSON data from security repo
            "commit_urls": [
                f"https://git.kernel.org/stable/c/{commit}" for commit in found_commits
            ],
            "affected_files": [],  # Will be filled by analyzing patches
            "scraped_at": time.time(),
            "method": "direct_commits",
        }

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

        elapsed_time = time.time() - start_time
        self.timing_info["data_scraping"] = elapsed_time

        success = len(found_commits) > 0
        if success:
            logger.info(
                f"Successfully scraped {len(found_commits)}/{len(self.commit_hashes)} commits ({elapsed_time:.2f}s)"
            )
        else:
            logger.error(f"Failed to scrape any commits ({elapsed_time:.2f}s)")

        return success

    def scrape_cve_data(self) -> bool:
        """Scrape data for the CVE - either from commits or traditional method."""
        # If we have specific commit hashes, use them directly
        if self.commit_hashes:
            return self.scrape_from_commits()

        # else, fall back to traditional scraping
        start_time = time.time()
        logger.info(f"Scraping data for {self.cve_id} using traditional method...")

        try:
            scraper = CVEDataScraper(base_data_dir=str(self.predict_data_dir.parent))
            scraper.vulns_repo_path = self.predict_data_dir / "linux_security_vulns"
            scraper.linux_repo_path = self.predict_data_dir / "linux_kernel_repo"

            # Now that paths are correct, ensure stable remote is available
            scraper._ensure_stable_remote_available()

            success = scraper.process_single_cve(self.cve_id)

            elapsed_time = time.time() - start_time
            self.timing_info["data_scraping"] = elapsed_time

            if success:
                logger.info(f"Successfully scraped {self.cve_id} ({elapsed_time:.2f}s)")
            else:
                logger.error(f"Failed to scrape {self.cve_id} ({elapsed_time:.2f}s)")
            return success
        except Exception as e:
            logger.error(f"Error during data scraping for {self.cve_id}: {e}")
            return False

    def extract_features(self) -> Optional[Dict]:
        """Extract features for the CVE using the feature extraction script."""
        start_time = time.time()
        logger.info(f"Extracting features for {self.cve_id}...")

        try:
            engineer = CVEFeatureExtractor(data_dir=str(self.predict_data_dir.parent))
            features_result = engineer.process_cve_with_timeout(
                self.cve_id, timeout_seconds=60
            )

            elapsed_time = time.time() - start_time
            self.timing_info["feature_extraction"] = elapsed_time

            if features_result:
                logger.info(f"Features extracted successfully! ({elapsed_time:.2f}s)")
                clean_features = {
                    k: v
                    for k, v in features_result.items()
                    if k in self.feature_columns
                }
                return clean_features
            else:
                logger.error(
                    f"Failed to extract features for {self.cve_id}! ({elapsed_time:.2f}s)"
                )
                return None
        except Exception as e:
            logger.error(f"Error during feature extraction for {self.cve_id}: {e}")
            return None

    def predict_severity(self) -> Optional[Dict]:
        """Predict severity for the CVE based on extracted features."""
        start_time = time.time()

        features = self.extract_features()
        if not features:
            return None

        logger.info(f"Making severity prediction for {self.cve_id}...")

        # Prepare feature vector
        feature_vector = []
        for feature_name in self.feature_columns:
            value = features.get(feature_name, False)
            if isinstance(value, bool):
                value = int(value)
            feature_vector.append(value)

        # Make prediction
        try:
            X = np.array(feature_vector).reshape(1, -1)
            prediction = self.model.predict(X)[0]
            prediction_proba = self.model.predict_proba(X)[0]
        except Exception as e:
            logger.error(f"Error during model prediction for {self.cve_id}: {e}")
            return None

        elapsed_time = time.time() - start_time
        self.timing_info["prediction"] = elapsed_time

        predicted_severity_label = self.label_mapping.get(prediction, "UNKNOWN")

        result = {
            "cve_id": self.cve_id,
            "predicted_severity": predicted_severity_label,
            "confidence": float(max(prediction_proba)),
            "probabilities": {
                self.label_mapping[i]: float(prob)
                for i, prob in enumerate(prediction_proba)
            },
            "active_features": [
                feature
                for feature in self.feature_columns
                if features.get(feature, False)
            ],
            "input_method": "direct_commits"
            if self.commit_hashes
            else "traditional_scraping",
            "commit_hashes": self.commit_hashes,
        }
        return result

    def predict(self) -> Optional[Dict]:
        """Run the complete prediction pipeline for a single CVE."""
        total_start_time = time.time()
        logger.info(f"Starting prediction for {self.cve_id}")

        if self.commit_hashes:
            logger.info(
                f"Using direct commit mode with {len(self.commit_hashes)} commits"
            )

        try:
            self.load_model()
            self.setup_predict_environment()

            if not self.scrape_cve_data():
                logger.error("Could not scrape CVE data. Aborting prediction.")
                return None

            result = self.predict_severity()

            total_elapsed_time = time.time() - total_start_time
            self.timing_info["total_processing"] = total_elapsed_time

            if result:
                result["timing_info"] = self.timing_info.copy()

                result_file = self.predict_data_dir / f"{self.cve_id}_prediction.json"
                with open(result_file, "w") as f:
                    json.dump(result, f, indent=2)

                logger.info(
                    f"Prediction for {self.cve_id} completed in {total_elapsed_time:.2f}s"
                )
            return result
        except Exception as e:
            logger.error(
                f"Unexpected error during prediction pipeline for {self.cve_id}: {e}"
            )
            return None

    def print_result(self, result: Optional[Dict]):
        """Print the prediction result and detailed timing information."""
        if not result:
            print("❌ Prediction failed or was skipped!")
            return

        print("\n" + "=" * 60)
        print("CVE SEVERITY PREDICTION")
        print("=" * 60)
        print(f"CVE ID: {result['cve_id']}")
        print(f"Input Method: {result.get('input_method', 'traditional_scraping')}")
        if result.get("commit_hashes"):
            print(
                f"Commit Hashes: {', '.join([h[:8] + '...' for h in result['commit_hashes']])}"
            )
        print(f"Predicted Severity: {result['predicted_severity']}")
        print(f"Confidence: {result['confidence']:.2%}")
        print()
        print("All Probabilities:")
        for severity, prob in result["probabilities"].items():
            print(f"  {severity}: {prob:.2%}")
        print()
        print("Top Active Features:")
        for feature in result["active_features"][:10]:
            print(f"  ✓ {feature}")

        if "timing_info" in result:
            print("\nProcessing Times:")
            timing = result["timing_info"]
            print(f"  Model Loading: {timing.get('model_loading', 0):.2f}s")
            print(f"  Environment Setup: {timing.get('environment_setup', 0):.2f}s")
            print(f"  Data Scraping: {timing.get('data_scraping', 0):.2f}s")
            print(f"  Feature Extraction: {timing.get('feature_extraction', 0):.2f}s")
            print(f"  Prediction: {timing.get('prediction', 0):.2f}s")
            print(f"  Total Processing: {timing.get('total_processing', 0):.2f}s")

        print("=" * 60)


def main():
    parser = argparse.ArgumentParser(description="Predict CVE severity")

    parser.add_argument("cve_id", type=str, help="CVE ID to predict")
    parser.add_argument(
        "--commit", type=str, help="Single commit hash to analyze directly"
    )
    parser.add_argument(
        "--commits", type=str, help="Comma-separated list of commit hashes"
    )

    args = parser.parse_args()

    # Parse commit hashes
    commit_hashes = []
    if args.commit:
        commit_hashes = [args.commit.strip()]
    elif args.commits:
        commit_hashes = [c.strip() for c in args.commits.split(",") if c.strip()]

    # Validate commit hashes
    for commit in commit_hashes:
        if not re.match(r"^[0-9a-fA-F]{40}$", commit):
            logger.error(f"Invalid commit hash format: {commit}")
            logger.error("Commit hashes must be 40-character hexadecimal strings")
            return

    # Create predictor instance
    predictor = CVEPredictor(args.cve_id, commit_hashes)

    # Run prediction
    result = predictor.predict()

    # Print results
    predictor.print_result(result)


if __name__ == "__main__":
    main()
