from __future__ import annotations
import argparse
import json
import os
import sys
import time
import warnings
from glob import glob
from pathlib import Path
import numpy as np
import pandas as pd
os.environ.setdefault('JAX_PLATFORMS', 'cpu')
warnings.filterwarnings('ignore')
import optax
from pzflow import Flow
from sklearn.metrics import average_precision_score, roc_auc_score
from sklearn.preprocessing import StandardScaler
REPO = Path(__file__).resolve().parents[2]
IDS2017_FEATURES = ['Flow Duration', 'Total Fwd Packets', 'Total Backward Packets', 'Total Length of Fwd Packets', 'Total Length of Bwd Packets', 'Fwd Packet Length Min', 'Fwd Packet Length Mean', 'Fwd Packet Length Std', 'Bwd Packet Length Max', 'Bwd Packet Length Min', 'Bwd Packet Length Mean', 'Flow IAT Mean', 'Flow IAT Std', 'Flow IAT Max', 'Flow IAT Min', 'Fwd IAT Total', 'Fwd IAT Mean', 'Fwd IAT Std', 'Fwd IAT Max', 'Fwd IAT Min', 'Bwd IAT Total', 'Bwd IAT Mean', 'Bwd IAT Std', 'Bwd IAT Max', 'Bwd IAT Min', 'Fwd Header Length', 'Bwd Header Length', 'Fwd Packets/s', 'Bwd Packets/s', 'Min Packet Length', 'Max Packet Length', 'Packet Length Mean', 'Packet Length Std', 'Packet Length Variance', 'SYN Flag Count', 'PSH Flag Count', 'ACK Flag Count', 'URG Flag Count', 'Down/Up Ratio', 'Average Packet Size', 'Avg Fwd Segment Size', 'Avg Bwd Segment Size', 'Subflow Fwd Packets', 'Subflow Fwd Bytes', 'Subflow Bwd Packets', 'Subflow Bwd Bytes', 'Init_Win_bytes_forward', 'Init_Win_bytes_backward', 'act_data_pkt_fwd', 'min_seg_size_forward', 'Active Mean', 'Active Std', 'Active Max', 'Active Min', 'Idle Mean', 'Idle Std', 'Idle Max', 'Idle Min']
TOR2016_FEATURES = ['Protocol', 'Flow Duration', 'Flow Bytes/s', 'Flow Packets/s', 'Flow IAT Mean', 'Flow IAT Std', 'Flow IAT Max', 'Flow IAT Min', 'Fwd IAT Mean', 'Fwd IAT Std', 'Fwd IAT Max', 'Fwd IAT Min', 'Bwd IAT Mean', 'Bwd IAT Std', 'Bwd IAT Max', 'Bwd IAT Min', 'Active Mean', 'Active Std', 'Active Max', 'Active Min', 'Idle Mean', 'Idle Std', 'Idle Max', 'Idle Min']
CICIOT5_FEATURES = ['HTTPS', 'Protocol Type', 'Magnitude', 'Variance', 'fin_count']
CICIDS_BEST5_FEATURES = ['Bwd Packet Length Mean', 'Fwd Packets/s', 'ACK Flag Count', 'Total Length of Bwd Packets', 'Flow Duration']
TOR_BEST4_FEATURES = ['Flow IAT Std', 'Flow Bytes/s', 'Flow Packets/s', 'Bwd IAT Max']
COLUMN_ALIASES = {'Total Fwd Packets': ['Total Fwd Packet'], 'Total Backward Packets': ['Total Bwd packets'], 'Total Length of Fwd Packets': ['Total Length of Fwd Packet'], 'Total Length of Bwd Packets': ['Total Length of Bwd Packet'], 'Fwd Header Length': ['Fwd Header Length.1'], 'Init_Win_bytes_forward': ['FWD Init Win Bytes', 'Init Win Bytes Fwd'], 'Init_Win_bytes_backward': ['Bwd Init Win Bytes', 'Init Win Bytes Bwd'], 'act_data_pkt_fwd': ['Fwd Act Data Pkts'], 'min_seg_size_forward': ['Fwd Seg Size Min'], 'Avg Fwd Segment Size': ['Fwd Segment Size Avg'], 'Avg Bwd Segment Size': ['Bwd Segment Size Avg'], 'Min Packet Length': ['Packet Length Min'], 'Max Packet Length': ['Packet Length Max']}
DATASETS = {'iscxtor': {'csv_glob': str(REPO / 'datasets/iscxtor2016/raw/csv/Scenario-*-merged_5s.csv'), 'label_col': 'label', 'benign_values': ['nonTOR'], 'drop_patterns': [], 'feature_set': TOR_BEST4_FEATURES}, 'cicids2017': {'csv_glob': str(REPO / 'datasets/cicids2017/raw/csv/*.csv'), 'label_col': 'Label', 'benign_values': ['BENIGN', 'Benign', 'benign'], 'drop_patterns': [' - Attempted', '- Attempted'], 'feature_set': CICIDS_BEST5_FEATURES}, 'cicddos2019': {'csv_glob': str(REPO / 'datasets/cicddos2019/raw/csv/**/*.csv'), 'label_col': 'Label', 'benign_values': ['BENIGN', 'Benign', 'benign'], 'drop_patterns': [], 'feature_set': CICIDS_BEST5_FEATURES}, 'ciciot2023': {'csv_glob': str(REPO / 'datasets/ciciot2023/raw/csv/CSV/*/*.pcap.csv'), 'label_col': None, 'benign_folder': 'Benign_Final', 'drop_patterns': [], 'feature_set': CICIOT5_FEATURES}}
PROTOCOL_CONFIG = {'iscxtor_within': ('iscxtor', 'iscxtor', {'n_train': 10000, 'n_val': 10000, 'n_attack': None}), 'cicids_within': ('cicids2017', 'cicids2017', {'n_train': 10000, 'n_val': 10000, 'n_attack': 30000}), 'cicddos_within': ('cicddos2019', 'cicddos2019', {'n_train': 10000, 'n_val': 10000, 'n_attack': 20000}), 'ciciot_within': ('ciciot2023', 'ciciot2023', {'n_train': 10000, 'n_val': 10000, 'n_attack': 30000}), 'forward_cross': ('cicids2017', 'cicddos2019', {'n_train': 10000, 'n_val': 10000, 'n_attack': 10000}), 'reverse_cross': ('cicddos2019', 'cicids2017', {'n_train': 10000, 'n_val': 10000, 'n_attack': 10000})}

