import os
import subprocess
import sys

def createOrActivateEnv():
    envPath = "env"
    if not os.path.exists(envPath):
        print("Creating virtual environment...")
        try:
            subprocess.run([sys.executable, "-m", "venv", envPath], check=True)
        except subprocess.CalledProcessError as e:
            print(f"Error creating virtual environment: {e}")
            sys.exit(1)
    
    activateScript = os.path.join(envPath, "Scripts", "activate") if os.name == "nt" else os.path.join(envPath, "bin", "activate")
    return activateScript

def installMissingPackages(scriptPath, importToInstall):
    envPath = "env"
    pipExecutable = os.path.join(envPath, "Scripts", "pip") if os.name == "nt" else os.path.join(envPath, "bin", "pip")


    while True:
        try:
            result = subprocess.run([sys.executable, scriptPath], capture_output=True, text=True)
            print(result.stdout)
            if result.stderr:
                print(result.stderr)

            if "ModuleNotFoundError" in result.stderr:
                missingModule = result.stderr.split("No module named ")[-1].strip().replace("'", "")
                installModule = importToInstall.get(missingModule, missingModule)

                print(f"Installing missing module: {installModule} (Imported as: {missingModule})")

                try:
                    subprocess.run([pipExecutable, "install", installModule], check=True)

                    print(f"Rerunning script {scriptPath} after installing {installModule}...")
                    result = subprocess.run([sys.executable, scriptPath], capture_output=True, text=True)
                    print(result.stdout)
                    if result.stderr:
                        print(result.stderr)

                    if "ModuleNotFoundError" not in result.stderr:
                        return
                    
                except subprocess.CalledProcessError as e:
                    print(f"Error installing module {installModule}: {e}")
                    sys.exit(1)

            else:
                break

        except FileNotFoundError as e:
            print(f"Error: Script not found: {scriptPath}")
            sys.exit(1)
        except Exception as e:
            print(f"An unexpected error occurred: {e}")
            sys.exit(1)

if __name__ == "__main__":
    scriptPath = sys.argv[1] if len(sys.argv) > 1 else "main.py"
    activateScript = createOrActivateEnv()
    print(f"Run the following command to activate the environment: {activateScript}")
    importToInstall = {
        "yaml": "PyYAML",
        "cv2": "opencv-python",
        "sklearn": "scikit-learn",
        "PIL": "Pillow",
        "dotenv": "python-dotenv"
    }
    installMissingPackages(scriptPath, importToInstall)