"""
preprocessing.py
================
Unified preprocessing pipeline for REST API anomaly detection research.

Supports three datasets:
  - CSIC 2010       : Raw HTTP logs (text format)
  - CIC-IDS2018     : Network flow CSVs (CICFlowMeter output)
  - UNSW-NB15       : Network flow CSVs (9 attack categories)

Each dataset is converted into sliding-window sessions ready for LSTM-Autoencoder training. Training sessions contain ONLY normal traffic.
Test sessions contain both normal and attack traffic with binary labels.

Author : K.A.D.S.D. Kandanaarachchi (2020/ICT/19)
Project: Detecting Anomalous REST API Traffic — IT4216
"""

import json
import logging
import re
from collections import Counter
from pathlib import Path
from urllib.parse import parse_qs, unquote, urlparse

import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler

# Logging

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)


# SHARED UTILITIES


def build_flow_sessions(
    X: np.ndarray, window_size: int = 5, step: int = 1
) -> np.ndarray:
    """
    Slide a window over a sequence of feature vectors.

    Parameters
    ----------
    X           : (n_flows, n_features) float array
    window_size : number of consecutive flows per session
    step        : stride between windows (1 = fully overlapping)

    Returns
    -------
    (n_sessions, window_size, n_features) float32 array
    """
    sessions = []
    for i in range(0, len(X) - window_size + 1, step):
        sessions.append(X[i : i + window_size])
    return np.array(sessions, dtype=np.float32)


def save_arrays(
    out_dir: Path,
    prefix: str,
    X_train: np.ndarray,
    X_test: np.ndarray,
    y_test: np.ndarray,
) -> None:
    """Save train/test numpy arrays and print a summary."""
    out_dir.mkdir(parents=True, exist_ok=True)
    np.save(out_dir / f"X_train_{prefix}.npy", X_train)
    np.save(out_dir / f"X_test_{prefix}.npy", X_test)
    np.save(out_dir / f"y_test_{prefix}.npy", y_test)
    log.info("Saved  X_train_%s %s", prefix, X_train.shape)
    log.info("Saved  X_test_%s  %s", prefix, X_test.shape)
    log.info("Saved  y_test_%s  %s", prefix, y_test.shape)


# DATASET 1 — CSIC 2010

# HTTP parser


def _parse_http_log(filepath: Path, label: str) -> pd.DataFrame:
    """
    Parse a raw CSIC 2010 HTTP log file into a DataFrame.

    Each HTTP request is separated by a blank line.
    Handles Latin-1 encoding used by the original dataset.
    """
    with open(filepath, "r", encoding="latin-1") as fh:
        content = fh.read()

    records = []
    for raw in content.strip().split("\n\n"):
        lines = raw.strip().split("\n")
        if not lines or not lines[0].strip():
            continue

        record: dict = {"raw": raw, "label": label}

        # Request line: METHOD /path HTTP/1.x
        parts = lines[0].strip().split(" ")
        record["method"] = parts[0] if parts else ""
        record["url"] = parts[1] if len(parts) > 1 else ""

        if record["method"] not in {
            "GET",
            "POST",
            "PUT",
            "DELETE",
            "HEAD",
            "OPTIONS",
            "PATCH",
        }:
            continue

        # Headers and optional body
        in_body, body_lines = False, []
        for line in lines[1:]:
            if line.strip() == "":
                in_body = True
                continue
            if in_body:
                body_lines.append(line)
            elif ":" in line:
                key, _, val = line.partition(":")
                record[key.strip().lower().replace("-", "_")] = val.strip()

        record["body"] = "\n".join(body_lines).strip()
        record["has_body"] = len(record["body"]) > 0
        record["raw_length"] = len(raw)
        records.append(record)

    return pd.DataFrame(records)


# URL abstraction

_SUSPICIOUS = {
    "'",
    '"',
    "<",
    ">",
    ";",
    "--",
    "DROP",
    "SELECT",
    "UNION",
    "INSERT",
    "DELETE",
    "SCRIPT",
    "ALERT",
    "../",
    "%27",
    "%3C",
    "%3E",
}


