from __future__ import annotations
import argparse
import sys
import time
from pathlib import Path
import numpy as np
import pandas as pd
REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO / 'Packet_CFM'))
from packet_store import PacketShardStore
SHAFIR5_FEATURE_NAMES = ('HTTPS', 'Protocol_Type', 'Magnitude', 'Variance', 'fin_count')

def _compute_batch(tokens: np.ndarray, lens: np.ndarray, dst_ports: np.ndarray, protocols: np.ndarray) -> np.ndarray:
    (B, T, _) = tokens.shape
    out = np.zeros((B, 5), dtype=np.float32)
    arange = np.arange(T)[None, :]
    mask = arange < lens[:, None]
    log_size = tokens[:, :, 0]
    sizes = np.expm1(np.maximum(log_size, 0.0))
    sizes = np.where(mask, sizes, 0.0)
    n = lens.astype(np.float32)
    n_safe = np.maximum(n, 1.0)
    sum_sq = (sizes * sizes).sum(axis=1)
    mean = sizes.sum(axis=1) / n_safe
    mean_sq = sum_sq / n_safe
    magnitude = np.sqrt(np.maximum(mean_sq, 0.0))
    variance = np.maximum(mean_sq - mean * mean, 0.0)
    fin_flags = tokens[:, :, 4]
    fin_flags = np.where(mask, fin_flags, 0.0)
    fin_count = fin_flags.sum(axis=1)
    https = (dst_ports == 443).astype(np.float32)
    proto_type = protocols.astype(np.float32)
    out[:, 0] = https
    out[:, 1] = proto_type
    out[:, 2] = magnitude
    out[:, 3] = variance
    out[:, 4] = fin_count
    return out

def main():
    p = argparse.ArgumentParser()
    p.add_argument('--source-store', type=Path, required=True)
    p.add_argument('--flows-parquet', type=Path, required=True)
    p.add_argument('--out', type=Path, required=True)
    p.add_argument('--T', type=int, default=None, help='Truncate to first T packets (default = stored).')
    p.add_argument('--batch', type=int, default=100000)
    args = p.parse_args()
    print(f'[read] {args.flows_parquet}')
    flows = pd.read_parquet(args.flows_parquet, columns=['flow_id', 'label', 'dst_port', 'protocol'])
    flow_id = flows['flow_id'].to_numpy(dtype=np.uint64)
    labels = flows['label'].astype(str).to_numpy()
    dst_ports = flows['dst_port'].to_numpy(dtype=np.uint32)
    protocols = flows['protocol'].to_numpy(dtype=np.uint8)
    store = PacketShardStore.open(args.source_store)
    store_fid = store.read_flows(columns=['flow_id'])['flow_id'].to_numpy(dtype=np.uint64)
    if len(store_fid) != len(flow_id) or not np.array_equal(store_fid, flow_id):
        raise ValueError('store flow_id ordering differs from flows.parquet')
    T_stored = int(store.manifest['packet_length'].max())
    T = args.T if args.T is not None else T_stored
    n = len(flows)
    feats = np.zeros((n, 5), dtype=np.float32)
    print(f'[stream] {n:,} flows × T={T} (stored {T_stored}), batch={args.batch}')
    t0 = time.time()
    for start in range(0, n, args.batch):
        end = min(start + args.batch, n)
        idx = np.arange(start, end, dtype=np.int64)
        (tok, ll) = store.read_packets(idx, T=T)
        ll = np.minimum(ll, T).astype(np.int32)
        feats[start:end] = _compute_batch(tok.astype(np.float32), ll, dst_ports[start:end], protocols[start:end])
        if start // args.batch % 20 == 0 or end == n:
            dt = time.time() - t0
            rate = end / max(dt, 1e-06)
            eta = (n - end) / max(rate, 1.0)
            print(f'[stream] {end:,}/{n:,}  dt={dt:.1f}s  rate={rate:.0f} flows/s  ETA={eta:.0f}s', flush=True)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    df = pd.DataFrame({'flow_id': flow_id, 'label': labels})
    for (i, name) in enumerate(SHAFIR5_FEATURE_NAMES):
        df[name] = feats[:, i]
    df.to_parquet(args.out, compression='snappy', index=False)
    print(f'[write] {args.out}  rows={len(df):,}  cols={list(df.columns)}')
    print(f'[stats] HTTPS=1 fraction: {(feats[:, 0] > 0).mean():.4f}')
    print(f'[stats] Protocol_Type unique values: {np.unique(feats[:, 1].astype(int))[:10]}')
    print(f'[stats] Magnitude mean={feats[:, 2].mean():.1f} median={np.median(feats[:, 2]):.1f}')
    print(f'[stats] Variance mean={feats[:, 3].mean():.1f}')
    print(f'[stats] fin_count mean={feats[:, 4].mean():.3f}')
if __name__ == '__main__':
    main()