def _resolve_columns(df: pd.DataFrame, names: list[str]) -> tuple[list[str], list[str]]:
    df.columns = [c.strip() if isinstance(c, str) else c for c in df.columns]
    (resolved, missing) = ([], [])
    for n in names:
        if n in df.columns:
            resolved.append(n)
            continue
        found = None
        for alias in COLUMN_ALIASES.get(n, []):
            if alias in df.columns:
                found = alias
                break
        if found is None:
            low = {c.lower(): c for c in df.columns}
            if n.lower() in low:
                found = low[n.lower()]
        if found is None:
            missing.append(n)
        else:
            resolved.append(found)
    return (resolved, missing)

def _load_csvs(dataset_name: str, return_paths: bool=False):
    cfg = DATASETS[dataset_name]
    paths = sorted(glob(cfg['csv_glob'], recursive=True))
    if not paths:
        raise FileNotFoundError(f"no CSVs match {cfg['csv_glob']}")
    print(f'  [csv] {dataset_name}: {len(paths)} files')
    return paths if return_paths else paths

def _attach_labels(df: pd.DataFrame, dataset_name: str, source_path: str | None=None) -> pd.DataFrame:
    cfg = DATASETS[dataset_name]
    if cfg.get('label_col') is None:
        folder = Path(source_path).parent.name
        df = df.copy()
        df['cls_label'] = folder
        df['binary_label'] = 0 if folder == cfg['benign_folder'] else 1
    else:
        lbl_col = cfg['label_col'].strip()
        match = None
        for c in df.columns:
            if isinstance(c, str) and c.strip() == lbl_col:
                match = c
                break
        if match is None:
            raise KeyError(f'label column {lbl_col!r} not found in {source_path}')
        df = df.copy()
        df['cls_label'] = df[match].astype(str).str.strip()
        for pat in cfg['drop_patterns']:
            df = df[~df['cls_label'].str.contains(pat, na=False, regex=False)]
        df['binary_label'] = df['cls_label'].apply(lambda x: 0 if x in cfg['benign_values'] else 1)
    return df

