import os
import sys
import subprocess

def run_step(step_name, command_args):
    print(f"\n>>> Running: {step_name}...")
    try:
        # Run python module or script in the current environment
        result = subprocess.run(
            [sys.executable] + command_args,
            cwd=os.getcwd(),
            check=True
        )
        print(f"Success: {step_name} completed.")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error in {step_name}: {str(e)}")
        return False

def setup():
    print("=" * 60)
    print("Phi-3 Text-to-SQL Studio Setup Assistant")
    print("=" * 60)
    
    # 1. Install remaining dependencies
    print("\n>>> Checking and installing remaining python libraries...")
    libraries = [
        "transformers", 
        "peft", 
        "bitsandbytes", 
        "trl", 
        "datasets", 
        "accelerate", 
        "flask", 
        "pandas"
    ]
    
    try:
        # Run pip install inside our virtual environment
        pip_path = os.path.join(".venv", "Scripts", "pip") if os.path.exists(".venv") else "pip"
        cmd = [pip_path, "install", "-U"] + libraries
        print(f"Executing: {' '.join(cmd)}")
        subprocess.run(cmd, check=True)
        print("Success: Libraries installed successfully.")
    except Exception as e:
        print(f"Warning during library installation: {str(e)}")
        print("Please ensure your virtual environment is activated and run: pip install -r requirements.txt")

    # 2. Setup SQLite Database
    success = run_step("Database Initialization", ["src/database.py"])
    if not success:
        sys.exit(1)
        
    # 3. Setup Dataset Generation
    success = run_step("Dataset Generation", ["src/dataset.py"])
    if not success:
        sys.exit(1)

    print("\n" + "=" * 60)
    print("ALL ENVIRONMENT COMPONENT SETUP IS COMPLETE!")
    print("=" * 60)
    print("\nYou are now ready to run:")
    print("1. Fine-tuning: Run 'python src/train.py' to train your model.")
    print("2. Web Studio Dashboard: Run 'python src/server.py' and open http://127.0.0.1:5000 in your browser.")
    print("=" * 60)

if __name__ == "__main__":
    setup()