"""Verify or download Lyra-2.0 checkpoints.

Three cases handled, in priority order:
1. /data/Lyra-2/checkpoints/model is already populated (read-only direct-mount
   OR warm bucket): skip download, just symlink.
2. /data is writable + empty: download from HF Hub.
3. /data is unwritable + empty: fall back to ephemeral cache.
"""

from __future__ import annotations

import logging
import os
import shutil
import sys
from pathlib import Path

from huggingface_hub import snapshot_download

log = logging.getLogger("lyra2-space.checkpoints")
logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(message)s")

REPO_ID = "nvidia/Lyra-2.0"
DATA_ROOT = Path(os.environ.get("LYRA2_DATA_ROOT", "/data"))
APP_ROOT = Path("/home/user/app")
APP_CHECKPOINTS = APP_ROOT / "Lyra-2" / "checkpoints"


def _already_populated(root: Path) -> bool:
    """True if root/Lyra-2/checkpoints/model has content."""
    marker = root / "Lyra-2" / "checkpoints" / "model"
    return marker.is_dir() and any(marker.iterdir())


def _is_writable(root: Path) -> bool:
    try:
        root.mkdir(parents=True, exist_ok=True)
        probe = root / ".writable_probe"
        probe.write_text("ok")
        probe.unlink()
        return True
    except OSError:
        return False


def _symlink_app_checkpoints(lyra2_root: Path) -> None:
    """Point /home/user/app/Lyra-2/checkpoints at <lyra2_root>/checkpoints."""
    src = lyra2_root / "checkpoints"
    dst = APP_CHECKPOINTS
    if dst.is_symlink() and dst.resolve() == src.resolve():
        log.info("Checkpoints symlink already points to %s", src)
        return
    if dst.is_symlink() or dst.is_file():
        dst.unlink()
    elif dst.exists():
        shutil.rmtree(dst)
    dst.parent.mkdir(parents=True, exist_ok=True)
    dst.symlink_to(src, target_is_directory=True)
    log.info("Linked %s -> %s", dst, src)


def main() -> int:
    # Case 1: /data already has the checkpoints (direct-mount OR warm bucket)
    if DATA_ROOT.is_dir() and _already_populated(DATA_ROOT):
        log.info("Checkpoints present at %s/Lyra-2 — skipping download.", DATA_ROOT)
        _symlink_app_checkpoints(DATA_ROOT / "Lyra-2")
        return 0

    # Case 2: /data writable but empty — download
    if _is_writable(DATA_ROOT):
        lyra2_root = DATA_ROOT / "Lyra-2"
        lyra2_root.mkdir(parents=True, exist_ok=True)
        log.info("Downloading %s checkpoints/* to %s (first cold start only).",
                 REPO_ID, lyra2_root)
        snapshot_download(
            repo_id=REPO_ID,
            allow_patterns=["checkpoints/*"],
            local_dir=str(lyra2_root),
            max_workers=8,
        )
        log.info("Checkpoint download complete.")
        _symlink_app_checkpoints(lyra2_root)
        return 0

    # Case 3: /data unwritable + empty — ephemeral fallback
    fallback = APP_ROOT / "_checkpoints_cache"
    log.warning("Cannot use %s (read-only AND empty). Falling back to ephemeral %s",
                DATA_ROOT, fallback)
    fallback.mkdir(parents=True, exist_ok=True)
    lyra2_root = fallback / "Lyra-2"
    lyra2_root.mkdir(parents=True, exist_ok=True)
    if not _already_populated(fallback):
        log.info("Downloading %s checkpoints/* to %s.", REPO_ID, lyra2_root)
        snapshot_download(
            repo_id=REPO_ID,
            allow_patterns=["checkpoints/*"],
            local_dir=str(lyra2_root),
            max_workers=8,
        )
    _symlink_app_checkpoints(lyra2_root)
    return 0


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