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, compute_flow_features_from_packets

def _from_npz(args: argparse.Namespace) -> tuple[np.ndarray, 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)
    packet_flow_id = pz['flow_id'] if 'flow_id' in pz.files else None
    T_stored = tokens.shape[1]
    if args.T is not None:
        if args.T > T_stored:
            raise ValueError(f'requested T={args.T} > stored {T_stored}')
        tokens = tokens[:, :args.T, :]
        lens = np.minimum(lens, args.T).astype(np.int32)
    print(f'[read] {args.flows_parquet}')
    flows = pd.read_parquet(args.flows_parquet, columns=['flow_id', 'label'])
    if len(flows) != len(tokens):
        raise ValueError(f'row count mismatch: packets={len(tokens):,} flows={len(flows):,}')
    flow_id = np.asarray(flows['flow_id'].to_numpy(), dtype=np.uint64)
    if packet_flow_id is not None:
        if not np.array_equal(flow_id, packet_flow_id.astype(np.uint64)):
            raise ValueError('packets.npz flow_id != flows.parquet flow_id')
    labels = flows['label'].astype(str).to_numpy()
    print(f'[compute] {len(tokens):,} flows × T={tokens.shape[1]} → {len(CANONICAL_FLOW_FEATURE_NAMES)} features ...')
    t0 = time.time()
    feats = compute_flow_features_from_packets(tokens, lens)
    dt = time.time() - t0
    print(f'[compute] {dt:.1f}s  ({len(tokens) / max(dt, 1e-06):.0f} flows/s)')
    return (feats, flow_id, labels, np.array([T_stored if args.T is None else args.T]))

def _from_store(args: argparse.Namespace) -> tuple[np.ndarray, 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}')
    print(f'[read] {args.flows_parquet}')
    flows = pd.read_parquet(args.flows_parquet, columns=['flow_id', 'label'])
    n = len(flows)
    store_flows = store.read_flows(columns=['flow_id'])
    if len(store_flows) != n:
        raise ValueError(f'store has {len(store_flows):,} rows but flows.parquet has {n:,}')
    if not np.array_equal(store_flows['flow_id'].to_numpy(dtype=np.uint64), flows['flow_id'].to_numpy(dtype=np.uint64)):
        raise ValueError('store flow_id ordering differs from flows.parquet')
    flow_id = flows['flow_id'].to_numpy(dtype=np.uint64)
    labels = flows['label'].astype(str).to_numpy()
    feats = np.zeros((n, len(CANONICAL_FLOW_FEATURE_NAMES)), 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)
        feats[start:end] = compute_flow_features_from_packets(tok, lens)
        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 (feats, flow_id, labels, np.array([T]))

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)
    p.add_argument('--out', type=Path, required=True)
    p.add_argument('--T', type=int, default=None, help='Truncate packet sequences to first T positions (default: use stored T_full).')
    p.add_argument('--batch', type=int, default=100000, help='Batch size when streaming from --source-store.')
    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')
    if args.packets_npz is not None:
        (feats, flow_id, labels, _) = _from_npz(args)
    else:
        (feats, flow_id, labels, _) = _from_store(args)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    df = pd.DataFrame({'flow_id': flow_id, 'label': labels})
    for (i, name) in enumerate(CANONICAL_FLOW_FEATURE_NAMES):
        df[name] = feats[:, i]
    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(df):,} rows × {len(df.columns)} cols)')
if __name__ == '__main__':
    main()
