import argparse
import logging
import shutil
from pathlib import Path

import evaluate
import numpy as np
from codecarbon import EmissionsTracker
from datasets import load_dataset
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments,
)

"""
Automatically classify new vulnerabilities based on their descriptions,
even if they don't have CVSS scores.

Currently tested with distilbert-base-uncased and roberta-base.
"""

# Logging setup
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Define severity label mapping
SEVERITY_MAPPING = {"Low": 0, "Medium": 1, "High": 2, "Critical": 3}


def compute_metrics(eval_pred):
    """Compute accuracy and F1-score for model evaluation."""
    metric = evaluate.load("accuracy")
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)


# Define severity mapping function
def map_cvss_to_severity(example, score_strategy="first"):
    """
    Map a CVSS score to a severity label and add it to the example.

    Parameters
    ----------
    example : dict
        A dictionary containing CVSS scores. Possible keys include:
        - "cvss_v4_0"
        - "cvss_v3_1"
        - "cvss_v3_0"
        - "cvss_v2_0"

    score_strategy : str, optional
        Strategy to select the CVSS score to use. Options are:
        - "first"  : Use the first non-null score as it appears in the CVE data (original field order).
        - "latest" : Use the score from the most recent CVSS version available
                     (v4.0 → v3.1 → v3.0 → v2.0).
        - "mean"   : Compute the arithmetic mean of all available non-null scores.

    Returns
    -------
    dict
        The input `example` dictionary with an added "severity_label" field,
        which can be one of: "Critical", "High", "Medium", "Low", or "Unknown".
    """

    def to_float(value):
        try:
            return float(value) if value is not None else None
        except ValueError:
            return None

    version_priority = ["cvss_v4_0", "cvss_v3_1", "cvss_v3_0", "cvss_v2_0"]

    if score_strategy == "first":
        severity_score = None
        for key, value in example.items():
            if key in version_priority:
                score = to_float(value)
                if score is not None:
                    severity_score = score
                    break

    elif score_strategy == "latest":
        for key in version_priority:
            score = to_float(example.get(key))
            if score is not None:
                severity_score = score
                break
        else:
            severity_score = None

    elif score_strategy == "mean":
        scores = [to_float(example.get(k)) for k in version_priority]
        filtered_scores = [s for s in scores if s is not None]
        severity_score = (
            sum(filtered_scores) / len(filtered_scores) if filtered_scores else None
        )

    else:
        raise ValueError(f"Unknown score_strategy: {score_strategy}")

    if severity_score is None:
        severity_label = "Unknown"
    elif severity_score >= 9.0:
        severity_label = "Critical"
    elif severity_score >= 7.0:
        severity_label = "High"
    elif severity_score >= 4.0:
        severity_label = "Medium"
    else:
        severity_label = "Low"

    example["severity_label"] = severity_label
    return example


def train(
    base_model,
    dataset_id,
    repo_id,
    model_save_dir="./vulnerability-classify",
    push_to_hub=True,
    use_cache=False,
    codecarbon=True,
):
    # Load dataset from Hugging Face
    dataset = load_dataset(dataset_id)

    # Map severity labels
    dataset = dataset.map(map_cvss_to_severity)

    # Filter out entries with no severity_label and with unknown keys
    dataset = dataset.filter(lambda x: "severity_label" in x)
    dataset = dataset.filter(lambda x: x["severity_label"] in SEVERITY_MAPPING)

    # Tokenization with labels
    tokenizer = AutoTokenizer.from_pretrained(base_model)
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    def tokenize_function(elem):
        tokenized = tokenizer(
            elem["description"],
            padding="max_length",
            truncation=True,
        )

        # Convert list of severity labels to integers
        tokenized["labels"] = [
            SEVERITY_MAPPING[label] for label in elem["severity_label"]
        ]

        return tokenized

    tokenized_datasets = dataset.map(tokenize_function, batched=True)
    # print(tokenized_datasets["test"])

    # Define model
    num_labels = len(SEVERITY_MAPPING)  # 4 classes

    model = AutoModelForSequenceClassification.from_pretrained(
        base_model,
        num_labels=num_labels,
        id2label={
            v: k for k, v in SEVERITY_MAPPING.items()
        },  # Mapping indices to labels
        label2id=SEVERITY_MAPPING,  # Mapping labels to indices
    )

    # Define training arguments
    training_args = TrainingArguments(
        output_dir=model_save_dir,
        eval_strategy="epoch",
        save_strategy="epoch",
        learning_rate=3e-5,
        per_device_train_batch_size=8 if "roberta" in base_model else 16,
        per_device_eval_batch_size=8 if "roberta" in base_model else 16,
        num_train_epochs=5,
        weight_decay=0.01,
        logging_dir="./logs",
        logging_steps=10,
        save_total_limit=2,
        load_best_model_at_end=True,
        use_cache=use_cache,
        hub_model_id=repo_id,
    )

    # Create Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_datasets["train"],
        eval_dataset=tokenized_datasets["test"],
        data_collator=data_collator,
        compute_metrics=compute_metrics,
    )

    # Train model
    tracker = None
    if codecarbon:
        tracker = EmissionsTracker(project_name="VulnTrain", allow_multiple_runs=True)
        tracker.start()
    try:
        trainer.train()
    finally:
        if tracker is not None:
            tracker.stop()
        model.save_pretrained(model_save_dir)
        tokenizer.save_pretrained(model_save_dir)

    if push_to_hub:
        trainer.push_to_hub()
        tokenizer.push_to_hub(repo_id)


def main():
    parser = argparse.ArgumentParser(
        description="Train a vulnerability classification model with a mapping on the severity."
    )
    parser.add_argument(
        "--base-model",
        dest="base_model",
        default="distilbert-base-uncased",
        choices=["distilbert-base-uncased", "roberta-base"],
        help="Base model to use.",
    )
    parser.add_argument(
        "--dataset-id",
        dest="dataset_id",
        default="CIRCL/vulnerability-scores",
        help="Path of the dataset. Local dataset or repository on the HF hub.",
    )
    parser.add_argument(
        "--repo-id",
        dest="repo_id",
        required=True,
        help="The name of the repository you want to push your object to. It should contain your organization name when pushing to a given organization.",
    )
    parser.add_argument(
        "--model-save-dir",
        dest="model_save_dir",
        default="results",
        help="The path to a directory where the tokenizer and the model will be saved.",
    )
    parser.add_argument(
        "--no-codecarbon",
        dest="no_codecarbon",
        action="store_true",
        default=False,
        help="Disable CodeCarbon emissions tracking.",
    )
    parser.add_argument(
        "--no-push",
        dest="no_push",
        action="store_true",
        default=False,
        help="Disable pushing the model and tokenizer to Hugging Face Hub.",
    )
    parser.add_argument(
        "--no-cache",
        dest="no_cache",
        action="store_true",
        default=False,
        help="Disable cache for the model during training.",
    )

    args = parser.parse_args()

    dir_path = Path(args.model_save_dir)
    if dir_path.exists() and dir_path.is_dir():
        shutil.rmtree(dir_path)

    logger.info(f"Using base model: {args.base_model}")
    logger.info(f"Dataset ID: {args.dataset_id}")
    logger.info(f"Repo ID: {args.repo_id}")
    logger.info(f"Saving model to: {args.model_save_dir}")
    logger.info("Starting the training process…")

    train(
        args.base_model,
        args.dataset_id,
        args.repo_id,
        args.model_save_dir,
        push_to_hub=not args.no_push,
        use_cache=not args.no_cache,
        codecarbon=not args.no_codecarbon,
    )


if __name__ == "__main__":
    main()
