#!/usr/bin/env python3

import os
import sys
import shutil
import zipfile
import tempfile
from pathlib import Path

from huggingface_hub import HfApi, hf_hub_download

HF_REPO_ID    = os.environ.get("HF_REPO_ID",   "eaglercraft/world")
HF_REPO_TYPE  = "dataset"
HF_TOKEN      = os.environ.get("HF_TOKEN",     "")
WORLD_DIR     = Path("/opt/server/backend/world")
ZIP_REPO_PATH = "world.zip"
TMP_DIR       = Path("/tmp/hf_world")



def _api() -> HfApi:
    return HfApi(token=HF_TOKEN)


def _ensure_repo():
    try:
        _api().create_repo(
            repo_id=HF_REPO_ID,
            repo_type=HF_REPO_TYPE,
            exist_ok=True,
            private=True,
        )
    except Exception:
        pass


def _repo_has_world() -> bool:
    try:
        files = list(_api().list_repo_files(
            repo_id=HF_REPO_ID,
            repo_type=HF_REPO_TYPE,
            token=HF_TOKEN,
        ))
        return ZIP_REPO_PATH in files
    except Exception as e:
        print(f"[sync_world] Could not list repo files: {e}")
        return False



def download() -> bool:
    if not HF_TOKEN:
        print("[sync_world] WARNING: HF_TOKEN not set, skipping download")
        return False

    if not _repo_has_world():
        print("[sync_world] world.zip not found in repo — first-run setup required")
        return False

    try:
        TMP_DIR.mkdir(parents=True, exist_ok=True)
        zip_path = hf_hub_download(
            repo_id=HF_REPO_ID,
            filename=ZIP_REPO_PATH,
            repo_type=HF_REPO_TYPE,
            token=HF_TOKEN,
            local_dir=str(TMP_DIR),
        )

        if WORLD_DIR.exists():
            shutil.rmtree(WORLD_DIR)
        WORLD_DIR.mkdir(parents=True, exist_ok=True)

        with zipfile.ZipFile(zip_path, "r") as zf:
            members = zf.namelist()
            top_dirs = {m.split("/")[0] for m in members if m.strip("/")}
            if len(top_dirs) == 1:
                top = top_dirs.pop()
                for member in members:
                    rel = member[len(top):].lstrip("/")
                    if not rel:
                        continue
                    dest = WORLD_DIR / rel
                    if member.endswith("/"):
                        dest.mkdir(parents=True, exist_ok=True)
                    else:
                        dest.parent.mkdir(parents=True, exist_ok=True)
                        with zf.open(member) as src, open(dest, "wb") as dst:
                            shutil.copyfileobj(src, dst)
            else:
                zf.extractall(WORLD_DIR)

        size_mb = os.path.getsize(zip_path) / 1_048_576
        print(f"[sync_world] World downloaded and extracted ({size_mb:.1f} MB)")
        return True

    except Exception as e:
        print(f"[sync_world] Download failed: {e}")
        return False


def upload() -> bool:
    if not HF_TOKEN:
        print("[sync_world] WARNING: HF_TOKEN not set, skipping upload")
        return False

    if not WORLD_DIR.exists():
        print(f"[sync_world] {WORLD_DIR} does not exist, skipping upload")
        return False

    tmp_zip = Path(tempfile.mktemp(suffix=".zip"))
    try:
        print("[sync_world] Zipping world...")
        with zipfile.ZipFile(tmp_zip, "w", zipfile.ZIP_DEFLATED) as zf:
            for fpath in WORLD_DIR.rglob("*"):
                if fpath.is_file():
                    zf.write(fpath, fpath.relative_to(WORLD_DIR.parent))

        size_mb = tmp_zip.stat().st_size / 1_048_576
        print(f"[sync_world] Uploading world.zip ({size_mb:.1f} MB)...")

        _ensure_repo()
        _api().upload_file(
            path_or_fileobj=str(tmp_zip),
            path_in_repo=ZIP_REPO_PATH,
            repo_id=HF_REPO_ID,
            repo_type=HF_REPO_TYPE,
            token=HF_TOKEN,
            commit_message=f"World autosave ({size_mb:.1f} MB)",
        )
        print(f"[sync_world] Upload complete.")
        return True

    except Exception as e:
        print(f"[sync_world] Upload failed: {e}")
        return False
    finally:
        tmp_zip.unlink(missing_ok=True)


def store_uploaded_zip(local_zip_path: str) -> bool:
    if not HF_TOKEN:
        print("[sync_world] WARNING: HF_TOKEN not set")
        return False

    try:
        size_mb = os.path.getsize(local_zip_path) / 1_048_576
        print(f"[sync_world] Storing user-supplied zip ({size_mb:.1f} MB)...")

        _ensure_repo()
        _api().upload_file(
            path_or_fileobj=local_zip_path,
            path_in_repo=ZIP_REPO_PATH,
            repo_id=HF_REPO_ID,
            repo_type=HF_REPO_TYPE,
            token=HF_TOKEN,
            commit_message=f"Initial world upload ({size_mb:.1f} MB)",
        )
        print("[sync_world] Zip stored on HF.")

        if WORLD_DIR.exists():
            shutil.rmtree(WORLD_DIR)
        WORLD_DIR.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(local_zip_path, "r") as zf:
            members = zf.namelist()
            top_dirs = {m.split("/")[0] for m in members if m.strip("/")}
            if len(top_dirs) == 1:
                top = top_dirs.pop()
                for member in members:
                    rel = member[len(top):].lstrip("/")
                    if not rel:
                        continue
                    dest = WORLD_DIR / rel
                    if member.endswith("/"):
                        dest.mkdir(parents=True, exist_ok=True)
                    else:
                        dest.parent.mkdir(parents=True, exist_ok=True)
                        with zf.open(member) as src, open(dest, "wb") as dst:
                            shutil.copyfileobj(src, dst)
            else:
                zf.extractall(WORLD_DIR)
        print("[sync_world] World extracted locally.")
        return True

    except Exception as e:
        print(f"[sync_world] store_uploaded_zip failed: {e}")
        return False



if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    action = sys.argv[1].lower()
    if action == "download":
        sys.exit(0 if download() else 1)
    elif action == "upload":
        sys.exit(0 if upload() else 1)
    elif action == "store" and len(sys.argv) >= 3:
        sys.exit(0 if store_uploaded_zip(sys.argv[2]) else 1)
    else:
        print(f"Unknown action: {action}")
        sys.exit(1)