#!/usr/bin/env python3
"""
Downloads Meridian ML models from a Hugging Face model repository.

Usage:
    python scripts/download_models.py

Environment variables:
    HF_MODEL_REPO  — HF repo ID (default: your-username/meridian-models)
    HF_TOKEN       — Optional HF token for private repos

Expected repo structure on HF Hub:
    your-username/meridian-models/
    ├── mentalroberta_onnx/
    │   ├── model.onnx          (473 MB)
    │   ├── config.json
    │   ├── tokenizer.json
    │   ├── tokenizer_config.json
    │   ├── vocab.json
    │   ├── merges.txt
    │   └── special_tokens_map.json
    ├── svm_depression.joblib   (20 KB)
    └── scaler.joblib           (19 KB)

To create this repo:
    1. pip install huggingface_hub
    2. huggingface-cli login
    3. huggingface-cli upload your-username/meridian-models backend/models/ . --repo-type model
"""
import os
import logging
from huggingface_hub import snapshot_download

logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger(__name__)

REPO_ID = os.getenv("HF_MODEL_REPO", "astroknotsheep/meridian-models")
TOKEN = os.getenv("HF_TOKEN", None)
LOCAL_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "backend", "models")


def main():
    logger.info(f"📦 Downloading models from: {REPO_ID}")
    logger.info(f"📁 Target directory: {LOCAL_DIR}")

    os.makedirs(LOCAL_DIR, exist_ok=True)

    snapshot_download(
        repo_id=REPO_ID,
        repo_type="model",
        local_dir=LOCAL_DIR,
        token=TOKEN,
        # Only download the files we need (skip README, .gitattributes etc.)
        allow_patterns=[
            "mentalroberta_onnx/*",
            "svm_depression.joblib",
            "scaler.joblib",
        ],
    )

    # Verify critical files exist
    checks = [
        os.path.join(LOCAL_DIR, "mentalroberta_onnx", "model.onnx"),
        os.path.join(LOCAL_DIR, "svm_depression.joblib"),
        os.path.join(LOCAL_DIR, "scaler.joblib"),
    ]
    for path in checks:
        if os.path.isfile(path):
            size_mb = os.path.getsize(path) / (1024 * 1024)
            logger.info(f"  ✓ {os.path.basename(path)} ({size_mb:.1f} MB)")
        else:
            logger.error(f"  ✗ MISSING: {path}")
            raise FileNotFoundError(f"Model file not found after download: {path}")

    logger.info("✅ All models downloaded successfully")


if __name__ == "__main__":
    main()