from __future__ import annotations
import argparse
import sys
import zipfile
from pathlib import Path
from typing import BinaryIO
import numpy as np
import pandas as pd
from numpy.lib import format as npy_format
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from packet_store import PacketShardStore, PacketShardWriter

def _read_npy_header(fp: BinaryIO) -> tuple[tuple[int, ...], np.dtype, bool]:
    version = npy_format.read_magic(fp)
    if version == (1, 0):
        (shape, fortran_order, dtype) = npy_format.read_array_header_1_0(fp)
    elif version == (2, 0):
        (shape, fortran_order, dtype) = npy_format.read_array_header_2_0(fp)
    else:
        raise ValueError(f'unsupported npy version {version}')
    return (tuple((int(v) for v in shape)), np.dtype(dtype), bool(fortran_order))

def _read_exact(fp: BinaryIO, n_bytes: int) -> bytes:
    chunks: list[bytes] = []
    remaining = int(n_bytes)
    while remaining:
        chunk = fp.read(remaining)
        if not chunk:
            raise EOFError(f'expected {n_bytes} bytes, missing {remaining}')
        chunks.append(chunk)
        remaining -= len(chunk)
    return b''.join(chunks)

def _open_member(zf: zipfile.ZipFile, name: str) -> tuple[BinaryIO, tuple[int, ...], np.dtype]:
    fp = zf.open(name)
    (shape, dtype, fortran_order) = _read_npy_header(fp)
    if fortran_order:
        fp.close()
        raise ValueError(f'{name} uses Fortran order, expected C order')
    return (fp, shape, dtype)

def _iter_npz_rows(npz_path: Path, rows: int, chunk_rows: int):
    with zipfile.ZipFile(npz_path) as zf:
        (token_fp, token_shape, token_dtype) = _open_member(zf, 'packet_tokens.npy')
        (length_fp, length_shape, length_dtype) = _open_member(zf, 'packet_lengths.npy')
        try:
            if len(token_shape) != 3:
                raise ValueError(f'packet_tokens.npy must be 3-D, got {token_shape}')
            if length_shape != (token_shape[0],):
                raise ValueError(f'packet_lengths.npy shape {length_shape} does not match tokens {token_shape}')
            if rows > token_shape[0]:
                raise ValueError(f'requested {rows} rows, but {npz_path} has {token_shape[0]}')
            row_values = int(np.prod(token_shape[1:], dtype=np.int64))
            token_row_bytes = row_values * token_dtype.itemsize
            length_row_bytes = length_dtype.itemsize
            emitted = 0
            while emitted < rows:
                take = min(int(chunk_rows), rows - emitted)
                token_bytes = _read_exact(token_fp, take * token_row_bytes)
                length_bytes = _read_exact(length_fp, take * length_row_bytes)
                tokens = np.frombuffer(token_bytes, dtype=token_dtype).reshape(take, token_shape[1], token_shape[2])
                lengths = np.frombuffer(length_bytes, dtype=length_dtype).reshape(take)
                yield (emitted, tokens, lengths)
                emitted += take
        finally:
            token_fp.close()
            length_fp.close()

def _npz_token_shape(npz_path: Path) -> tuple[int, int, int]:
    with zipfile.ZipFile(npz_path) as zf:
        (fp, shape, _dtype) = _open_member(zf, 'packet_tokens.npy')
        fp.close()
    if len(shape) != 3:
        raise ValueError(f'packet_tokens.npy must be 3-D, got {shape}')
    return shape

def convert(args: argparse.Namespace) -> None:
    pairs = list(zip(args.packets_npz, args.flows_parquet, strict=True))
    first_shape = _npz_token_shape(pairs[0][0])
    total_rows = 0
    with PacketShardWriter(args.out_store, shard_size=args.shard_size, T_full=first_shape[1], D=first_shape[2], overwrite=args.overwrite) as writer:
        for (split_id, (npz_path, flows_path)) in enumerate(pairs):
            token_shape = _npz_token_shape(npz_path)
            if token_shape[1:] != first_shape[1:]:
                raise ValueError(f'{npz_path} shape {token_shape} does not match {first_shape}')
            flows = pd.read_parquet(flows_path)
            rows = min(len(flows), token_shape[0])
            if args.max_rows_per_split > 0:
                rows = min(rows, args.max_rows_per_split)
            if len(flows) != token_shape[0]:
                raise ValueError(f'{flows_path} has {len(flows)} rows but {npz_path} has {token_shape[0]}')
            print(f'[split {split_id}] npz={npz_path} flows={flows_path} rows={rows:,} shape={token_shape}', flush=True)
            for (start, tokens, lengths) in _iter_npz_rows(npz_path, rows, args.chunk_rows):
                end = start + len(lengths)
                writer.add_batch(tokens, lengths, flows.iloc[start:end].reset_index(drop=True))
                total_rows += len(lengths)
                if total_rows % args.report_every < len(lengths) or end == rows:
                    print(f'[split {split_id}] emitted={end:,}/{rows:,} total={total_rows:,}', flush=True)
    store = PacketShardStore.open(args.out_store)
    flows = store.read_flows(columns=['label'])
    print(f"[done] store={args.out_store} rows={store.n_flows:,} shards={store.metadata['n_shards']}")
    print(flows['label'].value_counts().to_string())

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--packets-npz', type=Path, nargs='+', required=True)
    parser.add_argument('--flows-parquet', type=Path, nargs='+', required=True)
    parser.add_argument('--out-store', type=Path, required=True)
    parser.add_argument('--shard-size', type=int, default=50000)
    parser.add_argument('--chunk-rows', type=int, default=10000)
    parser.add_argument('--report-every', type=int, default=250000)
    parser.add_argument('--max-rows-per-split', type=int, default=0)
    parser.add_argument('--overwrite', action='store_true')
    args = parser.parse_args()
    if len(args.packets_npz) != len(args.flows_parquet):
        raise SystemExit('--packets-npz and --flows-parquet must have the same count')
    convert(args)
if __name__ == '__main__':
    main()
