from __future__ import annotations
import argparse
import json
import sys as _sys
import time
from pathlib import Path
import numpy as np
import pandas as pd
import torch
from sklearn.metrics import average_precision_score, roc_auc_score
REPO = Path(__file__).resolve().parents[1]
_sys.path.insert(0, str(REPO))
from common.data_contract import PACKET_CONTINUOUS_CHANNEL_IDX, PACKET_BINARY_CHANNEL_IDX, zscore as _zscore
from common.packet_store import PacketShardStore
_sys.path.insert(0, str(Path(__file__).resolve().parent))
from model import MixedCFMConfig, MixedTokenCFM

def _device(arg: str) -> torch.device:
    if arg == 'auto':
        return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    return torch.device(arg)

def _score_batch(model, flow_z, cont_z, disc_int, lens, device, batch_size=256, n_steps=16):
    out: dict[str, list[np.ndarray]] = {}
    for start in range(0, len(flow_z), batch_size):
        sl = slice(start, start + batch_size)
        f = torch.from_numpy(flow_z[sl]).float().to(device)
        c = torch.from_numpy(cont_z[sl]).float().to(device)
        d = torch.from_numpy(disc_int[sl]).long().to(device)
        l = torch.from_numpy(lens[sl]).long().to(device)
        with torch.no_grad():
            traj = model.trajectory_metrics(f, c, d, l, n_steps=n_steps)
            nll = model.disc_nll_score(f, c, d, l)
        for src in (traj, nll):
            for (k, v) in src.items():
                out.setdefault(k, []).append(v.detach().cpu().numpy())
        if start // batch_size % 20 == 0:
            print(f'[score] {min(start + batch_size, len(flow_z)):,}/{len(flow_z):,}', flush=True)
    return {k: np.concatenate(v) for (k, v) in out.items()}

