Param (
    [string]$DistroRepo="https://svsupport.blob.core.windows.net/files/VisionSense_Trainer/Packages",
    [string]$DistroName="VisionSense-PreviewTrainer",
    [string]$DistroHost="trainer.local"
)

function Check-EULA {
    Write-Host "Welcome to $DistroName Installer" -ForegroundColor DarkCyan
    Write-Host "This script will help you install $DistroName on your system."
    Write-Host "Please follow the prompts to complete the installation."
    Write-Host ""

    Write-Host "End User License Agreement (EULA)" -ForegroundColor DarkCyan
    Write-Host "Please read the EULA carefully before proceeding."
    Write-Host ""
    Write-Host "    $DistroRepo/visionsense-previewtrainer.eula.txt" -ForegroundColor Yellow
    Write-Host ""
    Write-Host "By continuing, you agree to the terms and conditions of the EULA."
    Write-Host ""

    $response = Read-Host "Do you accept the EULA? (y/N)"
    return $response -eq 'y' -or $response -eq 'Y'
}

if (-not (Check-EULA)) {
    Write-Host "You must accept the EULA to proceed with the installation. The installation will now abort." -ForegroundColor Red
    Write-Host ""
    exit 1
}

function Check-WSL {
    wsl --version > $null 2>&1
    return $LASTEXITCODE -eq 0
}

function Check-NVGPU {
    nvidia-smi > $null 2>&1
    return $LASTEXITCODE -eq 0
}

if (-not (Check-WSL)) {
    Write-Host "WSL is not installed. Please install WSL first." -ForegroundColor Red
    Write-Host "You can follow the instructions here: https://learn.microsoft.com/en-us/windows/wsl/install"
    Write-Host ""
    exit 1
}
if (-not (Check-NVGPU)) {
    Write-Host "NVIDIA GPU not detected. Please ensure you have a compatible NVIDIA GPU and the latest drivers installed." -ForegroundColor Red
    Write-Host "You can download the latest drivers from: https://www.nvidia.com/Download/index.aspx"
    Write-Host ""
    exit 1
}

function Setup-Distro {
    $localhosts = "$env:SystemRoot\System32\drivers\etc\hosts"
    $ipaddress = "127.0.0.1"

    if (-not (Select-String -Path $localhosts -Pattern "$DistroHost")) {
        Add-Content -Path $localhosts -Value "`r`n$ipaddress $DistroHost`r`n"
        Write-Host "Added DNS record for $DistroHost" -ForegroundColor Green
        Write-Host ""
    } else {
        Write-Host "DNS record for $DistroHost already exists" -ForegroundColor Yellow
        Resolve-DnsName $DistroHost
        Write-Host ""
    }
}

function Fetch-Distro {
    Write-Host "Download $DistroName" -ForegroundColor DarkCyan
    Write-Host ""
    curl.exe -L "$DistroRepo/visionsense-previewtrainer.sha512" -o "$env:TEMP\visionsense-previewtrainer.sha512"
    curl.exe -L "$DistroRepo/visionsense-previewtrainer.wsl" -o "$env:TEMP\visionsense-previewtrainer.wsl"
    Write-Host ""

    if (-not (Test-Path "$env:TEMP\visionsense-previewtrainer.sha512")) {return $false}
    if (-not (Test-Path "$env:TEMP\visionsense-previewtrainer.wsl")) {return $false}

    $expectedHash = Get-Content "$env:TEMP\visionsense-previewtrainer.sha512"
    $actualHash = (Get-FileHash -Path "$env:TEMP\visionsense-previewtrainer.wsl" -Algorithm SHA512).Hash
    Write-Host "$actualHash"
    $verificationResult = $expectedHash -eq $actualHash
    return $verificationResult
}

function Import-Distro {
    Write-Host "Import $DistroName" -ForegroundColor DarkCyan
    Write-Host "Remove existing distro if any"
    wsl --unregister $DistroName
    Write-Host "Create new distro"
    wsl --import $DistroName "$env:ProgramData\$DistroName" "$env:TEMP\visionsense-previewtrainer.wsl"
    $result = $LASTEXITCODE -eq 0
    Write-Host ""
    return $result
}

function Launch-Distro {
    Write-Host "Launch $DistroName" -ForegroundColor DarkCyan
    Write-Host ""
    wsl -d $DistroName -u cvmt --cd /opt/cvmt -- bash -c "docker compose up -d"
    Start-Sleep -Seconds 10
    wsl -d $DistroName -u cvmt --cd /opt/cvmt -- bash -c "docker compose down"
    wsl -d $DistroName -u cvmt --cd /opt/cvmt -- bash -c "docker compose up -d"
    return $LASTEXITCODE -eq 0
}

try {
    if (-not (Fetch-Distro)) {
        Write-Host "Checksum verification failed! The downloaded file may be corrupted." -ForegroundColor Red
        Write-Host ""
        exit 1
    } else {
        Write-Host "Checksum verification passed." -ForegroundColor Green
        Write-Host ""
    }
    if (-not (Import-Distro)) {
        Write-Host "Failed to import $DistroName" -ForegroundColor Red
        Write-Host ""
    } elseif (-not (Launch-Distro)) {
        Write-Host "Failed to launch $DistroName" -ForegroundColor Red
        Write-Host ""
    } else {
        Write-Host ""
        Setup-Distro
        Start-Sleep -Seconds 10

        Write-Host "Thank you for trying $DistroName!" -ForegroundColor DarkCyan
        Write-Host ""
        Write-Host "You can now access the trainer at http://$DistroHost"
        Write-Host ""

        Start-Process "http://$DistroHost"
        #wsl -d $DistroName -u cvmt --cd /opt/cvmt -- bash -c "docker compose logs -f"
    }
} catch {
    Write-Host "An error occurred during the setup process." -ForegroundColor Red
    Write-Host ""
} finally {
    Remove-Item "$env:TEMP\visionsense-previewtrainer.wsl"
    Remove-Item "$env:TEMP\visionsense-previewtrainer.sha512"
}