import argparse
import sys
import os

# Define available models and their Hugging Face Repository IDs
MODELS = {
    "MicroStructFormer_Swin_Large": "Dnq2025/MicroStructFormer", # TO-DO: Replace with actual HF repo
}

def list_models():
    print("Available models:")
    for model_name in MODELS.keys():
        print(f"  - {model_name}")

def download_model(model_name, dest_dir=None):
    if model_name not in MODELS:
        print(f"Error: Model '{model_name}' not found.")
        print("Use --list to see available models.")
        sys.exit(1)
        
    repo_id = MODELS[model_name]
    
    if dest_dir is None:
        # Default to a folder named after the model in the current working directory
        dest_dir = os.path.join(os.getcwd(), model_name)
        
    print(f"Downloading model '{model_name}' from Hugging Face...")
    print(f"Target directory: {dest_dir}")
    
    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        print("Error: 'huggingface_hub' is not installed. Please install it using 'pip install huggingface_hub'.")
        sys.exit(1)
        
    try:
        # Download the repository cleanly
        # We ignore .git files and READMEs to only get the model weights and configs
        snapshot_download(
            repo_id=repo_id,
            local_dir=dest_dir,
            ignore_patterns=["*.md", ".git*"]
        )
        print(f"\nSuccessfully downloaded '{model_name}' to {dest_dir}")
        print("You can now set MODEL_DIR to this path in your inference scripts.")
    except Exception as e:
        print(f"An error occurred during download: {e}")
        sys.exit(1)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="MicroStructFormer Model Downloader")
    parser.add_argument("--list", action="store_true", help="List available models")
    parser.add_argument("-m", "--model", type=str, help="Name of the model to download")
    parser.add_argument("-d", "--dest", type=str, help="Destination directory (optional)")
    
    args = parser.parse_args()
    
    if args.list:
        list_models()
        sys.exit(0)
        
    if args.model:
        download_model(args.model, args.dest)
    else:
        parser.print_help()