def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument('--model-dir', type=Path, required=True)
    p.add_argument('--target-store', type=Path, default=None, help='Sharded packet store (mutually exclusive with --target-packets-npz).')
    p.add_argument('--target-packets-npz', type=Path, default=None, help='Monolithic packets.npz (for datasets without full_store).')
    p.add_argument('--target-flows', type=Path, required=True)
    p.add_argument('--target-flow-features', type=Path, required=True)
    p.add_argument('--out', type=Path, required=True)
    p.add_argument('--n-benign', type=int, default=10000)
    p.add_argument('--n-attack', type=int, default=10000)
    p.add_argument('--benign-label', type=str, default='normal', help="Label string of benign class in target dataset (e.g. 'nontor' for ISCXTor2016).")
    p.add_argument('--seed', type=int, default=42)
    p.add_argument('--T', type=int, default=64)
    p.add_argument('--batch-size', type=int, default=256)
    p.add_argument('--n-steps', type=int, default=16)
    p.add_argument('--device', type=str, default='auto')
    args = p.parse_args()
    if (args.target_store is None) == (args.target_packets_npz is None):
        p.error('pass exactly one of --target-store or --target-packets-npz')
    device = _device(args.device)
    ckpt = torch.load(args.model_dir / 'model.pt', map_location='cpu', weights_only=False)
    model_cfg = MixedCFMConfig(**ckpt['model_cfg'])
    model = MixedTokenCFM(model_cfg).to(device)
    model.load_state_dict(ckpt['model_state_dict'])
    model.eval()
    cont_mean = np.asarray(ckpt['cont_mean'], dtype=np.float32)
    cont_std = np.asarray(ckpt['cont_std'], dtype=np.float32)
    flow_mean = np.asarray(ckpt['flow_mean'], dtype=np.float32)
    flow_std = np.asarray(ckpt['flow_std'], dtype=np.float32)
    flow_names = [str(n) for n in ckpt['flow_feature_names']]
    print(f'[model] T={args.T} flow_dim={model_cfg.flow_dim}')
    flows = pd.read_parquet(args.target_flows, columns=['flow_id', 'label'])
    ff = pd.read_parquet(args.target_flow_features)
    if not np.array_equal(flows['flow_id'].to_numpy(dtype=np.uint64), ff['flow_id'].to_numpy(dtype=np.uint64)):
        raise ValueError('flows and flow_features not row-aligned')
    labels = flows['label'].astype(str).to_numpy()
    print(f'[data] {len(flows):,} target rows')
    rng = np.random.default_rng(args.seed)
    benign_idx = np.flatnonzero(labels == args.benign_label)
    attack_idx = np.flatnonzero(labels != args.benign_label)
    n_benign = min(args.n_benign, len(benign_idx))
    if n_benign < args.n_benign:
        print(f'[warn] only {len(benign_idx)} benign rows available (asked {args.n_benign})')
    b_sel = np.sort(rng.choice(benign_idx, size=n_benign, replace=False))
    atk_classes = sorted(set(labels[attack_idx]))
    per_class = max(1, args.n_attack // len(atk_classes))
    a_sel_chunks = []
    for cls in atk_classes:
        pool = attack_idx[labels[attack_idx] == cls]
        k = min(per_class, len(pool))
        if k:
            a_sel_chunks.append(rng.choice(pool, size=k, replace=False))
    a_sel = np.sort(np.concatenate(a_sel_chunks))
    if len(a_sel) > args.n_attack:
        a_sel = np.sort(rng.choice(a_sel, size=args.n_attack, replace=False))
    print(f'[sample] benign={len(b_sel):,} attack={len(a_sel):,} ({len(atk_classes)} classes)')
    cont_idx = list(PACKET_CONTINUOUS_CHANNEL_IDX)
    disc_idx = list(PACKET_BINARY_CHANNEL_IDX)
    if args.target_store is not None:
        store = PacketShardStore.open(args.target_store)
        npz_tokens = None
        npz_lens = None
    else:
        store = None
        pz = np.load(args.target_packets_npz)
        npz_tokens = pz['packet_tokens'][:, :args.T].astype(np.float32)
        npz_lens = np.minimum(pz['packet_lengths'], args.T).astype(np.int32)

    def _materialize(idx: np.ndarray):
        if store is not None:
            (tok, l) = store.read_packets(idx, T=args.T)
        else:
            tok = npz_tokens[idx]
            l = npz_lens[idx]
        l = np.minimum(l, args.T).astype(np.int32)
        f = ff.iloc[idx][flow_names].to_numpy(dtype=np.float64)
        f = np.nan_to_num(f, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
        return (tok.astype(np.float32), l, f)
    print('[read] benign...')
    (b_tok, b_len, b_flow) = _materialize(b_sel)
    print('[read] attack...')
    (a_tok, a_len, a_flow) = _materialize(a_sel)
    if cont_mean.shape == (9,):
        cm = cont_mean[cont_idx]
        cs = cont_std[cont_idx]
    else:
        cm = cont_mean
        cs = cont_std

    def _prep(tok, lens):
        cont = tok[..., cont_idx]
        disc = tok[..., disc_idx].astype(np.int8)
        z = _zscore(cont, cm, cs)
        m = np.arange(args.T)[None, :] < lens[:, None]
        cont_z = (z * m[:, :, None]).astype(np.float32)
        return (cont_z, disc)
    (b_cont, b_disc) = _prep(b_tok, b_len)
    (a_cont, a_disc) = _prep(a_tok, a_len)
    b_flow_z = ((b_flow - flow_mean) / np.maximum(flow_std, 1e-06)).astype(np.float32)
    a_flow_z = ((a_flow - flow_mean) / np.maximum(flow_std, 1e-06)).astype(np.float32)
    t0 = time.time()
    print('[eval] benign...')
    b_scores = _score_batch(model, b_flow_z, b_cont, b_disc, b_len, device, batch_size=args.batch_size, n_steps=args.n_steps)
    print(f'[eval] benign done {time.time() - t0:.1f}s')
    t0 = time.time()
    print('[eval] attack...')
    a_scores = _score_batch(model, a_flow_z, a_cont, a_disc, a_len, device, batch_size=args.batch_size, n_steps=args.n_steps)
    print(f'[eval] attack done {time.time() - t0:.1f}s')
    keys = sorted(set(b_scores) & set(a_scores))
    overall = {}
    for k in keys:
        y = np.r_[np.zeros(len(b_scores[k])), np.ones(len(a_scores[k]))]
        s = np.r_[b_scores[k], a_scores[k]]
        s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
        overall[k] = {'auroc': float(roc_auc_score(y, s)), 'auprc': float(average_precision_score(y, s))}
    a_labels = labels[a_sel]
    per_cls = {}
    for cls in sorted(set(a_labels)):
        m = a_labels == cls
        per_cls[cls] = {'_n': float(m.sum())}
        for k in keys:
            y = np.r_[np.zeros(len(b_scores[k])), np.ones(int(m.sum()))]
            s = np.r_[b_scores[k], a_scores[k][m]]
            s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
            try:
                per_cls[cls][k] = float(roc_auc_score(y, s))
            except ValueError:
                per_cls[cls][k] = float('nan')
    out = {'model_dir': str(args.model_dir), 'target_store': str(args.target_store), 'n_benign': len(b_sel), 'n_attack': len(a_sel), 'n_score_keys': len(keys), 'overall': overall, 'per_class': per_cls}
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(out, indent=2))
    npz = args.out.with_suffix('.npz')
    save = {'a_labels': a_labels.astype(str)}
    for k in keys:
        save[f'b_{k}'] = b_scores[k].astype(np.float32)
        save[f'a_{k}'] = a_scores[k].astype(np.float32)
    np.savez(npz, **save)
    print(f'[saved] {args.out}')
if __name__ == '__main__':
    main()
