"""Prepare fixed NSL-KDD train/test tables for Module 6 notebooks."""

from __future__ import annotations

import argparse
from pathlib import Path

import pandas as pd


HERE = Path(__file__).resolve().parent
RAW_DIR = HERE / "raw"

FEATURE_NAMES = [
    "duration", "protocol_type", "service", "flag", "src_bytes", "dst_bytes",
    "land", "wrong_fragment", "urgent", "hot", "num_failed_logins",
    "logged_in", "num_compromised", "root_shell", "su_attempted", "num_root",
    "num_file_creations", "num_shells", "num_access_files", "num_outbound_cmds",
    "is_host_login", "is_guest_login", "count", "srv_count", "serror_rate",
    "srv_serror_rate", "rerror_rate", "srv_rerror_rate", "same_srv_rate",
    "diff_srv_rate", "srv_diff_host_rate", "dst_host_count",
    "dst_host_srv_count", "dst_host_same_srv_rate", "dst_host_diff_srv_rate",
    "dst_host_same_src_port_rate", "dst_host_srv_diff_host_rate",
    "dst_host_serror_rate", "dst_host_srv_serror_rate", "dst_host_rerror_rate",
    "dst_host_srv_rerror_rate",
]

CATEGORICAL_FEATURES = ["protocol_type", "service", "flag"]

ATTACK_GROUPS = {
    "normal": "Normal",
    "back": "DoS", "land": "DoS", "neptune": "DoS", "pod": "DoS",
    "smurf": "DoS", "teardrop": "DoS", "apache2": "DoS", "mailbomb": "DoS",
    "processtable": "DoS", "udpstorm": "DoS", "worm": "DoS",
    "ipsweep": "Probe", "nmap": "Probe", "portsweep": "Probe", "satan": "Probe",
    "mscan": "Probe", "saint": "Probe",
    "ftp_write": "R2L", "guess_passwd": "R2L", "imap": "R2L",
    "multihop": "R2L", "phf": "R2L", "spy": "R2L", "warezclient": "R2L",
    "warezmaster": "R2L", "sendmail": "R2L", "named": "R2L",
    "snmpgetattack": "R2L", "snmpguess": "R2L", "xlock": "R2L",
    "xsnoop": "R2L", "httptunnel": "R2L",
    "buffer_overflow": "U2R", "loadmodule": "U2R", "perl": "U2R",
    "rootkit": "U2R", "ps": "U2R", "sqlattack": "U2R", "xterm": "U2R",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--full-train",
        action="store_true",
        help="Use KDDTrain+.txt instead of KDDTrain+_20Percent.txt.",
    )
    return parser.parse_args()


def find_one(filename: str) -> Path:
    matches = list(RAW_DIR.rglob(filename)) if RAW_DIR.exists() else []
    if len(matches) != 1:
        raise FileNotFoundError(
            f"Expected exactly one {filename!r} below {RAW_DIR}. "
            "Follow data/README.md and verify the extracted files."
        )
    return matches[0]


def load_nsl_kdd(path: Path) -> pd.DataFrame:
    columns = FEATURE_NAMES + ["attack_label", "difficulty"]
    frame = pd.read_csv(path, header=None, names=columns, low_memory=False)
    frame["attack_label"] = (
        frame["attack_label"].astype(str).str.strip().str.rstrip(".").str.lower()
    )

    unknown = sorted(set(frame["attack_label"]) - set(ATTACK_GROUPS))
    if unknown:
        raise ValueError(f"Unknown attack labels in {path.name}: {unknown}")

    for column in CATEGORICAL_FEATURES:
        frame[column] = frame[column].astype(str).str.strip().str.lower()

    numeric_features = [name for name in FEATURE_NAMES if name not in CATEGORICAL_FEATURES]
    for column in numeric_features + ["difficulty"]:
        frame[column] = pd.to_numeric(frame[column], errors="raise")

    frame["attack_group"] = frame["attack_label"].map(ATTACK_GROUPS)
    frame["binary_label"] = (frame["attack_group"] != "Normal").astype(int)
    return frame[FEATURE_NAMES + ["attack_label", "attack_group", "binary_label", "difficulty"]]


def main() -> None:
    args = parse_args()
    train_name = "KDDTrain+.txt" if args.full_train else "KDDTrain+_20Percent.txt"
    train = load_nsl_kdd(find_one(train_name))
    test = load_nsl_kdd(find_one("KDDTest+.txt"))

    train_path = HERE / "nsl_kdd_train.csv"
    test_path = HERE / "nsl_kdd_test.csv"
    manifest_path = HERE / "feature_manifest.csv"

    train.to_csv(train_path, index=False)
    test.to_csv(test_path, index=False)

    manifest = pd.DataFrame({
        "feature": FEATURE_NAMES,
        "feature_type": [
            "categorical" if name in CATEGORICAL_FEATURES else "numeric"
            for name in FEATURE_NAMES
        ],
        "model_feature": True,
    })
    manifest.to_csv(manifest_path, index=False)

    print(f"Created {train_path} with shape {train.shape}")
    print(f"Created {test_path} with shape {test.shape}")
    print(f"Created {manifest_path}")
    print("Train attack groups:")
    print(train["attack_group"].value_counts())
    print("Test attack groups:")
    print(test["attack_group"].value_counts())


if __name__ == "__main__":
    main()