"""
Download and save a HuggingFace model locally.

Usage:
    python code_migration/download_model.py
    MODEL_NAME=Qwen/Qwen3.5-9B python code_migration/download_model.py
    MODEL_NAME=Qwen/Qwen3.5-4B SAVE_DIR=./qwen3.5-4b python code_migration/download_model.py
"""

import os
import sys

from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen3.5-4B")
SAVE_DIR = os.getenv("SAVE_DIR", None)


def main():
    model_name = MODEL_NAME
    save_dir = SAVE_DIR or f"./{model_name.split('/')[-1].lower()}"

    print(f"Downloading model: {model_name}")
    print(f"Save directory:    {save_dir}")
    print()

    # Download tokenizer
    print("Downloading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

    # Download model (full precision, no quantization — save the raw weights)
    print("Downloading model weights...")
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype="auto",
        trust_remote_code=True,
    )

    # Save locally
    print(f"Saving to {save_dir}...")
    os.makedirs(save_dir, exist_ok=True)
    model.save_pretrained(save_dir)
    tokenizer.save_pretrained(save_dir)

    print(f"Done. Model saved to {save_dir}")
    print(f"Use it with: MODEL_NAME={save_dir} python code_migration/inference.py")


if __name__ == "__main__":
    main()