def _abstract_url(url: str) -> str:
    """
    Replace variable URL segments and parameter values with typed placeholders.

    Examples
    --------
    /api/users/1054/profile → /api/users/{INT}/profile
    /shop/add?id=3&qty=2&name=Wine → /shop/add?id={INT}&qty={INT}&name={STR}

    /login?user=admin%27%3B+DROP+TABLE → /login?user={INJECT}
    """
    if not url or pd.isna(url):
        return ""

    try:
        parsed = urlparse(url)

        # Abstract numeric path segments
        abstracted_parts = []
        for part in parsed.path.split("/"):
            if re.match(r"^\d+$", part):
                abstracted_parts.append("{INT}")
            elif re.match(r"^[0-9a-f-]{32,}$", part, re.I):
                abstracted_parts.append("{HASH}")
            else:
                abstracted_parts.append(part)
        abs_path = "/".join(abstracted_parts)

        if not parsed.query:
            return abs_path

        # Abstract query parameter values
        abs_params = []
        for param in parsed.query.split("&"):
            if "=" not in param:
                abs_params.append(param)
                continue
            key, _, value = param.partition("=")
            decoded = unquote(value)
            if re.match(r"^\d+$", decoded):
                abs_params.append(f"{key}={{INT}}")
            elif re.match(r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+$", decoded):
                abs_params.append(f"{key}={{EMAIL}}")
            elif len(decoded) == 0:
                abs_params.append(f"{key}={{EMPTY}}")
            elif any(s in decoded.upper() for s in _SUSPICIOUS):
                abs_params.append(f"{key}={{INJECT}}")
            else:
                abs_params.append(f"{key}={{STR}}")

        return f"{abs_path}?{'&'.join(abs_params)}"

    except Exception:
        return url


def _extract_url_features(url: str) -> dict:
    """Extract numeric features from a raw URL string."""
    empty = {
        "path": "",
        "query_string": "",
        "param_count": 0,
        "path_depth": 0,
        "has_suspicious_chars": False,
        "query_length": 0,
    }
    if not url or pd.isna(url):
        return empty
    try:
        parsed = urlparse(url)
        params = parse_qs(parsed.query)
        decoded = unquote(url).upper()
        return {
            "path": parsed.path,
            "query_string": parsed.query,
            "param_count": len(params),
            "path_depth": len([p for p in parsed.path.split("/") if p]),
            "has_suspicious_chars": any(s in decoded for s in _SUSPICIOUS),
            "query_length": len(parsed.query),
        }
    except Exception:
        return empty


# Vocabulary & encoding


def _build_vocab(sessions: list[list[str]], min_freq: int = 2) -> dict:
    """
    Build an integer vocabulary from URL token sequences.

    Special tokens
    --------------
    <PAD> = 0   padding (unused with fixed-length windows)
    <UNK> = 1   URLs not seen during training (strong attack signal)
    """
    counts = Counter(url for session in sessions for url in session)
    vocab = {"<PAD>": 0, "<UNK>": 1}
    for url, cnt in counts.most_common():
        if cnt >= min_freq:
            vocab[url] = len(vocab)
    return vocab


def _encode_sessions(sessions: list[list[str]], vocab: dict) -> list[list[int]]:
    return [[vocab.get(url, 1) for url in session] for session in sessions]


def _build_url_sessions(
    df: pd.DataFrame, window_size: int, label_filter: str
) -> list[list[str]]:
    """Slide a window over the abstracted URL column for one label."""
    subset = df[df["label"] == label_filter]["url_abstracted"].reset_index(drop=True)
    return [
        subset.iloc[i : i + window_size].tolist()
        for i in range(len(subset) - window_size + 1)
    ]


# Public entry point


def preprocess_csic2010(
    data_dir: str | Path,
    output_dir: str | Path,
    window_size: int = 5,
    min_freq: int = 2,
) -> dict:
    """
    Full preprocessing pipeline for CSIC 2010.

    Expected files in data_dir
    --------------------------
    normalTrafficTraining.txt
    normalTrafficTest.txt
    anomalousTrafficTest.txt

    Output files in output_dir
    --------------------------
    X_train_csic2010_w{window_size}.npy   (n_sessions, window_size)   int32
    X_test_csic2010_w{window_size}.npy    (n_sessions, window_size)   int32
    y_test_csic2010_w{window_size}.npy    (n_sessions,)               int32
    vocab_csic2010_w{window_size}.json

    Filenames are suffixed with the window size so outputs from
    different --window runs coexist on disk instead of overwriting
    each other, letting you compare results across window sizes.

    Returns
    -------
    dict with shapes and vocabulary size
    """
    data_dir = Path(data_dir)
    output_dir = Path(output_dir)
    prefix = f"csic2010_w{window_size}"

    # Parse raw HTTP logs
    log.info("CSIC 2010 — parsing HTTP logs...")
    df_train = _parse_http_log(data_dir / "normalTrafficTraining.txt", "normal")
    df_normal = _parse_http_log(data_dir / "normalTrafficTest.txt", "normal")
    df_attack = _parse_http_log(data_dir / "anomalousTrafficTest.txt", "attack")

    df_train["split"] = "train"
    df_normal["split"] = "test"
    df_attack["split"] = "test"

    df = pd.concat([df_train, df_normal, df_attack], ignore_index=True)
    log.info(
        "  Total records: %d  (normal=%d  attack=%d)",
        len(df),
        len(df[df.label == "normal"]),
        len(df[df.label == "attack"]),
    )

    # URL abstraction
    log.info("  Abstracting URLs...")
    df["url_abstracted"] = df["url"].apply(_abstract_url)

    url_feats = df["url"].apply(_extract_url_features).apply(pd.Series)
    df = pd.concat([df, url_feats], axis=1)

    vocab_raw = df["url"].nunique()
    vocab_abs = df["url_abstracted"].nunique()
    log.info(
        "  Vocabulary: %d raw → %d abstracted  (%.1f%% reduction)",
        vocab_raw,
        vocab_abs,
        (1 - vocab_abs / vocab_raw) * 100,
    )

    # Build sliding-window sessions
    log.info("  Building sessions (window=%d)...", window_size)
    train_sessions = _build_url_sessions(df[df.split == "train"], window_size, "normal")
    test_n_sessions = _build_url_sessions(df[df.split == "test"], window_size, "normal")
    test_a_sessions = _build_url_sessions(df[df.split == "test"], window_size, "attack")

    # Vocabulary and encoding
    vocab_data = _build_vocab(train_sessions, min_freq=min_freq)
    log.info("  Vocabulary size: %d tokens", len(vocab_data))

    X_train = np.array(_encode_sessions(train_sessions, vocab_data), dtype=np.int32)
    X_test_normal = np.array(_encode_sessions(test_n_sessions, vocab_data), dtype=np.int32)
    X_test_attack = np.array(_encode_sessions(test_a_sessions, vocab_data), dtype=np.int32)

    X_test = np.concatenate([X_test_normal, X_test_attack], axis=0)
    y_test = np.array(
        [0] * len(X_test_normal) + [1] * len(X_test_attack),
        dtype=np.int32,
    )

    # Save
    save_arrays(output_dir, prefix, X_train, X_test, y_test)

    # CSIC 2010 has no scaler JSON to carry window metadata, so the
    # vocab file doubles as the metadata carrier: wrap the token
    # mapping alongside window_size instead of saving the bare dict.
    vocab_out = {"window_size": window_size, "vocab": vocab_data}
    vocab_path = output_dir / f"vocab_{prefix}.json"
    with open(vocab_path, "w") as fh:
        json.dump(vocab_out, fh, indent=2)
    log.info("  Vocab saved → %s", vocab_path.name)

    unk_in_attack = int((X_test_attack == 1).sum())
    log.info("  UNK tokens in attack test: %d / %d", unk_in_attack, X_test_attack.size)

    return {
        "dataset": "csic2010",
        "window": window_size,
        "vocab_size": len(vocab_data),
        "X_train": X_train.shape,
        "X_test": X_test.shape,
        "y_test": y_test.shape,
        "unk_in_attack": unk_in_attack,
    }


# DATASET 2 — CIC-IDS2018

# Features most relevant to web/API attack detection
_CICIDS_FEATURES = [
    "Tot Fwd Pkts",
    "Tot Bwd Pkts",
    "TotLen Fwd Pkts",
    "TotLen Bwd Pkts",
    "Pkt Len Mean",
    "Pkt Len Std",
    "Pkt Len Max",
    "Pkt Size Avg",
    "Flow Duration",
    "Flow IAT Mean",
    "Flow IAT Std",
    "Fwd Pkts/s",
    "Bwd Pkts/s",
    "Flow Byts/s",
    "Flow Pkts/s",
    "SYN Flag Cnt",
    "RST Flag Cnt",
    "PSH Flag Cnt",
    "ACK Flag Cnt",
    "Init Fwd Win Byts",
    "Init Bwd Win Byts",
    "Dst Port",
    "Protocol",
]

# Days that contain web-relevant attacks
_CICIDS_WEB_FILES = [
    "02-22-2018.csv",  # Brute Force Web, XSS, SQL Injection
    "02-23-2018.csv",  # Brute Force Web, XSS, SQL Injection (continued)
]


def preprocess_cicids2018(
    data_dir: str | Path,
    output_dir: str | Path,
    window_size: int = 5,
    train_ratio: float = 0.8,
    features: list[str] | None = None,
) -> dict:
    """
    Full preprocessing pipeline for CIC-IDS2018.

    Uses only the two CSV days containing web-based attacks (SQL Injection, XSS, Brute Force Web).

    Expected files in data_dir
    --------------------------
    02-22-2018.csv
    02-23-2018.csv

    Output files in output_dir
    --------------------------
    X_train_cicids2018_w{window_size}.npy   (n_sessions, window_size, n_features)  float32
    X_test_cicids2018_w{window_size}.npy    (n_sessions, window_size, n_features)  float32
    y_test_cicids2018_w{window_size}.npy    (n_sessions,)                          int32
    scaler_cicids2018_w{window_size}.json

    Filenames are suffixed with the window size so outputs from
    different --window runs coexist on disk instead of overwriting
    each other, letting you compare results across window sizes.

    Returns
    -------
    dict with shapes and feature count
    """
    data_dir = Path(data_dir)
    output_dir = Path(output_dir)
    feat_cols = features or _CICIDS_FEATURES
    prefix = f"cicids2018_w{window_size}"

    # Load and combine web-attack days
    log.info("CIC-IDS2018 — loading CSV files...")
    frames = []
    for fname in _CICIDS_WEB_FILES:
        fpath = data_dir / fname
        if not fpath.exists():
            log.warning("  File not found, skipping: %s", fname)
            continue
        log.info("  Reading %s...", fname)
        frames.append(pd.read_csv(fpath, low_memory=False, encoding="utf-8"))

    if not frames:
        raise FileNotFoundError(f"No CIC-IDS2018 web-attack files found in {data_dir}")

    df = pd.concat(frames, ignore_index=True)

    # Drop repeated header rows (known CIC artifact)
    df = df[df["Label"] != "Label"].reset_index(drop=True)

    log.info("  Combined shape: %s", df.shape)
    log.info("  Label counts:\n%s", df["Label"].value_counts().to_string())

    # Binary label
    df["label"] = df["Label"].apply(lambda x: "normal" if x == "Benign" else "attack")

    # Fix data types — CIC CSVs have mixed-type columns
    log.info("  Cleaning feature columns...")
    for col in feat_cols:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")

    df[feat_cols] = (
        df[feat_cols]
        .replace([np.inf, -np.inf], np.nan)
        .fillna(df[feat_cols].median(numeric_only=True))
    )

    # Sort by timestamp (best-effort)
    if "Timestamp" in df.columns:
        df["Timestamp"] = pd.to_datetime(
            df["Timestamp"], dayfirst=True, errors="coerce"
        )

        # Drop known corrupted rows: CIC-IDS2018 has a handful of rows
        # with an epoch-fallback timestamp (parses to well before the
        # capture period). Verified via verify_filter_fix.py — these
        # are not just unparseable (NaT) but successfully parse to a
        # bogus pre-2018 date, so they'd otherwise sort to the very
        # front of the sequence and contaminate the start of training.
        n_before = len(df)
        valid_mask = df["Timestamp"] >= "2018-01-01"
        n_dropped = int((~valid_mask).sum())
        if n_dropped:
            log.warning(
                "  Dropping %d row(s) with corrupted epoch-fallback timestamp "
                "(< 2018-01-01)", n_dropped
            )
        df = df[valid_mask].reset_index(drop=True)
        log.info("  Rows after timestamp filter: %d (was %d)", len(df), n_before)

        df = df.sort_values("Timestamp").reset_index(drop=True)

    # Split normal / attack
    df_normal = df[df.label == "normal"].reset_index(drop=True)
    df_attack = df[df.label == "attack"].reset_index(drop=True)
    log.info("  Normal flows: %d  Attack flows: %d", len(df_normal), len(df_attack))

    # Scale — fit ONLY on normal traffic
    log.info("  Fitting scaler on normal traffic...")
    scaler = StandardScaler()
    X_normal = scaler.fit_transform(df_normal[feat_cols].values)
    X_attack = scaler.transform(df_attack[feat_cols].values)

    # Save scaler parameters
    output_dir.mkdir(parents=True, exist_ok=True)
    scaler_params = {
        "window": window_size,
        "mean": scaler.mean_.tolist(),
        "scale": scaler.scale_.tolist(),
        "feature_names": feat_cols,
    }
    scaler_path = output_dir / f"scaler_{prefix}.json"
    with open(scaler_path, "w") as fh:
        json.dump(scaler_params, fh, indent=2)
    log.info("  Scaler saved → %s", scaler_path.name)

    # Build sliding-window sessions
    log.info("  Building sessions (window=%d)...", window_size)
    split_idx = int(len(X_normal) * train_ratio)

    train_sessions = build_flow_sessions(X_normal[:split_idx], window_size)
    test_normal_sessions = build_flow_sessions(X_normal[split_idx:], window_size)
    test_attack_sessions = build_flow_sessions(X_attack, window_size)

    X_test = np.concatenate([test_normal_sessions, test_attack_sessions], axis=0)
    y_test = np.array(
        [0] * len(test_normal_sessions) + [1] * len(test_attack_sessions),
        dtype=np.int32,
    )

    # Save
    save_arrays(output_dir, prefix, train_sessions, X_test, y_test)

    return {
        "dataset": "cicids2018",
        "window": window_size,
        "n_features": len(feat_cols),
        "X_train": train_sessions.shape,
        "X_test": X_test.shape,
        "y_test": y_test.shape,
        "normal_flows": len(df_normal),
        "attack_flows": len(df_attack),
    }


# DATASET 3 — UNSW-NB15

# Features selected for network anomaly detection
# UNSW-NB15 has 49 features — these 20 are the most informative
# based on the dataset's own feature importance analysis
_UNSW_FEATURES = [
    # Basic flow
    "dur",
    "spkts",
    "dpkts",
    "sbytes",
    "dbytes",
    # Rate features
    "rate",
    "sload",
    "dload",
    # Packet stats
    "smean",
    "dmean",
    # Connection behavior
    "ct_state_ttl",
    "ct_dst_ltm",
    "ct_src_dport_ltm",
    "ct_dst_sport_ltm",
    "ct_dst_src_ltm",
    # Service and protocol (encoded)
    "proto",
    "service",
    "state",
    # Time features
    "sinpkt",
    "dinpkt",
]

# UNSW-NB15 normal label varies by file — handle both conventions
_UNSW_NORMAL_LABELS = {"Normal", "normal", "0", 0}


def _encode_categorical(df: pd.DataFrame, cols: list[str]) -> pd.DataFrame:
    """Label-encode categorical columns in place."""
    for col in cols:
        if col in df.columns and df[col].dtype == object:
            df[col] = pd.Categorical(df[col]).codes.astype(np.float32)
    return df


def preprocess_unsw(
    data_dir: str | Path,
    output_dir: str | Path,
    window_size: int = 5,
    train_ratio: float = 0.8,
    features: list[str] | None = None,
    sample_n: int | None = 500_000,
) -> dict:
    """
    Full preprocessing pipeline for UNSW-NB15.

    Accepts either the pre-split training/testing CSVs or the four raw partition files (UNSW-NB15_1.csv … _4.csv).

    Expected files in data_dir (any of these layouts work)
    -------------------------------------------------------
    Layout A — pre-split:
        UNSW_NB15_training-set.csv
        UNSW_NB15_testing-set.csv

    Layout B — raw partitions:
        UNSW-NB15_1.csv  UNSW-NB15_2.csv
        UNSW-NB15_3.csv  UNSW-NB15_4.csv

    Output files in output_dir
    --------------------------
    X_train_unsw_w{window_size}.npy    (n_sessions, window_size, n_features)  float32
    X_test_unsw_w{window_size}.npy     (n_sessions, window_size, n_features)  float32
    y_test_unsw_w{window_size}.npy     (n_sessions,)                          int32
    scaler_unsw_w{window_size}.json

    Filenames are suffixed with the window size so outputs from
    different --window runs coexist on disk instead of overwriting
    each other, letting you compare results across window sizes.

    Parameters
    ----------
    sample_n : if not None, randomly sample this many normal rows to keep memory usage manageable on constrained hardware.

    Returns
    -------
    dict with shapes and feature count
    """
    data_dir = Path(data_dir)
    output_dir = Path(output_dir)
    feat_cols = features or _UNSW_FEATURES
    prefix = f"unsw_w{window_size}"

    # Detect layout and load
    log.info("UNSW-NB15 — detecting file layout in %s...", data_dir)

    # Check both the given dir and a "Training and Testing Sets" subfolder
    subdir = data_dir / "Training and Testing Sets"
    search_dir = subdir if subdir.exists() else data_dir

    train_path = search_dir / "UNSW_NB15_training-set.csv"
    test_path = search_dir / "UNSW_NB15_testing-set.csv"

    # Only treat numeric-suffixed files as partition files (exclude LIST_EVENTS etc.)
    part_files = sorted(
        f for f in data_dir.glob("UNSW-NB15_*.csv") if f.stem.split("_")[-1].isdigit()
    )

    if train_path.exists() and test_path.exists():
        log.info("  Layout A detected (pre-split CSVs)")
        df_tr = pd.read_csv(train_path, low_memory=False)
        df_te = pd.read_csv(test_path, low_memory=False)
        df = pd.concat([df_tr, df_te], ignore_index=True)

    elif part_files:
        log.info("  Layout B detected (%d partition files)", len(part_files))
        frames = []
        for pf in part_files:
            log.info("  Reading %s...", pf.name)
            frames.append(
                pd.read_csv(pf, low_memory=False, encoding="latin-1", header=None)
            )
        df = pd.concat(frames, ignore_index=True)

        # Raw partition files have no header — assign from feature list
        # (49 features + label columns; keep only what we need later)
        log.info("  Raw partition files loaded — no header row")

    else:
        raise FileNotFoundError(
            f"No UNSW-NB15 files found in {data_dir}. "
            "Expected UNSW_NB15_training-set.csv / testing-set.csv "
            "or UNSW-NB15_1.csv … _4.csv"
        )

    log.info("  Raw shape: %s", df.shape)
    log.info("  Columns: %s ...", df.columns.tolist()[:10])

    # Identify label column
    label_col = None
    for candidate in ["label", "Label", "attack_cat", "class"]:
        if candidate in df.columns:
            label_col = candidate
            break
    if label_col is None:
        # Last column is typically the label in raw partition files
        label_col = df.columns[-1]
        log.warning(
            "  Label column not found by name — using last column: %s", label_col
        )

    log.info("  Label column: '%s'", label_col)
    log.info(
        "  Label distribution:\n%s", df[label_col].value_counts().head(12).to_string()
    )

    # Binary label
    df["label"] = df[label_col].apply(
        lambda x: "normal" if x in _UNSW_NORMAL_LABELS else "attack"
    )

    # Keep only available feature columns
    available = [f for f in feat_cols if f in df.columns]
    missing = [f for f in feat_cols if f not in df.columns]
    if missing:
        log.warning("  Features not found (will be skipped): %s", missing)
    feat_cols = available
    log.info("  Using %d features", len(feat_cols))

    # Encode categoricals, cast to numeric
    cat_cols = ["proto", "service", "state"]
    df = _encode_categorical(df, cat_cols)

    for col in feat_cols:
        df[col] = pd.to_numeric(df[col], errors="coerce")

    df[feat_cols] = (
        df[feat_cols]
        .replace([np.inf, -np.inf], np.nan)
        .fillna(df[feat_cols].median(numeric_only=True))
    )

    # Split normal / attack
    df_normal = df[df.label == "normal"].reset_index(drop=True)
    df_attack = df[df.label == "attack"].reset_index(drop=True)
    log.info("  Normal: %d  Attack: %d", len(df_normal), len(df_attack))

    # Optional subsampling to fit on constrained hardware
    if sample_n and len(df_normal) > sample_n:
        log.info("  Subsampling normal to %d rows (hardware limit)", sample_n)
        df_normal = df_normal.sample(n=sample_n, random_state=42).reset_index(drop=True)

    # Scale — fit ONLY on normal traffic
    log.info("  Fitting scaler on normal traffic...")
    scaler = StandardScaler()
    X_normal = scaler.fit_transform(df_normal[feat_cols].values)
    X_attack = scaler.transform(df_attack[feat_cols].values)

    output_dir.mkdir(parents=True, exist_ok=True)
    scaler_params = {
        "window": window_size,
        "mean": scaler.mean_.tolist(),
        "scale": scaler.scale_.tolist(),
        "feature_names": feat_cols,
    }
    scaler_path = output_dir / f"scaler_{prefix}.json"
    with open(scaler_path, "w") as fh:
        json.dump(scaler_params, fh, indent=2)
    log.info("  Scaler saved → %s", scaler_path.name)

    # Build sliding-window sessions
    log.info("  Building sessions (window=%d)...", window_size)
    split_idx = int(len(X_normal) * train_ratio)

    train_sessions = build_flow_sessions(X_normal[:split_idx], window_size)
    test_normal_sessions = build_flow_sessions(X_normal[split_idx:], window_size)
    test_attack_sessions = build_flow_sessions(X_attack, window_size)

    X_test = np.concatenate([test_normal_sessions, test_attack_sessions], axis=0)
    y_test = np.array(
        [0] * len(test_normal_sessions) + [1] * len(test_attack_sessions),
        dtype=np.int32,
    )

    # Save
    save_arrays(output_dir, prefix, train_sessions, X_test, y_test)

    return {
        "dataset": "unsw-nb15",
        "window": window_size,
        "n_features": len(feat_cols),
        "X_train": train_sessions.shape,
        "X_test": X_test.shape,
        "y_test": y_test.shape,
        "normal_flows": len(df_normal),
        "attack_flows": len(df_attack),
    }


# CLI — run all three pipelines in sequence

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Run preprocessing pipelines for API anomaly detection datasets"
    )
    parser.add_argument(
        "--base", default=".", help="Base project directory (default: current dir)"
    )
    parser.add_argument(
        "--window", type=int, default=5, help="Sliding window size (default: 5)"
    )
    parser.add_argument(
        "--dataset",
        choices=["csic2010", "cicids2018", "unsw", "all"],
        default="all",
        help="Which dataset to process (default: all)",
    )
    args = parser.parse_args()

    base = Path(args.base)
    raw_dir = base / "data" / "raw"
    output_dir = base / "data" / "processed"

    results = {}

    if args.dataset in ("csic2010", "all"):
        log.info("=" * 60)
        results["csic2010"] = preprocess_csic2010(
            data_dir=raw_dir / "csic2010",
            output_dir=output_dir,
            window_size=args.window,
        )

    if args.dataset in ("cicids2018", "all"):
        log.info("=" * 60)
        results["cicids2018"] = preprocess_cicids2018(
            data_dir=raw_dir / "cicids2018full",
            output_dir=output_dir,
            window_size=args.window,
        )

    if args.dataset in ("unsw", "all"):
        log.info("=" * 60)
        results["unsw"] = preprocess_unsw(
            data_dir=raw_dir / "unswnb15/Training and Testing Sets",
            output_dir=output_dir,
            window_size=args.window,
        )

    log.info("=" * 60)
    log.info("ALL DONE — summary:")
    for name, info in results.items():
        log.info("  %-12s  train=%s  test=%s", name, info["X_train"], info["X_test"])