from __future__ import annotations
import argparse
from pathlib import Path
import numpy as np
import pandas as pd
LABEL_ALIASES = {'UDP-lag': 'UDPLag'}

def _infer_flows_path(packets_path: Path) -> Path:
    name = packets_path.name
    if name.startswith('packets.'):
        flows_name = 'flows.' + name[len('packets.'):].removesuffix('.npz') + '.parquet'
    else:
        raise ValueError(f'Cannot infer flows path from {packets_path}')
    return packets_path.parent / flows_name

def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--in', dest='inputs', action='append', type=Path, required=True, help='packets.<shard>.npz path. Pass multiple times (one per shard). flows.<shard>.parquet is inferred.')
    ap.add_argument('--out-packets', type=Path, required=True)
    ap.add_argument('--out-flows', type=Path, required=True)
    args = ap.parse_args()
    tok_chunks: list[np.ndarray] = []
    len_chunks: list[np.ndarray] = []
    flow_dfs: list[pd.DataFrame] = []
    META_COLS = ['flow_id', 'label', 'start_ts', 'src_ip', 'dst_ip', 'src_port', 'dst_port', 'protocol', 'n_pkts']
    for pkt_path in args.inputs:
        flow_path = _infer_flows_path(pkt_path)
        if not pkt_path.exists():
            raise FileNotFoundError(pkt_path)
        if not flow_path.exists():
            raise FileNotFoundError(flow_path)
        p = np.load(pkt_path)
        available = set(pd.read_parquet(flow_path).columns)
        cols = [c for c in META_COLS if c in available]
        f = pd.read_parquet(flow_path, columns=cols)
        if len(p['flow_id']) != len(f):
            raise ValueError(f'{pkt_path.name}: row count mismatch with {flow_path.name}')
        if not np.array_equal(p['flow_id'], f['flow_id'].to_numpy()):
            raise ValueError(f'{pkt_path.name}: flow_id mismatch with {flow_path.name}')
        tok_chunks.append(np.asarray(p['packet_tokens']))
        len_chunks.append(np.asarray(p['packet_lengths']))
        flow_dfs.append(f)
        print(f"[load] {pkt_path.name} : {len(p['flow_id']):>10,} rows  cols={cols}")
    T_full_set = {t.shape[1] for t in tok_chunks}
    D_set = {t.shape[2] for t in tok_chunks}
    if len(T_full_set) != 1 or len(D_set) != 1:
        raise ValueError(f'inconsistent T/D across shards: T={T_full_set} D={D_set}')
    tokens = np.concatenate(tok_chunks, axis=0)
    lengths = np.concatenate(len_chunks, axis=0)
    flow_df = pd.concat(flow_dfs, ignore_index=True)
    del tok_chunks, len_chunks, flow_dfs
    if LABEL_ALIASES:
        flow_df['label'] = flow_df['label'].map(lambda s: LABEL_ALIASES.get(s, s)).astype(str)
    N = len(tokens)
    flow_id = np.arange(N, dtype=np.uint64)
    flow_df['flow_id'] = flow_id
    labels = flow_df['label'].to_numpy().astype(str)
    print(f'\n[merge] total rows: {N:,}')
    print(f'[merge] label distribution:')
    (ulabels, counts) = np.unique(labels, return_counts=True)
    for (lbl, cnt) in sorted(zip(ulabels, counts), key=lambda x: -x[1]):
        print(f'    {lbl:<40s}  {cnt:>10,}')
    args.out_packets.parent.mkdir(parents=True, exist_ok=True)
    np.savez(args.out_packets, packet_tokens=tokens, packet_lengths=lengths, flow_id=flow_id)
    print(f'\n[merge] wrote {args.out_packets} ({args.out_packets.stat().st_size / 1000000000.0:.2f} GB)')
    args.out_flows.parent.mkdir(parents=True, exist_ok=True)
    flow_df.to_parquet(args.out_flows, compression='snappy', index=False)
    print(f'[merge] wrote {args.out_flows} ({args.out_flows.stat().st_size / 1000000.0:.2f} MB)  cols={list(flow_df.columns)}')
if __name__ == '__main__':
    main()
