@echo off
setlocal enabledelayedexpansion

:: Set console colors
set "GREEN=[32m"
set "YELLOW=[33m"
set "RED=[31m"
set "RESET=[0m"

:: Function to print colored messages
call :define_print_functions

echo %GREEN%[INFO]%RESET% Starting WSL2 environment configuration...
echo.

:: Check for Administrator privileges
net session >nul 2>&1
if %errorLevel% neq 0 (
    echo %YELLOW%[WARN]%RESET% Administrator privileges required
    echo %GREEN%[INFO]%RESET% Please run this script as Administrator
    pause
    exit /b 1
)

echo %GREEN%[INFO]%RESET% Note: Please make sure you have manually updated WSL2 kernel before proceeding.
echo %GREEN%[INFO]%RESET% You can update WSL2 kernel by running "wsl --update" in an administrator PowerShell.
echo.

:: Check if Debian is installed
echo %GREEN%[INFO]%RESET% Checking Debian installation...
wsl -d Debian echo "Check Debian" >nul 2>&1
if %ERRORLEVEL% equ 0 (
    echo %GREEN%[INFO]%RESET% Found existing Debian installation
    
    :: Ask if user wants to remove existing installation
    set /p "remove=Do you want to remove the existing Debian installation? [y/N] "
    if /i "!remove!"=="y" (
        echo %GREEN%[INFO]%RESET% Removing existing Debian installation...
        wsl --unregister Debian
        if !ERRORLEVEL! neq 0 (
            echo %RED%[ERROR]%RESET% Failed to remove Debian
            pause
            exit /b 1
        )
        echo %GREEN%[INFO]%RESET% Existing Debian installation removed
        set "NEED_INSTALL=1"
    )
) else (
    echo %GREEN%[INFO]%RESET% No existing Debian installation found
    set "NEED_INSTALL=1"
)

:: Handle Debian installation if needed
if defined NEED_INSTALL (
    :: Get custom installation path
    set "DEFAULT_PATH=%USERPROFILE%\WSL\Debian"
    echo.
    echo Current WSL installation path: %DEFAULT_PATH%
    set /p "CUSTOM_PATH=Enter new installation path or press Enter to keep current: "
    
    if not "!CUSTOM_PATH!"=="" (
        echo %GREEN%[INFO]%RESET% Setting up custom installation path: !CUSTOM_PATH!
        
        :: Create directory if it doesn't exist
        if not exist "!CUSTOM_PATH!" (
            mkdir "!CUSTOM_PATH!" 2>nul
            if !ERRORLEVEL! neq 0 (
                echo %RED%[ERROR]%RESET% Failed to create directory: !CUSTOM_PATH!
                pause
                exit /b 1
            )
            echo %GREEN%[INFO]%RESET% Created directory: !CUSTOM_PATH!
        )
        
        :: Set WSL installation path
        echo %GREEN%[INFO]%RESET% Setting WSL installation path...
        setx WSLENV "WSLPATH/up:!CUSTOM_PATH!" >nul
        if !ERRORLEVEL! neq 0 (
            echo %RED%[ERROR]%RESET% Failed to set WSL installation path
            pause
            exit /b 1
        )
        echo %GREEN%[INFO]%RESET% WSL installation path set to: !CUSTOM_PATH!
    )
    
    :: Install Debian
    echo %GREEN%[INFO]%RESET% Installing Debian...
    wsl --install -d Debian
    if !ERRORLEVEL! neq 0 (
        echo %RED%[ERROR]%RESET% Failed to install Debian
        pause
        exit /b 1
    )
    echo %GREEN%[INFO]%RESET% Debian installation completed
    
    :: Wait for Debian to be ready
    echo %GREEN%[INFO]%RESET% Waiting for Debian to be ready...
    timeout /t 10 /nobreak >nul
)

:: Shutdown WSL to ensure clean state
echo %GREEN%[INFO]%RESET% Shutting down WSL to ensure clean state...
wsl --shutdown
if %ERRORLEVEL% neq 0 (
    echo %RED%[ERROR]%RESET% Failed to shutdown WSL
    pause
    exit /b 1
)
timeout /t 2 /nobreak >nul
echo %GREEN%[INFO]%RESET% WSL shutdown completed

:: Update Debian system
echo %GREEN%[INFO]%RESET% Updating Debian system packages...

echo %GREEN%[INFO]%RESET% Updating package lists...
wsl -d Debian -u root apt update
if %ERRORLEVEL% neq 0 (
    echo %RED%[ERROR]%RESET% Failed to update package lists
    pause
    exit /b 1
)

echo %GREEN%[INFO]%RESET% Upgrading packages...
wsl -d Debian -u root apt upgrade -y
if %ERRORLEVEL% neq 0 (
    echo %RED%[ERROR]%RESET% Failed to upgrade packages
    pause
    exit /b 1
)

echo %GREEN%[INFO]%RESET% Cleaning up...
wsl -d Debian -u root apt autoremove -y
wsl -d Debian -u root apt clean

echo %GREEN%[INFO]%RESET% Debian system update completed

:: Configure WSL memory limit
echo %GREEN%[INFO]%RESET% Configuring WSL memory limit...
set "CONFIG_FILE=%USERPROFILE%\.wslconfig"

:: Backup existing config if it exists
if exist "%CONFIG_FILE%" (
    echo %GREEN%[INFO]%RESET% Backing up existing .wslconfig...
    copy "%CONFIG_FILE%" "%CONFIG_FILE%.bak" >nul
)

:: Create new config file
(
    echo [wsl2]
    echo memory=8GB
    echo processors=4
    echo swap=2GB
) > "%CONFIG_FILE%"

if %ERRORLEVEL% neq 0 (
    echo %RED%[ERROR]%RESET% Failed to create .wslconfig
    pause
    exit /b 1
)
echo %GREEN%[INFO]%RESET% WSL memory limit configured

echo.
echo ============================================
echo          Configuration Complete!
echo ============================================
echo.
echo Your WSL2 environment has been configured:
echo.
echo %GREEN%[INFO]%RESET% 1. Debian system packages updated
echo %GREEN%[INFO]%RESET% 2. WSL memory limit configured (8GB RAM, 4 cores, 2GB swap)
if defined CUSTOM_PATH (
    echo %GREEN%[INFO]%RESET% 3. WSL installation path set to: !CUSTOM_PATH!
)
echo %GREEN%[INFO]%RESET% Please restart your computer to apply all changes.
echo ============================================

pause
exit /b 0

:define_print_functions
exit /b 0