"""
Split cve_dataset.csv into training and testing datasets based on JSON splits
"""

import pandas as pd
import json


def split_dataset():
    print("Loading datasets...")

    # Load the main dataset
    df = pd.read_csv("data/cve_dataset.csv")
    print(f"Loaded {len(df)} CVEs from cve_dataset.csv")
    print("Using train_kernel_cves.json and test_kernel_cves.json for split")

    # Load train and test CVE lists
    with open("data/train_kernel_cves.json", "r") as f:
        train_cves = json.load(f)

    with open("data/test_kernel_cves.json", "r") as f:
        test_cves = json.load(f)

    # Extract CVE IDs
    train_cve_ids = {cve["cve_id"] for cve in train_cves}
    test_cve_ids = {cve["cve_id"] for cve in test_cves}

    print(f"Training CVEs: {len(train_cve_ids)}")
    print(f"Testing CVEs: {len(test_cve_ids)}")

    # Split the dataset
    train_df = df[df["cve_id"].isin(train_cve_ids)]
    test_df = df[df["cve_id"].isin(test_cve_ids)]

    print(f"Training dataset: {len(train_df)} CVEs")
    print(f"Testing dataset: {len(test_df)} CVEs")

    # Save the split datasets
    train_df.to_csv("data/cve_training_dataset.csv", index=False)
    test_df.to_csv("data/cve_testing_dataset.csv", index=False)

    print("✅ Datasets saved:")
    print("   - data/cve_training_dataset.csv")
    print("   - data/cve_testing_dataset.csv")


if __name__ == "__main__":
    split_dataset()
