"""Upload model weights to Modal Volume for BoccaccioAI deployment.

Usage:
    python scripts/upload_weights_modal.py --input checkpoints/finetune/model.safetensors
"""

from __future__ import annotations

import argparse
from pathlib import Path

import modal


VOLUME_NAME = "boccaccio-weights"
REMOTE_PATH = "model.safetensors"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Upload model weights to Modal Volume",
    )
    parser.add_argument(
        "--input",
        type=str,
        default="checkpoints/finetune/model.safetensors",
        help="Path to the local safetensors file.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    local_path = Path(args.input)
    if not local_path.exists():
        raise FileNotFoundError(f"File not found: {local_path}")

    size_gb = local_path.stat().st_size / 1024**3
    print(f"Uploading {local_path} ({size_gb:.2f} GB) to Modal Volume '{VOLUME_NAME}'...")

    vol = modal.Volume.from_name(VOLUME_NAME, create_if_missing=True)

    with vol.batch_upload(force=True) as batch:
        batch.put_file(str(local_path), REMOTE_PATH)

    print(f"Done! Weights available at /weights/{REMOTE_PATH} in Modal containers.")


if __name__ == "__main__":
    main()