#!/usr/bin/env python
# coding=utf-8

"""
Script to install requirements in the correct order for the Phi-4 training project.
This ensures base requirements are installed first, followed by additional requirements.
"""

import os
import sys
import subprocess
import argparse
import logging
from pathlib import Path

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

def install_requirements(include_flash=False):
    """Install requirements in the correct order."""
    current_dir = Path(__file__).parent
    base_req_path = current_dir / "requirements-base.txt"
    main_req_path = current_dir / "requirements.txt"
    flash_req_path = current_dir / "requirements-flash.txt"
    
    if not base_req_path.exists():
        logger.error(f"Base requirements file not found: {base_req_path}")
        return False
    
    if not main_req_path.exists():
        logger.error(f"Main requirements file not found: {main_req_path}")
        return False
    
    logger.info("Installing dependencies in sequential order...")
    
    try:
        # Step 1: Install base requirements
        logger.info(f"Step 1: Installing base requirements from {base_req_path}")
        subprocess.run([sys.executable, "-m", "pip", "install", "-r", str(base_req_path)], 
                      check=True)
        logger.info("Base requirements installed successfully")
        
        # Step 2: Install main requirements
        logger.info(f"Step 2: Installing additional requirements from {main_req_path}")
        subprocess.run([sys.executable, "-m", "pip", "install", "-r", str(main_req_path)], 
                      check=True)
        logger.info("Additional requirements installed successfully")
        
        # Step 3: Optionally install flash-attention
        if include_flash and flash_req_path.exists():
            logger.info(f"Step 3: Installing flash-attention from {flash_req_path}")
            subprocess.run([sys.executable, "-m", "pip", "install", "-r", str(flash_req_path), "--no-build-isolation"], 
                          check=True)
            logger.info("Flash-attention installed successfully")
        elif include_flash:
            logger.warning(f"Flash requirements file not found: {flash_req_path}")
        
        logger.info("All required packages installed successfully!")
        return True
    
    except subprocess.CalledProcessError as e:
        logger.error(f"Error installing dependencies: {str(e)}")
        return False

def main():
    parser = argparse.ArgumentParser(description="Install requirements for Phi-4 training")
    parser.add_argument("--flash", action="store_true", help="Also install flash-attention (optional)")
    args = parser.parse_args()
    
    success = install_requirements(include_flash=args.flash)
    if success:
        logger.info("Installation completed successfully!")
    else:
        logger.error("Installation failed. Please check the logs for details.")
        sys.exit(1)

if __name__ == "__main__":
    main()