from __future__ import annotations
import argparse
import json
import sys as _sys
import time
from pathlib import Path
from pathlib import Path as _Path
import numpy as np
import torch
import yaml
from sklearn.metrics import average_precision_score, roc_auc_score
_sys.path.insert(0, str(_Path(__file__).resolve().parent))
from data import load_mixed_data
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: MixedTokenCFM, flow_np: np.ndarray, cont_np: np.ndarray, disc_np: np.ndarray, len_np: np.ndarray, device: torch.device, *, batch_size: int, n_steps: int) -> dict[str, np.ndarray]:
    out: dict[str, list[np.ndarray]] = {}
    for start in range(0, len(flow_np), batch_size):
        sl = slice(start, start + batch_size)
        flow = torch.from_numpy(flow_np[sl]).float().to(device)
        cont = torch.from_numpy(cont_np[sl]).float().to(device)
        disc = torch.from_numpy(disc_np[sl]).long().to(device)
        lens = torch.from_numpy(len_np[sl]).long().to(device)
        with torch.no_grad():
            traj = model.trajectory_metrics(flow, cont, disc, lens, n_steps=n_steps)
            nll = model.disc_nll_score(flow, cont, disc, lens)
        for d in (traj, nll):
            for (k, v) in d.items():
                out.setdefault(k, []).append(v.detach().cpu().numpy())
        print(f'[score] {min(start + batch_size, len(flow_np)):,}/{len(flow_np):,}', flush=True)
    return {k: np.concatenate(v, axis=0) for (k, v) in out.items()}

def _auroc_safe(y, s) -> float:
    try:
        return float(roc_auc_score(y, s))
    except ValueError:
        return float('nan')

def _auprc_safe(y, s) -> float:
    try:
        return float(average_precision_score(y, s))
    except ValueError:
        return float('nan')

def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument('--model-dir', type=Path, required=True)
    p.add_argument('--out-dir', type=Path, required=True)
    p.add_argument('--n-val-cap', type=int, default=None)
    p.add_argument('--n-atk-cap', type=int, default=None)
    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()
    device = _device(args.device)
    args.out_dir.mkdir(parents=True, exist_ok=True)
    cfg = yaml.safe_load((args.model_dir / 'config.yaml').read_text())
    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()
    print(f'[model] T={model_cfg.T} flow_dim={model_cfg.flow_dim}')
    data = load_mixed_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']), 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)), attack_cap=int(cfg['attack_cap']) if cfg.get('attack_cap') else None, val_cap=int(cfg['val_cap']) if cfg.get('val_cap') else None)
    print(f'[data] val={len(data.val_flow):,} attack={len(data.attack_flow):,}')
    rng = np.random.default_rng(0)
    (val_flow, val_cont, val_disc, val_len) = (data.val_flow, data.val_cont, data.val_disc, data.val_len)
    (atk_flow, atk_cont, atk_disc, atk_len) = (data.attack_flow, data.attack_cont, data.attack_disc, data.attack_len)
    atk_labels = data.attack_labels
    if args.n_val_cap is not None and len(val_flow) > args.n_val_cap:
        idx = np.sort(rng.choice(len(val_flow), size=args.n_val_cap, replace=False))
        (val_flow, val_cont, val_disc, val_len) = (val_flow[idx], val_cont[idx], val_disc[idx], val_len[idx])
    if args.n_atk_cap is not None and len(atk_flow) > args.n_atk_cap:
        idx = np.sort(rng.choice(len(atk_flow), size=args.n_atk_cap, replace=False))
        (atk_flow, atk_cont, atk_disc, atk_len) = (atk_flow[idx], atk_cont[idx], atk_disc[idx], atk_len[idx])
        atk_labels = atk_labels[idx]
    print(f'[eval] scoring val={len(val_flow):,} atk={len(atk_flow):,}')
    t0 = time.time()
    val = _score_batch(model, val_flow, val_cont, val_disc, val_len, device, batch_size=args.batch_size, n_steps=args.n_steps)
    print(f'[eval] val done {time.time() - t0:.1f}s')
    t0 = time.time()
    atk = _score_batch(model, atk_flow, atk_cont, atk_disc, atk_len, device, batch_size=args.batch_size, n_steps=args.n_steps)
    print(f'[eval] atk done {time.time() - t0:.1f}s')
    keys = sorted(set(val) & set(atk))
    overall: dict[str, dict[str, float]] = {}
    for k in keys:
        y = np.r_[np.zeros(len(val[k])), np.ones(len(atk[k]))]
        s = np.r_[val[k], atk[k]]
        s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
        overall[k] = {'auroc': _auroc_safe(y, s), 'auprc': _auprc_safe(y, s)}
    per_class: dict[str, dict[str, float]] = {}
    for c in sorted(set(atk_labels.tolist())):
        m = atk_labels == c
        per_class[c] = {'_n': float(m.sum())}
        for k in keys:
            y = np.r_[np.zeros(len(val[k])), np.ones(int(m.sum()))]
            s = np.r_[val[k], atk[k][m]]
            s = np.nan_to_num(s, nan=0.0, posinf=1000000000000.0, neginf=-1000000000000.0)
            per_class[c][k] = _auroc_safe(y, s)
    np.savez(args.out_dir / 'phase1_scores.npz', val_labels=np.array(['normal'] * len(val_flow)), atk_labels=atk_labels.astype(str), **{f'val_{k}': val[k] for k in keys}, **{f'atk_{k}': atk[k] for k in keys})
    json.dump({'overall': overall, 'per_class': per_class}, open(args.out_dir / 'phase1_summary.json', 'w'), indent=2)
    print(f'[wrote] {args.out_dir}/phase1_summary.json   keys={len(keys)}')
if __name__ == '__main__':
    main()
