from __future__ import annotations
import argparse
import sys
import time
from pathlib import Path
import numpy as np
import pandas as pd
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from common.data_contract import CANONICAL_FLOW_FEATURE_NAMES

def compute_spectral_features(packet_tokens: np.ndarray, packet_lengths: np.ndarray, n_bands: int=8) -> np.ndarray:
    (N, T, _) = packet_tokens.shape
    mask = (np.arange(T)[None, :] < packet_lengths[:, None]).astype(np.float32)
    sig = packet_tokens[..., :2].astype(np.float32) * mask[..., None]
    Z = np.fft.rfft(sig, axis=1)
    if n_bands > Z.shape[1]:
        raise ValueError(f'n_bands={n_bands} > {Z.shape[1]} available bins')
    Z_K = Z[:, :n_bands]
    size_re = Z_K[..., 0].real.astype(np.float32)
    size_im = Z_K[..., 0].imag.astype(np.float32)
    iat_re = Z_K[..., 1].real.astype(np.float32)
    iat_im = Z_K[..., 1].imag.astype(np.float32)
    out = np.concatenate([size_re, size_im, iat_re, iat_im], axis=1)
    return out

def _spectral_column_names(n_bands: int) -> list[str]:
    cols: list[str] = []
    for prefix in ('spec_size_re', 'spec_size_im', 'spec_iat_re', 'spec_iat_im'):
        for k in range(n_bands):
            cols.append(f'{prefix}_K{k}')
    return cols

def _from_npz(args: argparse.Namespace) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    print(f'[read] {args.packets_npz}')
    pz = np.load(args.packets_npz)
    tokens = pz['packet_tokens']
    lens = pz['packet_lengths'].astype(np.int32)
    if args.T is not None:
        if args.T > tokens.shape[1]:
            raise ValueError(f'requested T={args.T} > stored {tokens.shape[1]}')
        tokens = tokens[:, :args.T, :]
        lens = np.minimum(lens, args.T).astype(np.int32)
    flow_id = pz['flow_id'].astype(np.uint64) if 'flow_id' in pz.files else None
    print(f'[compute] {len(tokens):,} flows × T={tokens.shape[1]} → {4 * args.n_bands} spectral cols ...')
    t0 = time.time()
    spec = compute_spectral_features(tokens, lens, n_bands=args.n_bands)
    print(f'[compute] {time.time() - t0:.1f}s')
    return (spec, flow_id, lens)

def _from_store(args: argparse.Namespace) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'Packet_CFM'))
    from packet_store import PacketShardStore
    store = PacketShardStore.open(args.source_store)
    T_stored = int(store.manifest['packet_length'].max())
    T = args.T if args.T is not None else T_stored
    if T > T_stored:
        raise ValueError(f'requested T={T} > stored max {T_stored}')
    store_flows = store.read_flows(columns=['flow_id'])
    n = len(store_flows)
    flow_id = store_flows['flow_id'].to_numpy(dtype=np.uint64)
    spec = np.zeros((n, 4 * args.n_bands), dtype=np.float32)
    print(f'[stream] {n:,} flows × T={T} (full={T_stored}), batch={args.batch} ...')
    t0 = time.time()
    all_idx = np.arange(n, dtype=np.int64)
    for start in range(0, n, args.batch):
        end = min(start + args.batch, n)
        idx = all_idx[start:end]
        (tok, lens) = store.read_packets(idx, T=T)
        lens = np.minimum(lens, T).astype(np.int32)
        spec[start:end] = compute_spectral_features(tok, lens, n_bands=args.n_bands)
        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)
    return (spec, flow_id, None)

def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--packets-npz', type=Path, default=None, help='Monolithic packets.npz path (mutually exclusive with --source-store).')
    p.add_argument('--source-store', type=Path, default=None, help='PacketShardStore directory (mutually exclusive with --packets-npz).')
    p.add_argument('--flows-parquet', type=Path, required=True, help='flows.parquet for flow_id + label.')
    p.add_argument('--base-features', type=Path, required=True, help='Existing canonical flow_features.parquet (20-d).')
    p.add_argument('--out', type=Path, required=True)
    p.add_argument('--n-bands', type=int, default=8)
    p.add_argument('--T', type=int, default=None, help='Truncate to first T packets before FFT (default: stored T).')
    p.add_argument('--batch', type=int, default=100000)
    args = p.parse_args()
    if (args.packets_npz is None) == (args.source_store is None):
        p.error('pass exactly one of --packets-npz or --source-store')
    print(f'[read] {args.flows_parquet}')
    flows = pd.read_parquet(args.flows_parquet, columns=['flow_id', 'label'])
    n = len(flows)
    flow_id_flows = flows['flow_id'].to_numpy(dtype=np.uint64)
    labels = flows['label'].astype(str).to_numpy()
    print(f'[read] {args.base_features}')
    base = pd.read_parquet(args.base_features)
    if len(base) != n:
        raise ValueError(f'base features rows {len(base):,} != flows rows {n:,}')
    if 'flow_id' in base.columns:
        if not np.array_equal(base['flow_id'].to_numpy(dtype=np.uint64), flow_id_flows):
            raise ValueError('base flow_id != flows flow_id (row alignment broken)')
    if args.packets_npz is not None:
        (spec, flow_id_pkt, _) = _from_npz(args)
    else:
        (spec, flow_id_pkt, _) = _from_store(args)
    if flow_id_pkt is not None and (not np.array_equal(flow_id_pkt, flow_id_flows)):
        raise ValueError('packet flow_id != flows flow_id')
    out_df = pd.DataFrame({'flow_id': flow_id_flows, 'label': labels})
    for name in CANONICAL_FLOW_FEATURE_NAMES:
        if name not in base.columns:
            raise ValueError(f'base parquet missing canonical feature {name!r}')
        out_df[name] = base[name].to_numpy(dtype=np.float32)
    spec_cols = _spectral_column_names(args.n_bands)
    for (i, name) in enumerate(spec_cols):
        out_df[name] = spec[:, i]
    args.out.parent.mkdir(parents=True, exist_ok=True)
    out_df.to_parquet(args.out, compression='snappy', index=False)
    sz_mb = args.out.stat().st_size / 1000000.0
    print(f'[write] {args.out}  ({sz_mb:.2f} MB, {len(out_df):,} rows × {len(out_df.columns)} cols, +{4 * args.n_bands} spectral)')
if __name__ == '__main__':
    main()
