"""Потоковый прогон большого бэга прямо из tar-архива.

`new_data` — это 90 ГБ в 221 шарде, и распаковывать его целиком некуда. Скрипт
достаёт по одному шарду во временный каталог, прогоняет конвейер, забирает
дескрипторы кандидатов и удаляет шард. Состояние конвейера (поза, коридор,
собственное движение, треки) переносится между шардами, поэтому запись
обрабатывается как одна непрерывная поездка.

    python tools/stream_new_data.py --shards 60 --out data/cache/new_data_candidates.npz
"""
from __future__ import annotations

import argparse
import re
import shutil
import tarfile
import tempfile
import time
from pathlib import Path

import numpy as np

import _bootstrap as B  # noqa: F401
from flyguard.bag import Bag
from flyguard.mushroom_body import describe
from flyguard.pipeline import FlyGuard, Params

SHARD_RE = re.compile(r"_(\d+)\.db3$")


def shard_members(tar_path: Path):
    """Список (номер, TarInfo) шардов, отсортированный по номеру."""
    out = []
    with tarfile.open(tar_path, "r:") as t:
        for m in t:
            mm = SHARD_RE.search(m.name)
            if mm and m.isfile():
                out.append((int(mm.group(1)), m.name, m.size))
    out.sort()
    return out


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--tar", default=r"C:\Games\Study\AI_Lidar\датасет\new_data")
    ap.add_argument("--out", default=str(B.CACHE / "new_data_candidates.npz"))
    ap.add_argument("--shards", type=int, default=0, help="0 — все")
    ap.add_argument("--stride", type=int, default=1, help="прореживание кадров")
    ap.add_argument("--workdir", default=None)
    args = ap.parse_args()

    tar_path = Path(args.tar)
    B.CACHE.mkdir(parents=True, exist_ok=True)
    members = shard_members(tar_path)
    if args.shards:
        members = members[:args.shards]
    print(f"шардов к обработке: {len(members)} из архива {tar_path.name}")

    work = Path(args.workdir) if args.workdir else Path(tempfile.mkdtemp(prefix="flyguard_"))
    work.mkdir(parents=True, exist_ok=True)

    fg = FlyGuard(Params(), memory=None)
    rows, dists, frames_total = [], [], 0
    t_start = time.time()

    with tarfile.open(tar_path, "r:") as tar:
        for k, (num, name, size) in enumerate(members):
            shard_dir = work / "bag"
            if shard_dir.exists():
                shutil.rmtree(shard_dir)
            shard_dir.mkdir(parents=True)
            target = shard_dir / Path(name).name

            src = tar.extractfile(name)
            with open(target, "wb") as dst:
                shutil.copyfileobj(src, dst, length=1 << 22)

            n_here = 0
            try:
                bag = Bag(shard_dir)
                for _, pc in bag.frames(stride=args.stride):
                    res = fg.process(pc)
                    n_here += 1
                    if res is None:
                        continue
                    for c in res.candidates:
                        rows.append(describe(c))
                        dists.append(c.d)
            except Exception as exc:                      # шард может быть обрезан
                print(f"  шард {num}: пропущен ({exc})")
            frames_total += n_here
            target.unlink(missing_ok=True)

            if k % 5 == 0 or k == len(members) - 1:
                el = time.time() - t_start
                print(f"  шард {k+1}/{len(members)} (#{num}) кадров {n_here:4d} | "
                      f"всего кадров {frames_total:6d} кандидатов {len(rows):7d} | "
                      f"{el/60:5.1f} мин", flush=True)

    shutil.rmtree(work, ignore_errors=True)
    X = np.stack(rows).astype(np.float32) if rows else np.zeros((0, 21), np.float32)
    np.savez_compressed(args.out, X=X, d=np.asarray(dists, np.float32),
                        frames=frames_total)
    print(f"готово: {frames_total} кадров, {X.shape[0]} кандидатов → {args.out}")
    print(f"время: {(time.time()-t_start)/60:.1f} мин")


if __name__ == "__main__":
    main()