#!/usr/bin/env python3
"""
Build script for Hugging Face Spaces - downloads model files at build time
"""

import os
import sys
import subprocess
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def run_download_script():
    """Run the model download script"""
    try:
        logger.info("Running advanced model download script...")
        
        # Try the advanced download script first
        result = subprocess.run([sys.executable, "download_model_advanced.py"], 
                              capture_output=True, text=True, check=True)
        logger.info("Advanced model download completed successfully")
        logger.info(result.stdout)
        return True
        
    except subprocess.CalledProcessError as e:
        logger.warning(f"Advanced download failed: {e}")
        logger.warning("Falling back to basic download script...")
        
        try:
            # Fallback to basic download script
            result = subprocess.run([sys.executable, "download_model.py"], 
                                  capture_output=True, text=True, check=True)
            logger.info("Basic model download completed successfully")
            logger.info(result.stdout)
            return True
            
        except subprocess.CalledProcessError as e2:
            logger.error(f"Basic download also failed: {e2}")
            logger.error(e2.stderr)
            return False

def verify_build():
    """Verify that the build was successful"""
    try:
        logger.info("Verifying build results...")
        
        # Check if int4 directory exists
        if not os.path.exists("./int4"):
            logger.error("int4 directory not found")
            return False
        
        # Check for essential files
        essential_files = [
            "config.json",
            "pytorch_model.bin",
            "tokenizer.json",
            "tokenizer_config.json"
        ]
        
        missing_files = []
        for file in essential_files:
            file_path = os.path.join("./int4", file)
            if not os.path.exists(file_path):
                missing_files.append(file)
        
        if missing_files:
            logger.error(f"Missing essential files: {missing_files}")
            return False
        
        # Check file sizes
        total_size = 0
        for file in essential_files:
            file_path = os.path.join("./int4", file)
            if os.path.exists(file_path):
                file_size = os.path.getsize(file_path)
                total_size += file_size
                logger.info(f"✅ {file}: {file_size} bytes")
        
        logger.info(f"Total model size: {total_size / (1024*1024):.2f} MB")
        
        if total_size < 1000000:  # Less than 1MB
            logger.warning("Model files seem too small")
            return False
        
        logger.info("Build verification completed successfully")
        return True
        
    except Exception as e:
        logger.error(f"Error verifying build: {e}")
        return False

def main():
    """Main build function"""
    logger.info("Starting Hugging Face Space build process...")
    
    # Run the model download script
    if run_download_script():
        # Verify the build
        if verify_build():
            logger.info("Build process completed successfully")
            return True
        else:
            logger.error("Build verification failed")
            return False
    else:
        logger.error("Model download failed")
        return False

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)