from __future__ import annotations
import argparse
import json
import sys
import time
from collections import defaultdict
from pathlib import Path
import numpy as np
import pandas as pd
import yaml
if not hasattr(np, 'Inf'):
    np.Inf = np.inf
from sklearn.metrics import average_precision_score, roc_auc_score
REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO / 'baselines/Kitsune-py'))
sys.path.insert(0, str(REPO / 'Unified_CFM'))
from FeatureExtractor import FE
from KitNET.KitNET import KitNET
from data import load_unified_data
PCAP_GLOBS = {'iscxtor': str(REPO / 'datasets/iscxtor2016/raw/pcap_extracted/**/*.pcap'), 'cicids2017': str(REPO / 'datasets/cicids2017/raw/pcap/*.pcap'), 'cicddos2019': str(REPO / 'datasets/cicddos2019/raw/pcap/*')}
WITHIN_DIRS = {'iscxtor_within': ('phase25_multiseed_2026_04_25/iscxtor2016_lambda0p3_seed{seed}', 'iscxtor', {'n_val': 10000, 'n_atk': None}), 'cicids_within': ('phase25_sigma06_multiseed_2026_04_25/cicids2017_lambda0p3_sigma0p6_seed{seed}', 'cicids2017', {'n_val': 10000, 'n_atk': 30000}), 'cicddos_within': ('phase25_multiseed_2026_04_25/cicddos2019_lambda0p3_seed{seed}', 'cicddos2019', {'n_val': 10000, 'n_atk': 20000})}

def _canonical_key(src_ip, dst_ip, src_port, dst_port, protocol) -> tuple:
    a = (src_ip, src_port)
    b = (dst_ip, dst_port)
    if a <= b:
        return (a[0], b[0], a[1], b[1], int(protocol))
    return (b[0], a[0], b[1], a[1], int(protocol))

def _proto_from_kitsune(srcproto: str, dstproto: str) -> int:
    if srcproto == 'icmp':
        return 1
    if srcproto == 'arp':
        return 0
    return -1

class FEWithMeta(FE):

    def __init__(self, path, limit=np.inf):
        super().__init__(path, limit)
        self._last_ts = None
        self._last_5tuple = None
        self._last_framelen = None

    def get_next_vector(self):
        if self.curPacketIndx == self.limit:
            if self.parse_type == 'tsv':
                self.tsvinf.close()
            return []
        if self.parse_type == 'tsv':
            row = self.tsvin.__next__()
            IPtype = np.nan
            timestamp = row[0]
            framelen = row[1]
            srcIP = ''
            dstIP = ''
            if row[4] != '':
                (srcIP, dstIP, IPtype) = (row[4], row[5], 0)
            elif row[17] != '':
                (srcIP, dstIP, IPtype) = (row[17], row[18], 1)
            srcproto = row[6] + row[8]
            dstproto = row[7] + row[9]
            (srcMAC, dstMAC) = (row[2], row[3])
            if srcproto == '':
                if row[12] != '':
                    (srcproto, dstproto) = ('arp', 'arp')
                    (srcIP, dstIP, IPtype) = (row[14], row[16], 0)
                elif row[10] != '':
                    (srcproto, dstproto, IPtype) = ('icmp', 'icmp', 0)
                elif srcIP + srcproto + dstIP + dstproto == '':
                    (srcIP, dstIP) = (row[2], row[3])
        else:
            return []
        try:
            sp = int(srcproto) if srcproto.isdigit() else 0
            dp = int(dstproto) if dstproto.isdigit() else 0
        except Exception:
            (sp, dp) = (0, 0)
        try:
            self._last_ts = float(timestamp)
        except Exception:
            self._last_ts = np.nan
        self._last_5tuple = (srcIP, dstIP, sp, dp)
        try:
            self._last_framelen = int(framelen)
        except Exception:
            self._last_framelen = 0
        self.curPacketIndx += 1
        try:
            return self.nstat.updateGetStats(IPtype, srcMAC, dstMAC, srcIP, srcproto, dstIP, dstproto, int(framelen), float(timestamp))
        except Exception as e:
            print(f'  [warn] netStat error: {e}')
            return []

