#!/usr/bin/env python3
"""Download WeatherNext model weights from Google Cloud Storage."""

import os
import sys
from pathlib import Path
from config import WEIGHTS_DIR, MODEL_CONFIGS

try:
    from google.cloud import storage
except ImportError:
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "google-cloud-storage"])
    from google.cloud import storage

def download_blob(bucket_name, source_blob_name, destination):
    """Download a blob from GCS bucket."""
    print(f"Downloading {source_blob_name}...")
    storage_client = storage.Client.create_anonymous_client()
    bucket = storage_client.bucket(bucket_name)
    blob = bucket.blob(source_blob_name)
    blob.download_to_filename(str(destination))
    print(f"  -> {destination} ({os.path.getsize(destination) / 1e6:.1f} MB)")

def download_all():
    """Download all required model weights."""
    bucket = "dm_graphcast"

    weights = [
        "WeatherNextCyclones_Mini_<2024.npz",
        "WeatherNextCyclones_Mini_<2023.npz",
    ]

    for w in weights:
        dst = WEIGHTS_DIR / w
        if dst.exists():
            print(f"Skipping {w} (already exists)")
            continue
        try:
            download_blob(bucket, w, dst)
        except Exception as e:
            print(f"Failed to download {w}: {e}")

    print("\n=== Download complete ===")
    print(f"Weights stored in: {WEIGHTS_DIR}")

if __name__ == "__main__":
    download_all()