"""Make the pre-built bundle (deploy_data/: vector index + step thumbnails + manual figures) available
in ./deploy_data before the app starts.

Sources, in this order:
  1. FIXSENSE_BUNDLE_MOUNT (e.g. /data/bundle): the private dataset is *mounted* read-only into the Space by
     Hugging Face. Reading ~1000 small files through the mount is very slow, so the dataset also holds a single
     ``bundle.tar``; that one file is read sequentially and unpacked onto the local disk (the embedded vector
     store needs a writable directory anyway). Falls back to copying loose files if there is no tar.
  2. FIXSENSE_BUNDLE_REPO: download the private dataset with HF_TOKEN.
Neither set -> do nothing (local run, or the bundle was baked into the image).
"""
from __future__ import annotations

import os
import shutil
import sys
import tarfile
import time


def _ok(path: str) -> bool:
    return os.path.isfile(os.path.join(path, "index", "fixsense_index_meta.json"))


def _log(msg: str) -> None:
    print(f"[fetch_bundle] {msg}", flush=True)


def _from_mount(mount: str, target: str) -> int:
    t0 = time.time()
    tar = os.path.join(mount, "bundle.tar")
    if os.path.isfile(tar):
        _log(f"unpacking {tar} ({os.path.getsize(tar) / 1e6:.0f} MB) -> ./{target}")
        os.makedirs(target, exist_ok=True)
        with tarfile.open(tar, "r:") as tf:
            tf.extractall(target)
    elif _ok(mount):
        _log(f"no bundle.tar; copying loose files from {mount} (slow on a mount)")
        n = 0
        for dirpath, _dirs, files in os.walk(mount):
            rel = os.path.relpath(dirpath, mount)
            os.makedirs(os.path.join(target, rel), exist_ok=True)
            for f in files:
                if f == ".lock":
                    continue
                shutil.copyfile(os.path.join(dirpath, f), os.path.join(target, rel, f))
                n += 1
                if n % 100 == 0:
                    _log(f"  copied {n} files ({time.time() - t0:.0f}s)")
    else:
        _log(f"ERROR: {mount} holds neither bundle.tar nor index/fixsense_index_meta.json; contents: {os.listdir(mount)[:10]}")
        return 1
    if not _ok(target):
        _log(f"ERROR: ./{target}/index/fixsense_index_meta.json missing after unpacking")
        return 1
    _log(f"bundle ready in ./{target} ({time.time() - t0:.0f}s)")
    return 0


def _report_videos() -> None:
    """Startup diagnostics: is the repair-video bucket mounted, and how fast can it be read?"""
    vdir = os.environ.get("FIXSENSE_VIDEO_DIR", "").strip()
    if not vdir:
        return
    mp4 = [os.path.join(d, f) for d, _s, fs in os.walk(vdir) for f in fs if f.lower().endswith(".mp4")] if os.path.isdir(vdir) else []
    _log(f"video mount {vdir}: exists={os.path.isdir(vdir)} mp4 files={len(mp4)}")
    if mp4:
        t0 = time.time()
        with open(mp4[0], "rb") as fh:
            fh.seek(os.path.getsize(mp4[0]) // 2)
            n = len(fh.read(2_000_000))
        _log(f"  mid-file read of {n / 1e6:.1f} MB from {os.path.basename(mp4[0])[:40]}: {time.time() - t0:.2f}s")


def main() -> int:
    _report_videos()
    target = os.environ.get("FIXSENSE_BUNDLE_DIR", "deploy_data")
    mount = os.environ.get("FIXSENSE_BUNDLE_MOUNT", "").strip()
    repo = os.environ.get("FIXSENSE_BUNDLE_REPO", "").strip()

    if mount:
        _log(f"mount {mount}: exists={os.path.isdir(mount)}")
        return _from_mount(mount, target)

    if repo:
        from huggingface_hub import snapshot_download

        path = snapshot_download(repo_id=repo, repo_type="dataset", local_dir=target, token=os.environ.get("HF_TOKEN") or None)
        if not _ok(path):
            _log(f"ERROR: {repo} is not a bundle built by run_build_index_v2.py")
            return 1
        _log(f"bundle ready in ./{target} (downloaded from {repo})")
        return 0

    _log(f"no bundle source configured; using existing ./{target} if present")
    return 0


if __name__ == "__main__":
    sys.exit(main())