$EnvName = "default"
$PythonVersion = "3.9"
$RequirementsFile = ""
$Force = $false
$Reinstall = $false

try {
    # Clear the console only when no command is being executed
    if ($args.Count -eq 0) {
        Clear-Host
    }
    

    # Check for miniforge installer and download if not found
    if (-not (Test-Path -Path .\miniforge.exe)) {
        Write-Host "miniforge.exe not found. Downloading from GitHub releases..."
        
        $downloadUrl = "https://github.com/conda-forge/miniforge/releases/download/25.3.1-0/Miniforge3-Windows-x86_64.exe"
        $tempFile = "Miniforge3-Windows-x86_64.exe"
        
        try {
            Invoke-WebRequest -Uri $downloadUrl -OutFile $tempFile -UseBasicParsing
            if (Test-Path -Path $tempFile) {
                Move-Item $tempFile "miniforge.exe" -Force
                Write-Host "Download completed successfully."
            }
            else {
                Write-Error "Failed to download miniforge installer."
                exit 2
            }
        }
        catch {
            Write-Error "Failed to download miniforge installer: $_"
            exit 2
        }
    }

    # Delete the old mamba folder if it exists and Force is specified
    if ($Force -and (Test-Path -Path .\mamba)) {
        try {
            Remove-Item -Recurse -Force .\mamba
        }
        catch {
            Write-Error "Failed to remove existing mamba directory. Check permissions."
            exit 3
        }
    }

    # Extract the Miniconda package if mamba folder doesn't exist
    if (-not (Test-Path -Path .\mamba)) {
        .\miniforge.exe /InstallationType=JustMe /AddToPath=0 /S /RegisterPython=0 /NoRegistry=1 /NoScripts=1 /NoShortcuts=1 /D=$PWD\mamba
        
        # Wait for installation to complete
        $nid = (Get-Process miniforge -ErrorAction SilentlyContinue).id
        if ($nid) { Wait-Process -Id $nid }
        
        # Move _conda.exe to conda.exe
        if (Test-Path .\mamba\_conda.exe) {
            Move-Item .\mamba\_conda.exe .\mamba\conda.exe -Force
        }
    }

    # Check for conda.exe
    if (-not (Test-Path -Path .\mamba\conda.exe)) {
        Write-Error "conda.exe not found in mamba directory. Extraction may have failed."
        exit 4
    }

    # Temporary clean the environment variable for make sure no other conda in env
    # Clear all conda-related environment variables
    $env:CONDA_PREFIX = $null
    $env:CONDA_DEFAULT_ENV = $null
    $env:CONDA_PYTHON_EXE = $null
    $env:CONDA_EXE = $null
    $env:CONDA_PROMPT_MODIFIER = $null

    # Remove any conda paths from PATH
    $env:Path = ($env:Path -split ';' | Where-Object { $_ -notlike '*conda*' -and $_ -notlike '*mamba*' -and $_ -notlike '*miniconda*' -and $_ -notlike '*miniforge*' }) -join ';'

    # Kill any conda/mamba processes
    try {
        Get-Process | Where-Object { $_.ProcessName -like '*conda*' -or $_.ProcessName -like '*mamba*' } | Stop-Process -Force -ErrorAction SilentlyContinue
    }
    catch {}

    # Add conda to PATH
    $env:Path = "$PWD\mamba;$PWD\mamba\Scripts;$env:Path"

    # Check if environment already exists
    $envPath = "$PWD\mamba\envs\$EnvName"
    if (Test-Path -Path $envPath) {
        if ($Reinstall) {
            try {
                Remove-Item -Recurse -Force $envPath
            }
            catch {
                Write-Error "Failed to remove existing environment. Check permissions."
                exit 8
            }
        }
    }

    # Create conda environment if it doesn't exist or if Reinstall is specified
    if (-not (Test-Path -Path $envPath)) {
        # Try to create the environment
        .\mamba\conda.exe create --no-shortcuts -y -k --prefix $envPath python=$PythonVersion
        if ($LASTEXITCODE -ne 0) {
            Write-Error "Failed to create environment. Python version may be invalid or directory not writable."
            exit 5
        }
    }

    # Set environment variables for the conda environment manually
    $env:CONDA_PREFIX = $envPath
    $env:CONDA_DEFAULT_ENV = $EnvName
    $env:CONDA_PYTHON_EXE = "$envPath\python.exe"

    # Add conda environment to PATH
    $env:Path = "$envPath;$envPath\Scripts;$envPath\Library\bin;$PWD\mamba;$PWD\mamba\Scripts;$env:Path"

    # Check for pip
    if (-not (Test-Path -Path "$envPath\Scripts\pip.exe")) {
        Write-Error "pip not found in environment. Environment may not have been created correctly."
        exit 6
    }

    # Install requirements if specified
    if ($RequirementsFile) {
        if (Test-Path -Path $RequirementsFile) {
            & "$envPath\python.exe" -m pip install -r $RequirementsFile
            if ($LASTEXITCODE -ne 0) {
                Write-Error "Failed to install requirements from $RequirementsFile."
                exit 7
            }
        }
        else {
            Write-Error "Requirements file $RequirementsFile not found. Skipping requirements installation."
        }
    }

    # execute arguments: run any command passed to the script within the prepared environment
    if ($args.Count -gt 0) {
        $exe = $args[0]
        if ($args.Count -gt 1) {
            $exeArgs = $args[1..($args.Count - 1)]
            & $exe @($exeArgs)
        }
        else {
            & $exe
        }
        exit $LASTEXITCODE
    }
}
catch {
    Write-Error "Unexpected error: $_"
    exit 99
}