#!/usr/bin/env python3
"""
download_model.py - Utility script to download LionGuard2 model from Hugging Face
"""

import os
import argparse
from huggingface_hub import hf_hub_download


def download_lionguard2(repo_id, filename="LionGuard2.safetensors", token=None, output_dir="./"):
    """
    Download LionGuard2 model from Hugging Face private repository.
    
    Args:
        repo_id: The Hugging Face repository ID (e.g., "username/repo-name")
        filename: The filename to download (default: "LionGuard2.safetensors")
        token: Hugging Face access token for private repositories
        output_dir: Directory to save the downloaded file
    """
    if token is None:
        token = os.environ.get("HF_API_KEY")
        if not token:
            print("Error: No HF_API_KEY found in environment variables.")
            print("Please set your Hugging Face token:")
            print("export HF_API_KEY=your_token_here")
            return False
    
    try:
        print(f"Downloading {filename} from {repo_id}...")
        
        # Download the model file
        model_path = hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            token=token,
            local_dir=output_dir,
            local_dir_use_symlinks=False  # Download actual file, not symlink
        )
        
        print(f"✅ Model successfully downloaded to: {model_path}")
        return True
        
    except Exception as e:
        print(f"❌ Failed to download model: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(description="Download LionGuard2 model from Hugging Face")
    parser.add_argument("repo_id", help="Hugging Face repository ID (e.g., username/repo-name)")
    parser.add_argument("--filename", default="LionGuard2.safetensors", help="Filename to download")
    parser.add_argument("--token", help="Hugging Face access token (optional if HF_API_KEY env var is set)")
    parser.add_argument("--output-dir", default="./", help="Output directory for downloaded file")
    
    args = parser.parse_args()
    
    success = download_lionguard2(
        repo_id=args.repo_id,
        filename=args.filename,
        token=args.token,
        output_dir=args.output_dir
    )
    
    if success:
        print(f"\n🎉 Ready to use! The model has been downloaded and can now be used by the application.")
    else:
        print(f"\n💡 Make sure you have:")
        print(f"   1. Valid Hugging Face token with access to the private repository")
        print(f"   2. Correct repository ID: {args.repo_id}")
        print(f"   3. The model file exists in the repository")


if __name__ == "__main__":
    main()