def _stream_pcap_kitsune(pcap_path: Path, *, kit: KitNET, fm_grace: int, ad_grace: int, packet_limit: int, fivetuple_to_rmses: dict, n_packets_total: list) -> None:
    print(f'  [stream] {pcap_path.name}', flush=True)
    fe = FEWithMeta(str(pcap_path), limit=packet_limit)
    t0 = time.time()
    n_local = 0
    while True:
        x = fe.get_next_vector()
        if len(x) == 0:
            break
        n_local += 1
        n_packets_total[0] += 1
        rmse = kit.process(x)
        if rmse is None or rmse == 0:
            continue
        if fe._last_5tuple is None:
            continue
        (srcIP, dstIP, sp, dp) = fe._last_5tuple
        key = (srcIP, dstIP, sp, dp) if (srcIP, sp) <= (dstIP, dp) else (dstIP, srcIP, dp, sp)
        fivetuple_to_rmses[key].append(rmse)
        if n_local % 200000 == 0:
            print(f'    [{n_local:,}] elapsed {time.time() - t0:.0f}s ({n_local / max(time.time() - t0, 0.001):.0f} pkt/s)', flush=True)
    print(f'  [stream] {pcap_path.name} done: {n_local:,} packets in {time.time() - t0:.0f}s', flush=True)

