import os
import sys
import traceback

def download_model():
    try:
        print("Starting model download...")
        
        # Import transformers here to provide better error messages if import fails
        try:
            from transformers import GroundingDinoProcessor, GroundingDinoForObjectDetection
        except ImportError:
            print("ERROR: Failed to import from transformers. Make sure you have installed the required packages:")
            print("pip install transformers torch")
            return False
        
        # Get the directory where the script is located
        if getattr(sys, 'frozen', False):
            # If the application is frozen (e.g., PyInstaller)
            current_dir = os.path.dirname(sys.executable)
        else:
            # If running as a script
            current_dir = os.path.dirname(os.path.abspath(__file__))
        
        # Define model directory path
        model_dir = os.path.join(current_dir, "models", "grounding-dino")
        
        # Ensure the model directory exists
        os.makedirs(model_dir, exist_ok=True)
        print(f"Model will be saved to: {model_dir}")
        
        # Download from Hugging Face and save locally
        print("Downloading processor...")
        processor = GroundingDinoProcessor.from_pretrained("IDEA-Research/grounding-dino-base")
        processor.save_pretrained(model_dir)
        
        print("Downloading model...")
        model = GroundingDinoForObjectDetection.from_pretrained("IDEA-Research/grounding-dino-base")
        model.save_pretrained(model_dir)
        
        print(f"✅ Model and processor downloaded and saved to: {model_dir}")
        
        # Set the environment variable so the app can access the model directory
        os.environ["MODEL_DIR"] = model_dir
        print(f"✅ MODEL_DIR environment variable set to: {model_dir}")
        
        return True
    
    except Exception as e:
        print(f"ERROR: Failed to download model: {str(e)}")
        print("Traceback:")
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = download_model()
    if not success:
        sys.exit(1)  # Exit with error code if download failed