def _load_dataset(dataset_name: str, feature_set: list[str]) -> pd.DataFrame:
    cfg = DATASETS[dataset_name]
    paths = _load_csvs(dataset_name)
    dfs = []
    for p in paths:
        try:
            df = pd.read_csv(p, low_memory=False)
        except Exception as e:
            print(f'  [csv-warn] skip {p}: {e}')
            continue
        df = _attach_labels(df, dataset_name, source_path=p)
        (resolved, missing) = _resolve_columns(df, feature_set)
        if missing:
            if not hasattr(_load_dataset, '_warned'):
                _load_dataset._warned = set()
            key = (dataset_name, tuple(missing))
            if key not in _load_dataset._warned:
                _load_dataset._warned.add(key)
                print(f'  [warn] {Path(p).name}: missing {missing}')
        sub = df[resolved + ['binary_label', 'cls_label']].copy()
        rename = {r: n for (r, n) in zip(resolved, [f for f in feature_set if f not in missing])}
        sub = sub.rename(columns=rename)
        dfs.append(sub)
    if not dfs:
        raise RuntimeError(f'no usable CSVs for {dataset_name}')
    full = pd.concat(dfs, axis=0, ignore_index=True)
    for c in [c for c in feature_set if c in full.columns]:
        full[c] = pd.to_numeric(full[c], errors='coerce')
    full = full.replace([np.inf, -np.inf], np.nan)
    feat_cols = [c for c in feature_set if c in full.columns]
    full = full.dropna(subset=feat_cols).reset_index(drop=True)
    print(f'  [csv] {dataset_name} concat: {len(full):,} rows  benign={int((full.binary_label == 0).sum()):,}  attack={int((full.binary_label == 1).sum()):,}  features_kept={len(feat_cols)}')
    return (full, feat_cols)

