#!/usr/bin/env python
from __future__ import annotations

import os
import sys
from pathlib import Path

from huggingface_hub import HfApi

SOURCE_CACHE = (
    Path.home()
    / ".cache/huggingface/hub/models--ibm-granite--granite-speech-4.1-2b-plus"
)
DEST_REPO = "olivius/granite-speech-4.1-2b-plus-mlx"


def find_weights_dir(root: Path) -> Path | None:
    if not root.exists():
        return None
    if list(root.glob("*.safetensors")) or (root / "config.json").exists():
        return root
    snapshots = root / "snapshots"
    if snapshots.exists():
        candidates = [
            path
            for path in snapshots.iterdir()
            if path.is_dir() and (list(path.glob("*.safetensors")) or (path / "config.json").exists())
        ]
        if candidates:
            return sorted(candidates, key=lambda p: p.stat().st_mtime)[-1]
    return None


def print_manual_commands() -> None:
    print(f"MLX weights not found at {SOURCE_CACHE}")
    print("Create them first with:")
    print("mlxconv ibm-granite/granite-speech-4.1-2b-plus")
    print("mlxconv ibm-granite/granite-speech-4.1-2b-plus --dtype q4_k_4")


def main() -> int:
    weights_dir = find_weights_dir(SOURCE_CACHE)
    if weights_dir is None:
        print_manual_commands()
        return 1

    token = os.environ.get("HF_TOKEN")
    if not token:
        print("HF_TOKEN is required to upload.", file=sys.stderr)
        return 2

    api = HfApi(token=token)
    api.create_repo(DEST_REPO, repo_type="model", exist_ok=True)
    api.upload_folder(
        repo_id=DEST_REPO,
        repo_type="model",
        folder_path=str(weights_dir),
        commit_message="Upload Granite Speech 4.1-2b-plus MLX weights",
    )
    print(f"Uploaded {weights_dir} to {DEST_REPO}")
    return 0


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