# =============================================================================
# Drone Simulation - Windows Installation Script (PowerShell)
# =============================================================================
# Installs ROS 2 Humble, PyBullet, and all required dependencies
# Uses a Python virtual environment for pip packages
#
# Usage:
#   1. Open PowerShell as Administrator
#   2. Run: Set-ExecutionPolicy RemoteSigned -Scope CurrentUser
#   3. Run: .\install_windows.ps1
#
# Tested on: Windows 10/11
# =============================================================================

Write-Host "==============================================" -ForegroundColor Cyan
Write-Host "  Drone Simulation - Windows Installation" -ForegroundColor Cyan
Write-Host "==============================================" -ForegroundColor Cyan
Write-Host ""

# Get script and project paths
$ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path
$ProjectRoot = Split-Path -Parent $ScriptDir
$VenvDir = Join-Path $ProjectRoot "venv"

Write-Host "[INFO] Project root: $ProjectRoot" -ForegroundColor Gray
Write-Host "[INFO] Virtual environment: $VenvDir" -ForegroundColor Gray

# Check for Administrator privileges
$isAdmin = ([Security.Principal.WindowsPrincipal] [Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
if (-not $isAdmin) {
    Write-Host "[WARN] Not running as Administrator. Some installations may fail." -ForegroundColor Yellow
    Write-Host "       Consider running PowerShell as Administrator." -ForegroundColor Yellow
    Write-Host ""
}

# Function to refresh environment PATH
function Refresh-Path {
    $env:Path = [System.Environment]::GetEnvironmentVariable("Path","Machine") + ";" + [System.Environment]::GetEnvironmentVariable("Path","User")
    # Also add common Chocolatey paths
    $chocoPath = "$env:ProgramData\chocolatey\bin"
    if (Test-Path $chocoPath) {
        $env:Path = "$chocoPath;$env:Path"
    }
}

# -----------------------------------------------------------------------------
# Step 1: Install Chocolatey (Package Manager)
# -----------------------------------------------------------------------------
Write-Host "[STEP 1/7] Checking Chocolatey..." -ForegroundColor Green

$chocoInstalled = $false
try {
    $chocoVersion = choco --version 2>$null
    if ($chocoVersion) {
        $chocoInstalled = $true
        Write-Host "[INFO] Chocolatey already installed (version $chocoVersion)" -ForegroundColor Green
    }
} catch {
    $chocoInstalled = $false
}

if (-not $chocoInstalled) {
    Write-Host "[INFO] Installing Chocolatey..." -ForegroundColor Yellow
    
    try {
        Set-ExecutionPolicy Bypass -Scope Process -Force
        [System.Net.ServicePointManager]::SecurityProtocol = [System.Net.ServicePointManager]::SecurityProtocol -bor 3072
        Invoke-Expression ((New-Object System.Net.WebClient).DownloadString('https://community.chocolatey.org/install.ps1'))
        
        # Refresh environment to find choco
        Refresh-Path
        
        # Verify installation
        $chocoPath = "$env:ProgramData\chocolatey\bin\choco.exe"
        if (Test-Path $chocoPath) {
            Write-Host "[INFO] Chocolatey installed successfully" -ForegroundColor Green
        } else {
            Write-Host "[ERROR] Chocolatey installation failed. Please install manually:" -ForegroundColor Red
            Write-Host "        https://chocolatey.org/install" -ForegroundColor Yellow
            Write-Host "[INFO] After installing Chocolatey, close and reopen PowerShell, then run this script again." -ForegroundColor Yellow
            exit 1
        }
    } catch {
        Write-Host "[ERROR] Failed to install Chocolatey: $_" -ForegroundColor Red
        Write-Host "[INFO] Please install Chocolatey manually: https://chocolatey.org/install" -ForegroundColor Yellow
        exit 1
    }
}

# -----------------------------------------------------------------------------
# Step 2: Install Python
# -----------------------------------------------------------------------------
Write-Host ""
Write-Host "[STEP 2/7] Installing Python..." -ForegroundColor Green

$pythonInstalled = $false
try {
    $pythonVersion = python --version 2>$null
    if ($pythonVersion -match "Python 3") {
        $pythonInstalled = $true
        Write-Host "[INFO] Python already installed ($pythonVersion)" -ForegroundColor Green
    }
} catch {
    $pythonInstalled = $false
}

if (-not $pythonInstalled) {
    Write-Host "[INFO] Installing Python 3.11..." -ForegroundColor Yellow
    
    # Use full path to choco if needed
    $chocoExe = "$env:ProgramData\chocolatey\bin\choco.exe"
    if (Test-Path $chocoExe) {
        & $chocoExe install python311 -y
    } else {
        choco install python311 -y
    }
    
    Refresh-Path
    
    # Verify
    try {
        $pythonVersion = python --version
        Write-Host "[INFO] Python installed ($pythonVersion)" -ForegroundColor Green
    } catch {
        Write-Host "[ERROR] Python installation failed" -ForegroundColor Red
        Write-Host "[INFO] Please install Python 3.11 manually from https://python.org" -ForegroundColor Yellow
        exit 1
    }
}

# -----------------------------------------------------------------------------
# Step 3: Install Visual C++ Build Tools (optional but recommended)
# -----------------------------------------------------------------------------
Write-Host ""
Write-Host "[STEP 3/7] Checking Visual C++ Build Tools..." -ForegroundColor Green

# Check if cl.exe exists (Visual C++ compiler)
$vsInstalled = $false
$vsPaths = @(
    "C:\Program Files\Microsoft Visual Studio\2022\BuildTools\VC\Tools\MSVC",
    "C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Tools\MSVC",
    "C:\Program Files (x86)\Microsoft Visual Studio\2019\BuildTools\VC\Tools\MSVC"
)

foreach ($path in $vsPaths) {
    if (Test-Path $path) {
        $vsInstalled = $true
        Write-Host "[INFO] Visual C++ Build Tools found" -ForegroundColor Green
        break
    }
}

if (-not $vsInstalled) {
    Write-Host "[INFO] Visual C++ Build Tools not found" -ForegroundColor Yellow
    Write-Host "[INFO] Attempting to install (this may take 10-20 minutes)..." -ForegroundColor Yellow
    
    try {
        $chocoExe = "$env:ProgramData\chocolatey\bin\choco.exe"
        if (Test-Path $chocoExe) {
            & $chocoExe install visualstudio2022-workload-vctools -y
        } else {
            choco install visualstudio2022-workload-vctools -y
        }
        Write-Host "[INFO] Visual C++ Build Tools installed" -ForegroundColor Green
    } catch {
        Write-Host "[WARN] Could not install Visual C++ Build Tools automatically" -ForegroundColor Yellow
        Write-Host "[INFO] PyBullet may fail to install. If it does, install Visual Studio Build Tools manually:" -ForegroundColor Yellow
        Write-Host "       https://visualstudio.microsoft.com/visual-cpp-build-tools/" -ForegroundColor Yellow
    }
}

# -----------------------------------------------------------------------------
# Step 4: Download and Install ROS 2
# -----------------------------------------------------------------------------
Write-Host ""
Write-Host "[STEP 4/7] Installing ROS 2 Humble..." -ForegroundColor Green

$ros2Path = "C:\dev\ros2_humble"

if (-not (Test-Path $ros2Path)) {
    Write-Host "[INFO] Downloading ROS 2 Humble..." -ForegroundColor Yellow
    
    # Create installation directory
    New-Item -ItemType Directory -Force -Path "C:\dev" | Out-Null
    
    # Download ROS 2 binary
    $ros2Url = "https://github.com/ros2/ros2/releases/download/release-humble-20240523/ros2-humble-20240523-windows-release-amd64.zip"
    $ros2Zip = "$env:TEMP\ros2-humble.zip"
    
    Write-Host "[INFO] This may take a while (1-2 GB download)..." -ForegroundColor Yellow
    
    try {
        # Use BITS for more reliable download
        Start-BitsTransfer -Source $ros2Url -Destination $ros2Zip -DisplayName "Downloading ROS 2"
    } catch {
        # Fallback to Invoke-WebRequest
        Write-Host "[INFO] Using alternative download method..." -ForegroundColor Yellow
        Invoke-WebRequest -Uri $ros2Url -OutFile $ros2Zip
    }
    
    Write-Host "[INFO] Extracting ROS 2..." -ForegroundColor Yellow
    Expand-Archive -Path $ros2Zip -DestinationPath "C:\dev" -Force
    Remove-Item $ros2Zip
    
    Write-Host "[INFO] ROS 2 installed to $ros2Path" -ForegroundColor Green
} else {
    Write-Host "[INFO] ROS 2 already installed at $ros2Path" -ForegroundColor Green
}

# -----------------------------------------------------------------------------
# Step 5: Create Python Virtual Environment
# -----------------------------------------------------------------------------
Write-Host ""
Write-Host "[STEP 5/7] Creating Python virtual environment..." -ForegroundColor Green

# Remove existing venv if present
if (Test-Path $VenvDir) {
    Remove-Item -Recurse -Force $VenvDir
}

# Create virtual environment
python -m venv $VenvDir

Write-Host "[INFO] Virtual environment created at: $VenvDir" -ForegroundColor Green

# -----------------------------------------------------------------------------
# Step 6: Install Python Dependencies in venv
# -----------------------------------------------------------------------------
Write-Host ""
Write-Host "[STEP 6/7] Installing Python dependencies in virtual environment..." -ForegroundColor Green

# Activate venv and install packages
& "$VenvDir\Scripts\Activate.ps1"

python -m pip install --upgrade pip

$requirementsFile = Join-Path $ProjectRoot "requirements.txt"
if (Test-Path $requirementsFile) {
    pip install -r $requirementsFile
} else {
    pip install pybullet pyinstaller pillow
}

Write-Host "[INFO] Python packages installed in venv" -ForegroundColor Green

# -----------------------------------------------------------------------------
# Step 7: Create Activation Script
# -----------------------------------------------------------------------------
Write-Host ""
Write-Host "[STEP 7/7] Creating activation scripts..." -ForegroundColor Green

# Create batch file for cmd.exe
$activateBat = Join-Path $ProjectRoot "activate.bat"
@"
@echo off
REM =============================================================================
REM Drone Simulation - Environment Activation Script (Windows CMD)
REM =============================================================================
REM Usage: activate.bat
REM =============================================================================

echo Activating ROS 2 Humble...
call C:\dev\ros2_humble\local_setup.bat

echo Activating Python virtual environment...
call "%~dp0venv\Scripts\activate.bat"

echo.
echo [OK] Environment ready! You can now run:
echo   python simulation_host.py
echo   python ros_bridge.py
echo   python controllers.py
echo.
"@ | Out-File -FilePath $activateBat -Encoding ASCII

# Create PowerShell script
$activatePs1 = Join-Path $ProjectRoot "activate.ps1"
@"
# =============================================================================
# Drone Simulation - Environment Activation Script (Windows PowerShell)
# =============================================================================
# Usage: . .\activate.ps1
# =============================================================================

`$ScriptDir = Split-Path -Parent `$MyInvocation.MyCommand.Path

Write-Host "Activating ROS 2 Humble..." -ForegroundColor Yellow
& "C:\dev\ros2_humble\local_setup.ps1"

Write-Host "Activating Python virtual environment..." -ForegroundColor Yellow
& "`$ScriptDir\venv\Scripts\Activate.ps1"

Write-Host ""
Write-Host "[OK] Environment ready! You can now run:" -ForegroundColor Green
Write-Host "  python simulation_host.py"
Write-Host "  python ros_bridge.py"
Write-Host "  python controllers.py"
Write-Host ""
"@ | Out-File -FilePath $activatePs1 -Encoding UTF8

Write-Host "[INFO] Created activation scripts:" -ForegroundColor Green
Write-Host "       $activateBat (for CMD)" -ForegroundColor Gray
Write-Host "       $activatePs1 (for PowerShell)" -ForegroundColor Gray

# -----------------------------------------------------------------------------
# Verification
# -----------------------------------------------------------------------------
Write-Host ""
Write-Host "==============================================" -ForegroundColor Cyan
Write-Host "  Installation Complete!" -ForegroundColor Cyan
Write-Host "==============================================" -ForegroundColor Cyan
Write-Host ""
Write-Host "Verifying installation..." -ForegroundColor Yellow
Write-Host ""

try {
    python -c "import pybullet; print('  PyBullet: OK')"
} catch {
    Write-Host "  PyBullet: FAILED" -ForegroundColor Red
}

try {
    python -c "import PyInstaller; print('  PyInstaller: OK')"
} catch {
    Write-Host "  PyInstaller: FAILED" -ForegroundColor Red
}

try {
    python -c "import PIL; print('  Pillow: OK')"
} catch {
    Write-Host "  Pillow: FAILED" -ForegroundColor Red
}

Write-Host ""
Write-Host "==============================================" -ForegroundColor Cyan
Write-Host "  IMPORTANT: Activate the environment first!" -ForegroundColor Cyan
Write-Host "==============================================" -ForegroundColor Cyan
Write-Host ""
Write-Host "Before running any scripts, activate the environment:" -ForegroundColor Yellow
Write-Host "  CMD:        $activateBat" -ForegroundColor White
Write-Host "  PowerShell: . $activatePs1" -ForegroundColor White
Write-Host ""
Write-Host "Then run the simulation:" -ForegroundColor Yellow
Write-Host "  python simulation_host.py" -ForegroundColor White
Write-Host "  python ros_bridge.py" -ForegroundColor White
Write-Host "  python controllers.py" -ForegroundColor White
Write-Host ""