"""
Script to download YOLOv8 model and save it to the models directory
This ensures the model is included in the repository for Hugging Face Spaces
"""

import os
from ultralytics import YOLO

def download_model(model_name="yolov8n.pt", save_dir="models"):
    """
    Download a YOLOv8 model and save it to the specified directory
    
    Args:
        model_name: Name of the YOLOv8 model to download
        save_dir: Directory to save the model to
    """
    # Create models directory if it doesn't exist
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        print(f"Created directory: {save_dir}")
    
    # Full path to save the model
    model_path = os.path.join(save_dir, model_name)
    
    # Check if model already exists
    if os.path.exists(model_path):
        print(f"Model already exists at: {model_path}")
        # Load model to verify it's valid
        try:
            model = YOLO(model_path)
            print("Model loaded successfully")
            return model_path
        except Exception as e:
            print(f"Error loading existing model: {e}")
            print("Downloading new model...")
    
    # Download the model
    try:
        print(f"Downloading model: {model_name}")
        model = YOLO(model_name)
        
        # Save the model to the specified path
        model_file = model.export(format="torchscript")
        
        # Rename to the original model name
        if os.path.exists(model_file) and model_file != model_path:
            # Copy the model file to the specified path
            import shutil
            shutil.copy(model_file, model_path)
            print(f"Model saved to: {model_path}")
        
        return model_path
    except Exception as e:
        print(f"Error downloading model: {e}")
        return None

if __name__ == "__main__":
    # Download YOLOv8n model
    model_path = download_model()
    
    if model_path:
        print(f"Model successfully downloaded to: {model_path}")
    else:
        print("Failed to download model")