def _flows_to_key(flows_df: pd.DataFrame) -> np.ndarray:
    keys = []
    for (src_ip, dst_ip, sp, dp) in zip(flows_df['src_ip'], flows_df['dst_ip'], flows_df['src_port'], flows_df['dst_port']):
        if (str(src_ip), int(sp)) <= (str(dst_ip), int(dp)):
            k = (str(src_ip), str(dst_ip), int(sp), int(dp))
        else:
            k = (str(dst_ip), str(src_ip), int(dp), int(sp))
        keys.append(k)
    return np.asarray(keys, dtype=object)

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 main():
    p = argparse.ArgumentParser()
    p.add_argument('--protocol', required=True, choices=list(WITHIN_DIRS))
    p.add_argument('--seed', type=int, required=True)
    p.add_argument('--out-dir', type=Path, required=True)
    p.add_argument('--fm-grace', type=int, default=5000)
    p.add_argument('--ad-grace', type=int, default=50000)
    p.add_argument('--max-ae-size', type=int, default=10)
    p.add_argument('--lr', type=float, default=0.1)
    p.add_argument('--hidden-ratio', type=float, default=0.75)
    p.add_argument('--packet-limit-per-pcap', type=int, default=2000000, help='Cap per-pcap packets to keep runtime tractable. None = full.')
    p.add_argument('--max-pcaps', type=int, default=None, help='Cap number of pcap files processed (default: all).')
    args = p.parse_args()
    args.out_dir.mkdir(parents=True, exist_ok=True)
    (template, ds_name, caps) = WITHIN_DIRS[args.protocol]
    model_dir = REPO / 'artifacts' / template.format(seed=args.seed)
    print(f'[run] kitsune_path_a protocol={args.protocol} seed={args.seed}')
    print(f'[run] dataset={ds_name}  model_dir={model_dir}')
    cfg = yaml.safe_load((model_dir / 'config.yaml').read_text())
    data = load_unified_data(packets_npz=Path(cfg['packets_npz']) if cfg.get('packets_npz') else None, source_store=Path(cfg['source_store']) if cfg.get('source_store') else None, flows_parquet=Path(cfg['flows_parquet']), flow_features_path=Path(cfg['flow_features_path']) if cfg.get('flow_features_path') else None, flow_feature_columns=cfg.get('flow_feature_columns'), flow_features_align=str(cfg.get('flow_features_align', 'auto')), T=int(cfg['T']), split_seed=int(cfg.get('data_seed', cfg.get('seed', 42))), train_ratio=float(cfg.get('train_ratio', 0.8)), benign_label=str(cfg.get('benign_label', 'normal')), min_len=int(cfg.get('min_len', 2)), packet_preprocess=str(cfg.get('packet_preprocess', 'mixed_dequant')), attack_cap=int(cfg['attack_cap']) if cfg.get('attack_cap') else caps['n_atk'], val_cap=int(cfg['val_cap']) if cfg.get('val_cap') else caps['n_val'])
    flows_full = pd.read_parquet(cfg['flows_parquet'])
    print(f'[data] flows.parquet rows: {len(flows_full):,}; val={len(data.val_flow):,}  attack={len(data.attack_flow):,}')
    from glob import glob
    pcaps = sorted(glob(PCAP_GLOBS[ds_name], recursive=True))
    pcaps = [Path(p) for p in pcaps]
    if args.max_pcaps is not None:
        pcaps = pcaps[:args.max_pcaps]
    print(f'[pcap] discovered {len(pcaps)} pcap(s)')
    for p in pcaps[:5]:
        print(f'  {p}')
    if len(pcaps) > 5:
        print(f'  ...({len(pcaps) - 5} more)')
    kit = KitNET(n=100, max_autoencoder_size=args.max_ae_size, FM_grace_period=args.fm_grace, AD_grace_period=args.ad_grace, learning_rate=args.lr, hidden_ratio=args.hidden_ratio)
    fivetuple_to_rmses: dict = defaultdict(list)
    n_total = [0]
    t0 = time.time()
    for p in pcaps:
        _stream_pcap_kitsune(p, kit=kit, fm_grace=args.fm_grace, ad_grace=args.ad_grace, packet_limit=args.packet_limit_per_pcap, fivetuple_to_rmses=fivetuple_to_rmses, n_packets_total=n_total)
    elapsed = time.time() - t0
    print(f'[stream] total {n_total[0]:,} packets in {elapsed:.0f}s ({n_total[0] / max(elapsed, 0.001):.0f} pkt/s)')
    print(f'[stream] unique 5-tuples seen: {len(fivetuple_to_rmses):,}')
    keys_full = _flows_to_key(flows_full)
    print(f'[match] keying {len(keys_full):,} flows to 5-tuples')
    flow_score_mean = np.full(len(flows_full), np.nan, dtype=np.float64)
    flow_score_max = np.full(len(flows_full), np.nan, dtype=np.float64)
    flow_score_median = np.full(len(flows_full), np.nan, dtype=np.float64)
    n_matched = 0
    for (i, k) in enumerate(keys_full):
        rl = fivetuple_to_rmses.get(tuple(k))
        if rl:
            flow_score_mean[i] = float(np.mean(rl))
            flow_score_max[i] = float(np.max(rl))
            flow_score_median[i] = float(np.median(rl))
            n_matched += 1
    print(f'[match] flows with RMSE coverage: {n_matched:,}/{len(flows_full):,} ({100 * n_matched / max(len(flows_full), 1):.1f}%)')
    val_flow_ids = set((int(x) for x in data.val_flow_ids)) if hasattr(data, 'val_flow_ids') else None
    bin_labels = (flows_full['label'].astype(str) != cfg.get('benign_label', 'normal')).astype(int).to_numpy()
    keys = ['mean', 'max', 'median']
    score_arrs = {'mean': flow_score_mean, 'max': flow_score_max, 'median': flow_score_median}
    overall = {}
    for k in keys:
        s = score_arrs[k]
        valid = ~np.isnan(s)
        if valid.sum() < 10:
            overall[k] = {'auroc': float('nan'), 'auprc': float('nan'), 'n_valid': int(valid.sum())}
            continue
        y = bin_labels[valid]
        sv = s[valid]
        overall[k] = {'auroc': _safe_metric(roc_auc_score, y, sv), 'auprc': _safe_metric(average_precision_score, y, sv), 'n_valid': int(valid.sum())}
        print(f"  [{k}] AUROC={overall[k]['auroc']:.4f}  AUPRC={overall[k]['auprc']:.4f}  (n_valid={overall[k]['n_valid']:,})")
    out_json = args.out_dir / f'{args.protocol}_seed{args.seed}.json'
    out = {'method': 'kitsune_path_a', 'protocol': args.protocol, 'seed': args.seed, 'dataset': ds_name, 'n_pcaps': len(pcaps), 'n_total_packets': int(n_total[0]), 'n_unique_5tuples': int(len(fivetuple_to_rmses)), 'n_flows_total': int(len(flows_full)), 'n_flows_matched': int(n_matched), 'fm_grace': args.fm_grace, 'ad_grace': args.ad_grace, 'packet_limit_per_pcap': args.packet_limit_per_pcap, 'elapsed_sec': round(elapsed, 1), 'overall_by_agg': overall}
    out_json.write_text(json.dumps(out, indent=2))
    np.savez_compressed(out_json.with_suffix('.npz'), flow_score_mean=flow_score_mean.astype(np.float32), flow_score_max=flow_score_max.astype(np.float32), flow_score_median=flow_score_median.astype(np.float32), binary_label=bin_labels.astype(np.int8))
    print(f'[saved] {out_json}')
if __name__ == '__main__':
    main()