def _sample_within(df: pd.DataFrame, caps: dict, seed: int):
    rng = np.random.default_rng(seed)
    benign = df[df.binary_label == 0]
    attack = df[df.binary_label == 1]
    n_train = caps['n_train']
    n_val = caps['n_val']
    n_atk = caps['n_attack']
    needed_b = n_train + n_val
    if len(benign) < needed_b:
        raise RuntimeError(f'only {len(benign)} benign rows, need {needed_b}')
    b_idx = rng.permutation(len(benign))
    train = benign.iloc[b_idx[:n_train]]
    val = benign.iloc[b_idx[n_train:n_train + n_val]]
    if n_atk is None:
        atk = attack
    else:
        atk_classes = sorted(attack['cls_label'].unique())
        per = max(1, n_atk // len(atk_classes))
        chunks = []
        for cls in atk_classes:
            pool = attack[attack['cls_label'] == cls]
            k = min(per, len(pool))
            if k:
                chunks.append(pool.sample(n=k, random_state=seed))
        atk = pd.concat(chunks, axis=0, ignore_index=True)
        if len(atk) > n_atk:
            atk = atk.sample(n=n_atk, random_state=seed)
    return (train, val, atk)

def _sample_cross(src_df, tgt_df, caps, seed):
    rng = np.random.default_rng(seed + 1000)
    src_benign = src_df[src_df.binary_label == 0]
    if len(src_benign) < caps['n_train']:
        raise RuntimeError(f"src benign only {len(src_benign)}, need {caps['n_train']}")
    sb_idx = rng.permutation(len(src_benign))
    train = src_benign.iloc[sb_idx[:caps['n_train']]]
    rng2 = np.random.default_rng(seed)
    tgt_benign = tgt_df[tgt_df.binary_label == 0]
    tgt_attack = tgt_df[tgt_df.binary_label == 1]
    if len(tgt_benign) < caps['n_val']:
        raise RuntimeError(f'tgt benign only {len(tgt_benign)}')
    tb_idx = rng2.permutation(len(tgt_benign))
    val = tgt_benign.iloc[tb_idx[:caps['n_val']]]
    atk_classes = sorted(tgt_attack['cls_label'].unique())
    per = max(1, caps['n_attack'] // len(atk_classes))
    chunks = []
    for cls in atk_classes:
        pool = tgt_attack[tgt_attack['cls_label'] == cls]
        k = min(per, len(pool))
        if k:
            chunks.append(pool.sample(n=k, random_state=seed))
    atk = pd.concat(chunks, axis=0, ignore_index=True)
    if len(atk) > caps['n_attack']:
        atk = atk.sample(n=caps['n_attack'], random_state=seed)
    return (train, val, atk)

def _safe_metric(fn, y, s) -> float:
    s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
    try:
        return float(fn(y, s))
    except ValueError:
        return float('nan')

def _train_and_score(train, val, atk, feat_cols, *, epochs, lr, optimizer):
    raw_train = train[feat_cols].astype(np.float64).values
    keep = raw_train.std(axis=0) > 0
    if not keep.all():
        dropped = [c for (c, k) in zip(feat_cols, keep) if not k]
        print(f'  [train] dropping {len(dropped)} zero-variance cols: {dropped}')
        feat_cols = [c for (c, k) in zip(feat_cols, keep) if k]
        raw_train = raw_train[:, keep]
    raw_val = val[feat_cols].astype(np.float64).values
    raw_atk = atk[feat_cols].astype(np.float64).values
    scaler = StandardScaler()
    X_train = scaler.fit_transform(raw_train)
    X_val = scaler.transform(raw_val)
    X_atk = scaler.transform(raw_atk)
    clip_lim = 30.0
    X_train = np.clip(X_train, -clip_lim, clip_lim)
    X_val = np.clip(X_val, -clip_lim, clip_lim)
    X_atk = np.clip(X_atk, -clip_lim, clip_lim)
    df_train = pd.DataFrame(X_train.astype(np.float32), columns=[f'x{i}' for i in range(len(feat_cols))])
    df_val = pd.DataFrame(X_val.astype(np.float32), columns=df_train.columns)
    df_atk = pd.DataFrame(X_atk.astype(np.float32), columns=df_train.columns)
    if optimizer == 'sgd':
        opt = optax.sgd(learning_rate=lr)
    else:
        opt = optax.adam(learning_rate=lr)
    flow = Flow(df_train.columns.tolist())
    t0 = time.time()
    losses = flow.train(df_train, optimizer=opt, epochs=epochs, verbose=False)
    t_train = time.time() - t0
    t0 = time.time()
    lp_val = np.asarray(flow.log_prob(df_val))
    lp_atk = np.asarray(flow.log_prob(df_atk))
    t_score = time.time() - t0
    return {'score_val': (-lp_val).astype(np.float32), 'score_atk': (-lp_atk).astype(np.float32), 'losses': np.asarray(losses, dtype=np.float64), 't_train': t_train, 't_score': t_score}

def _per_class(val_score, atk_score, atk_labels):
    out = {}
    for cls in sorted(set(atk_labels)):
        m = atk_labels == cls
        n_c = int(m.sum())
        v_c = atk_score[m]
        y = np.r_[np.zeros(len(val_score)), np.ones(len(v_c))]
        s = np.r_[val_score, v_c]
        out[cls] = {'_n': float(n_c), 'auroc': _safe_metric(roc_auc_score, y, s)}
    return out

def main():
    p = argparse.ArgumentParser()
    p.add_argument('--protocol', required=True, choices=list(PROTOCOL_CONFIG))
    p.add_argument('--seed', type=int, required=True, choices=[42, 43, 44])
    p.add_argument('--out-dir', type=Path, required=True)
    p.add_argument('--epochs', type=int, default=100)
    p.add_argument('--lr', type=float, default=0.001)
    p.add_argument('--optimizer', choices=['sgd', 'adam'], default='sgd')
    args = p.parse_args()
    args.out_dir.mkdir(parents=True, exist_ok=True)
    (src_name, tgt_name, caps) = PROTOCOL_CONFIG[args.protocol]
    cross = src_name != tgt_name
    print(f'[run] shafir_nf_csv protocol={args.protocol} seed={args.seed}')
    print(f'      src={src_name}  tgt={tgt_name}  cross={cross}')
    feat_set = DATASETS[src_name]['feature_set']
    (src_df, src_feat_cols) = _load_dataset(src_name, feat_set)
    if cross:
        (tgt_df, tgt_feat_cols) = _load_dataset(tgt_name, feat_set)
        feat_cols = [c for c in feat_set if c in src_feat_cols and c in tgt_feat_cols]
        print(f'  [features] cross intersection: {len(feat_cols)} cols')
        (train, val, atk) = _sample_cross(src_df, tgt_df, caps, args.seed)
    else:
        feat_cols = src_feat_cols
        print(f'  [features] within: {len(feat_cols)} cols')
        (train, val, atk) = _sample_within(src_df, caps, args.seed)
    print(f'  [data] train={len(train):,}  val={len(val):,}  attack={len(atk):,}  D={len(feat_cols)}')
    res = _train_and_score(train, val, atk, feat_cols, epochs=args.epochs, lr=args.lr, optimizer=args.optimizer)
    (val_score, atk_score) = (res['score_val'], res['score_atk'])
    y = np.r_[np.zeros(len(val_score)), np.ones(len(atk_score))]
    s = np.r_[val_score, atk_score]
    overall = {'neg_log_prob': {'auroc': _safe_metric(roc_auc_score, y, s), 'auprc': _safe_metric(average_precision_score, y, s)}}
    a_labels = atk['cls_label'].astype(str).to_numpy()
    per_cls = _per_class(val_score, atk_score, a_labels)
    out = {'method': 'shafir_nf_csv', 'protocol': args.protocol, 'seed': args.seed, 'src_dataset': src_name, 'tgt_dataset': tgt_name, 'feature_set': feat_cols, 'n_features': len(feat_cols), 'n_train': len(train), 'n_val': len(val), 'n_atk': len(atk), 'epochs': args.epochs, 'lr': args.lr, 'optimizer': args.optimizer, 't_train_sec': round(res['t_train'], 2), 't_score_sec': round(res['t_score'], 2), 'loss_first_last': [float(res['losses'][0]), float(res['losses'][-1])], 'overall': overall, 'per_class': per_cls}
    out_json = args.out_dir / f'{args.protocol}_seed{args.seed}.json'
    out_json.write_text(json.dumps(out, indent=2))
    npz_path = out_json.with_suffix('.npz')
    np.savez_compressed(npz_path, b_neg_log_prob=val_score, a_neg_log_prob=atk_score, a_labels=a_labels.astype(str), losses=res['losses'])
    print(f'[saved] {out_json}')
    print(f"[result] AUROC={overall['neg_log_prob']['auroc']:.4f} AUPRC={overall['neg_log_prob']['auprc']:.4f} train={res['t_train']:.1f}s")
if __name__ == '__main__':
    main()
