"""
MalwareGuard — Benign URL Validation Dataset Generator (20k Unseen Legitimate URLs).
Reads Majestic Million domains from rank 300,000 onwards (completely unseen during training)
and generates 20,000 realistic nested path and parameter URLs.
"""

import os
import random
import pandas as pd

DATA_DIR = "data"
MAJESTIC_CSV = os.path.join(DATA_DIR, "majestic_million.csv")
OUTPUT_CSV = os.path.join(DATA_DIR, "benign_20k_urls.csv")

def main():
    print("[INFO] Loading Majestic Million domains from rank 300,000 onwards...")
    if not os.path.exists(MAJESTIC_CSV):
        print(f"[ERROR] {MAJESTIC_CSV} does not exist. Please run compile_megacorpus.py first.")
        return

    # Load 40,000 domains starting from row 300,000 to choose from
    df = pd.read_csv(MAJESTIC_CSV, skiprows=300000, nrows=40000, header=None, names=["GlobalRank", "TldRank", "Domain", "TLD", "RefSubNets", "RefIPs", "IDN_Domain", "IDN_TLD", "PrevGlobalRank", "PrevTldRank", "PrevRefSubNets", "PrevRefIPs"])
    
    domains = df["Domain"].dropna().astype(str).tolist()
    random.seed(42)
    random.shuffle(domains)

    # Path patterns to randomly apply to the domains
    nested_patterns = [
        "",
        "www.{}",
        "http://{}",
        "https://{}",
        "https://www.{}",
        "https://{}/login",
        "https://{}/signin",
        "https://www.{}/login",
        "https://www.{}/signin",
        "https://{}/index.html",
        "https://{}/about",
        "https://{}/contact",
        "https://{}/privacy-policy",
        "https://{}/terms-of-service",
        "https://{}/store/products",
        "https://{}/blog/posts",
        "https://{}/wiki/Main_Page",
        "https://{}/en-us/windows",
        "https://{}/trending",
        "https://{}/store/products/details",
        "https://{}/assets/js/main.js",
        "https://{}/category/item/view",
        "https://{}/search?q=query&category=1",
        "https://{}/view?id=12345&mode=details",
        "https://{}/profile/settings/edit",
        "https://{}/download/app/latest",
        "https://{}/docs/api/v1/intro",
        "https://{}/faq/help-center"
    ]

    print("[INFO] Generating 20,000 realistic, unseen benign URLs...")
    benign_urls = []
    
    # We will draw domains and assign a random nested pattern
    for i in range(20000):
        dom = domains[i % len(domains)].lower()
        pattern = random.choice(nested_patterns)
        if "{}" in pattern:
            url = pattern.format(dom)
        else:
            url = dom
        benign_urls.append(url)

    validation_df = pd.DataFrame({"url": benign_urls, "label": 0})
    
    # Save the file
    validation_df.to_csv(OUTPUT_CSV, index=False)
    print(f"[SUCCESS] Generated 20,000 benign URLs test set at {OUTPUT_CSV}")

if __name__ == "__main__":
    main()