# QSWAT+ installer / updater for Windows (per user, no administrator rights).
#
#   irm https://install.swat.tamu.edu/qsp.ps1 | iex
#
# Checks that QGIS is installed and that its Python has what QSWAT+ needs
# (Python 3.12 on x64, matplotlib, shapely 2), downloads the QSWAT+ setup from the
# download mirror, checks its SHA-256 and runs it silently. The setup puts the
# plugin into the default profile of QGIS 3 and of QGIS 4 (whichever are found)
# and the SWAT+ support data (Databases, TauDEM539Bin, Tools, ExampleDatasets)
# into %USERPROFILE%\SWATPlus, and QSWAT+ is enabled in those profiles. Re-running
# it updates QSWAT+ when the mirror has a newer version and leaves it alone
# otherwise.
#
# With options:
#
#   & ([scriptblock]::Create((irm https://install.swat.tamu.edu/qsp.ps1))) -WithSoils
#
#   -Check       only check the requirements; install nothing. Throws when one is
#                not met, returns normally when all are.
#   -WithSoils   also install the soils reference database (~40 MB), no prompt
#   -NoSoils     skip the soils database, no prompt
#   -Uninstall   remove the plugin from the QGIS profiles (the support data stays)
#
# QSWAT_SKIP_DEPCHECK=1 in the environment skips the QGIS Python check.
#
# Windows PowerShell 5.1 or PowerShell 7. Errors are thrown, never `exit`, which
# would close the window the command was pasted into.
param([switch]$Check, [switch]$WithSoils, [switch]$NoSoils, [switch]$Uninstall)

& {
$ErrorActionPreference = 'Stop'
$ProgressPreference = 'SilentlyContinue'   # Invoke-WebRequest is many times slower with it
[Net.ServicePointManager]::SecurityProtocol = [Net.ServicePointManager]::SecurityProtocol -bor [Net.SecurityProtocolType]::Tls12

# Download mirrors, tried in order.
$Bases = @(
    'https://install.swat.tamu.edu/downloads'
    'https://install.swatplus.link/downloads'
)
$Feed = 'qswatplus/windows'
$SetupName = 'QSWATPlus-win-x64-Setup.exe'
$DatasetsRel = '4.0/swatplus_datasets.sqlite'
$SoilsZip = 'swatplus_soils.zip'
$DbDir = Join-Path $env:USERPROFILE 'SWATPlus\Databases'
$UninstallRoot = 'HKCU:\Software\Microsoft\Windows\CurrentVersion\Uninstall'
# the setup's uninstall entry (QSWATPlus.nsi)
$OwnKey = "$UninstallRoot\QSWATPlus"
# the Python the Windows modules are built for (cp312-win_amd64)
$PyWant = '3.12'

function Say([string]$m)  { Write-Host '==> ' -ForegroundColor Cyan -NoNewline; Write-Host $m }
function Warn([string]$m) { Write-Host '!!  ' -ForegroundColor Yellow -NoNewline; Write-Host $m }

# A mirror can answer a missing file with a 200 and an HTML page, so every
# download is checked for what it should be before it is used.
function Test-Download([string]$kind, [string]$file) {
    $head = [byte[]]::new(16)
    $fs = [IO.File]::OpenRead($file)
    try { $n = $fs.Read($head, 0, 16) } finally { $fs.Close() }
    switch ($kind) {
        'exe'     { return $n -ge 2 -and $head[0] -eq 0x4D -and $head[1] -eq 0x5A }
        'zip'     { return $n -ge 2 -and $head[0] -eq 0x50 -and $head[1] -eq 0x4B }
        'sqlite'  { return $n -eq 16 -and [Text.Encoding]::ASCII.GetString($head, 0, 15) -eq 'SQLite format 3' }
        'sha256'  { return (Get-Content -Raw $file) -match '^[0-9a-fA-F]{64}' }
        'version' { return (Get-Content -Raw $file) -match '^v?\d+\.\d+' }
    }
    return $false
}

function Get-FromMirror([string]$rel, [string]$out, [string]$kind) {
    foreach ($base in $Bases) {
        try { Invoke-WebRequest -UseBasicParsing -Uri "$base/$rel" -OutFile $out -TimeoutSec 900 } catch { continue }
        if (Test-Download $kind $out) { return $true }
    }
    return $false
}

# QGIS installs that can run plugins: a QGIS folder in Program Files (the
# standalone installer) or C:\OSGeo4W, with bin\python-qgis*.bat, which sets up
# QGIS's environment and starts its Python. Folders left behind by an uninstalled
# QGIS have no such file and are skipped.
function Get-QgisInstall {
    $dirs = @()
    foreach ($pf in @($env:ProgramW6432, $env:ProgramFiles) | Where-Object { $_ } | Select-Object -Unique) {
        if (Test-Path $pf) { $dirs += @(Get-ChildItem $pf -Directory -Filter 'QGIS*' -ErrorAction SilentlyContinue) }
    }
    if (Test-Path 'C:\OSGeo4W') { $dirs += Get-Item 'C:\OSGeo4W' }
    foreach ($d in $dirs) {
        $bat = Get-ChildItem (Join-Path $d.FullName 'bin') -Filter 'python-qgis*.bat' -ErrorAction SilentlyContinue |
            Sort-Object Name -Descending | Select-Object -First 1
        if ($bat) { [pscustomobject]@{ Name = $d.Name; Dir = $d.FullName; Bat = $bat.FullName } }
    }
}

# Runs the probe with one QGIS's Python and returns its key:value lines as a
# hashtable, empty when that Python does not run. cmd handles the redirection,
# so nothing the probe writes to stderr reaches PowerShell.
function Invoke-QgisProbe($q, [string]$probe, [string]$outFile) {
    $cmdArgs = '/d /c ""' + $q.Bat + '" "' + $probe + '" > "' + $outFile + '" 2>&1"'
    $null = Start-Process -FilePath $env:ComSpec -ArgumentList $cmdArgs -Wait -NoNewWindow -PassThru
    $r = @{}
    if (Test-Path $outFile) {
        foreach ($line in Get-Content $outFile) {
            if ($line -match '^(python|platform|qgis|missing|numpy):(.*)$') { $r[$Matches[1]] = $Matches[2].Trim() }
        }
    }
    return $r
}

# Every QGIS found, each with the problems that keep QSWAT+ from running in it.
# QSWAT+ can be installed when at least one has none.
function Test-Qgis([string]$tmp) {
    $probe = Join-Path $tmp 'probe.py'
    Set-Content -Path $probe -Encoding ascii -Value @(
        'import sys, sysconfig',
        'missing = []',
        'try:',
        '    import matplotlib',
        'except Exception:',
        '    missing.append("matplotlib")',
        'try:',
        '    import shapely',
        '    v = tuple(int(x) for x in (shapely.__version__.split(".") + ["0", "0"])[:2])',
        '    if v < (2, 0):',
        '        missing.append("shapely>=2.0 (found %s)" % shapely.__version__)',
        'except Exception:',
        '    missing.append("shapely")',
        'print("python:" + sys.version.split()[0])',
        'print("platform:" + sysconfig.get_platform())',
        'print("missing:" + ", ".join(missing))',
        'try:',
        '    import numpy',
        '    print("numpy:" + numpy.__version__)',
        'except Exception:',
        '    pass',
        'try:',
        '    from qgis.core import Qgis',
        '    print("qgis:" + Qgis.version().split("-")[0])',   # the number; the release name may not be ASCII
        'except Exception:',
        '    pass'
    )
    $i = 0
    foreach ($q in @(Get-QgisInstall)) {
        $i++
        $r = Invoke-QgisProbe $q $probe (Join-Path $tmp "probe-$i.txt")
        $problems = @()
        if (-not $r.ContainsKey('missing')) {
            $problems += "its Python did not run ($($q.Bat))"
        } else {
            if ($r['platform'] -ne 'win-amd64') { $problems += "its Python is $($r['platform']); QSWAT+ has Windows modules for x64 only" }
            if ($r['python'] -notlike "$PyWant.*") { $problems += "its Python is $($r['python']); QSWAT+ has Windows modules for Python $PyWant only" }
            if ($r['missing']) { $problems += "its Python lacks $($r['missing'])" }
        }
        # QGIS 3 or 4: from qgis.core when it loaded, else from the folder name
        $major = if ($r['qgis'] -match '^(\d+)\.') { $Matches[1] } elseif ($q.Name -match '^QGIS[ -]?(\d+)') { $Matches[1] } else { '' }
        [pscustomobject]@{ Name = $q.Name; Dir = $q.Dir; Bat = $q.Bat; Major = $major; Version = $r['qgis']
                           Python = $r['python']; Numpy = $r['numpy']; Problems = $problems }
    }
}

# Uninstall entries of QSWAT+ and of the SWAT+ support data written by other
# installers (TAMU's Inno Setup ones), whose uninstallers delete files this setup
# also installs.
function Get-OtherInstall {
    if (-not (Test-Path $UninstallRoot)) { return }
    Get-ChildItem $UninstallRoot | Where-Object { $_.PSChildName -ne 'QSWATPlus' } |
        ForEach-Object { Get-ItemProperty $_.PSPath } |
        Where-Object { "$($_.DisplayName)" -match '^QSWAT' -or "$($_.DisplayName)" -like 'SWAT+ Tools*' }
}

# Sets QSWATPlus=true under [PythonPlugins] and returns created, enabled or
# unchanged. The file is read and written as Latin-1, which maps every byte to one
# character and back, so its UTF-8 text, line endings and order stay as they are;
# a file that already has QSWATPlus=true is not written. A new file gets CRLF line
# endings, as QGIS writes them on Windows.
function Enable-QswatPlusSetting([string]$Ini) {
    $latin1 = [Text.Encoding]::GetEncoding(28591)
    if (-not (Test-Path -LiteralPath $Ini)) {
        New-Item -ItemType Directory -Force -Path (Split-Path -Parent $Ini) | Out-Null
        [IO.File]::WriteAllBytes($Ini, $latin1.GetBytes("[PythonPlugins]`r`nQSWATPlus=true`r`n"))
        return 'created'
    }
    $text = $latin1.GetString([IO.File]::ReadAllBytes($Ini))
    $eol = if ($text.Contains("`r`n")) { "`r`n" } elseif ($text.Contains("`n")) { "`n" } else { "`r`n" }
    $hdr = [regex]::Match($text, '(?m)^\[PythonPlugins\]\r?$')
    if (-not $hdr.Success) {
        # no section: append one, after a blank line
        $sep = ''
        if ($text.Length -gt 0) { if (-not $text.EndsWith("`n")) { $sep += $eol }; $sep += $eol }
        $new = $text + $sep + '[PythonPlugins]' + $eol + 'QSWATPlus=true' + $eol
    } else {
        $start = $hdr.Index + $hdr.Length          # end of the header, before its line break
        $rest = $text.Substring($start)
        $next = [regex]::Match($rest, '(?m)^\[')
        $body = if ($next.Success) { $rest.Substring(0, $next.Index) } else { $rest }
        $key = [regex]::Match($body, '(?m)^[ \t]*QSWATPlus[ \t]*=([^\r\n]*)')
        if ($key.Success) {
            if ($key.Groups[1].Value.Trim() -eq 'true') { return 'unchanged' }
            # that line only; its line break stays
            $at = $start + $key.Index
            $new = $text.Substring(0, $at) + 'QSWATPlus=true' + $text.Substring($at + $key.Length)
        } else {
            # right after the header line; a header on the last line gets a line break first
            $brk = [regex]::Match($rest, '^\r?\n')
            if ($brk.Success) { $at = $start + $brk.Length; $ins = 'QSWATPlus=true' + $eol }
            else { $at = $start; $ins = $eol + 'QSWATPlus=true' + $eol }
            $new = $text.Substring(0, $at) + $ins + $text.Substring($at)
        }
    }
    [IO.File]::WriteAllBytes($Ini, $latin1.GetBytes($new))
    return 'enabled'
}

function Get-PluginVersion([string]$major) {
    $meta = Join-Path $env:APPDATA "QGIS\QGIS$major\profiles\default\python\plugins\QSWATPlus\metadata.txt"
    if (-not (Test-Path $meta)) { return $null }
    $line = Select-String -Path $meta -Pattern '^version=' | Select-Object -First 1
    if ($line) { return $line.Line.Substring(8).Trim() }
}

if ($Uninstall) {
    $own = Get-ItemProperty $OwnKey -ErrorAction SilentlyContinue
    if (-not $own) { Warn 'QSWAT+ from this installer is not installed.'; return }
    if (Get-Process -Name 'qgis*' -ErrorAction SilentlyContinue) { throw 'QGIS is running. Close it, then run this again.' }
    $exe = $own.UninstallString.Trim('"')
    Say "Removing QSWAT+ $($own.DisplayVersion)"
    # An NSIS uninstaller copies itself to %TEMP% and returns at once unless it is
    # given _?=<its folder>; then it cannot delete itself, so that is done here.
    $dir = Split-Path $exe
    $p = Start-Process -FilePath $exe -ArgumentList "/S _?=$dir" -Wait -PassThru
    if ($p.ExitCode -ne 0) { throw "The QSWAT+ uninstaller failed (exit code $($p.ExitCode))." }
    Remove-Item -Force $exe -ErrorAction SilentlyContinue
    Remove-Item $dir -ErrorAction SilentlyContinue
    if (Test-Path $OwnKey) { throw 'The uninstaller finished, but Windows still lists QSWAT+ as installed.' }
    Say 'Done. The support data in %USERPROFILE%\SWATPlus was left in place.'
    return
}

$tmp = Join-Path ([IO.Path]::GetTempPath()) ('swatplus-' + [guid]::NewGuid().ToString('N'))
New-Item -ItemType Directory -Path $tmp | Out-Null
try {
    Say "QSWAT+ installer (Windows, $env:PROCESSOR_ARCHITECTURE)"

    # --- requirements: a QGIS whose Python can run QSWAT+, and an installer on a mirror
    $usable = @()
    if ($env:QSWAT_SKIP_DEPCHECK -eq '1') {
        Warn 'Skipping the QGIS Python check (QSWAT_SKIP_DEPCHECK=1).'
    } else {
        $all = @(Test-Qgis $tmp)
        if ($all.Count -eq 0) { throw 'QGIS was not found. Install QGIS (3.40 or newer) from https://qgis.org, then run this again.' }
        foreach ($q in $all) {
            $label = if ($q.Version) { "QGIS $($q.Version) ($($q.Dir))" } else { $q.Dir }
            if ($q.Problems.Count -eq 0) {
                Say "$label`: Python $($q.Python), numpy $($q.Numpy), matplotlib and shapely present"
                $usable += $q
            } else {
                Warn "$label`: QSWAT+ will not run here: $($q.Problems -join '; ')."
            }
        }
        if ($usable.Count -eq 0) {
            throw 'No QGIS on this computer can run QSWAT+ (see above). Install QGIS 3.40 or newer from https://qgis.org; with the OSGeo4W installer also add the python3-matplotlib and python3-shapely packages.'
        }
    }
    $shaFile = Join-Path $tmp 'setup.sha256'
    if (-not (Get-FromMirror "$Feed/$SetupName.sha256" $shaFile 'sha256')) { throw 'Could not find the QSWAT+ Windows setup on any mirror.' }
    foreach ($o in @(Get-OtherInstall)) {
        Warn "$($o.DisplayName) from another installer is registered in $($o.InstallLocation). QSWAT+ installs over the same files, and that installer's uninstaller would delete them; remove '$($o.DisplayName)' in Settings > Apps first to avoid this."
    }
    if ($Check) {
        Say 'QSWAT+ requirements met.'
        return
    }

    # --- install or update
    $verFile = Join-Path $tmp 'version.txt'
    if (-not (Get-FromMirror "$Feed/version.txt" $verFile 'version')) { throw 'Could not read the QSWAT+ version from any mirror.' }
    $version = (Get-Content -Raw $verFile).Trim()

    $own = Get-ItemProperty $OwnKey -ErrorAction SilentlyContinue
    # up to date only when the registered version is the mirror's and the plugin
    # is there in the profile of every usable QGIS
    $majors = @($usable | ForEach-Object { $_.Major } | Where-Object { $_ } | Select-Object -Unique)
    $stale = @($majors | Where-Object { (Get-PluginVersion $_) -ne $version })
    if ($own -and $own.DisplayVersion -eq $version -and $stale.Count -eq 0) {
        Say "QSWAT+ $version is up to date."
    } else {
        if (Get-Process -Name 'qgis*' -ErrorAction SilentlyContinue) {
            throw 'QGIS is running. Close it, then run this again.'
        }
        $want = ((Get-Content -Raw $shaFile) -split '\s+')[0].ToLower()
        $from = if ($own) { "$($own.DisplayVersion) -> " } else { '' }
        Say "Downloading QSWAT+ $from$version"
        $setup = Join-Path $tmp $SetupName
        if (-not (Get-FromMirror "$Feed/$SetupName" $setup 'exe')) { throw 'Could not download the QSWAT+ setup from any mirror.' }
        if ((Get-FileHash $setup -Algorithm SHA256).Hash.ToLower() -ne $want) { throw "Checksum mismatch for $SetupName." }
        Say 'Checksum OK; installing'
        $p = Start-Process -FilePath $setup -ArgumentList '/S' -Wait -PassThru
        # exit code 2: the setup found no QGIS 3 or 4, neither a profile under
        # %APPDATA%\QGIS nor a QGIS 3.x or 4.x folder in Program Files
        if ($p.ExitCode -eq 2) {
            throw 'The QSWAT+ setup found no QGIS 3 or QGIS 4 (no profile in %APPDATA%\QGIS and no QGIS folder in Program Files), so it installed nothing. Install QGIS from https://qgis.org and start it once, then run this again.'
        }
        if ($p.ExitCode -ne 0) { throw "The QSWAT+ setup failed (exit code $($p.ExitCode))." }
        $now = Get-ItemProperty $OwnKey -ErrorAction SilentlyContinue
        if (-not $now -or $now.DisplayVersion -ne $version) { throw "The QSWAT+ setup finished, but Windows does not list QSWAT+ $version as installed." }
        foreach ($m in $majors) {
            $v = Get-PluginVersion $m
            if ($v -ne $version) { throw "The QSWAT+ setup finished, but the QGIS $m profile has plugin version '$v'." }
            Say "Installed QSWAT+ $version -> $env:APPDATA\QGIS\QGIS$m\profiles\default\python\plugins\QSWATPlus"
        }
    }

    # Enable QSWAT+ in every default profile that has it. The setup does this as
    # well; here it also covers a run where QSWAT+ was already up to date.
    $enabled = $false
    if (Get-Process -Name 'qgis*' -ErrorAction SilentlyContinue) {
        Warn 'QGIS is running, so QSWAT+ was not enabled. Enable it in Plugins > Manage and Install Plugins, or close QGIS and run this again.'
    } else {
        $enabled = $true
        foreach ($m in '3', '4') {
            if (-not (Get-PluginVersion $m)) { continue }
            $ini = Join-Path $env:APPDATA "QGIS\QGIS$m\profiles\default\QGIS\QGIS$m.ini"
            try {
                if ((Enable-QswatPlusSetting $ini) -eq 'unchanged') { Say "QSWAT+ already enabled in $ini" }
                else { Say "Enabled QSWAT+ in $ini" }
            } catch {
                Warn "Could not enable QSWAT+ in $ini ($($_.Exception.Message)); enable it in Plugins > Manage and Install Plugins."
                $enabled = $false
            }
        }
    }

    New-Item -ItemType Directory -Force -Path $DbDir | Out-Null
    # The datasets database is refreshed every time, as qsp.sh does: projects are
    # created from it, so an old schema breaks New Project.
    $ds = Join-Path $tmp 'datasets.sqlite'
    if (Get-FromMirror $DatasetsRel $ds 'sqlite') {
        Move-Item -Force $ds (Join-Path $DbDir 'swatplus_datasets.sqlite')
        Say "Datasets database -> $DbDir\swatplus_datasets.sqlite"
    } else {
        Warn 'Could not download the datasets database from any mirror; kept the copy from the setup.'
    }

    $soilsDb = Join-Path $DbDir 'swatplus_soils.sqlite'
    if (Test-Path $soilsDb) {
        Say "Soils database already present: $soilsDb"
    } else {
        $wantSoils = if ($WithSoils) { $true } elseif ($NoSoils -or -not [Environment]::UserInteractive) { $false }
                     else { (Read-Host 'Also install the SWAT+ soils reference database (~40 MB)? [y/N]') -match '^(y|yes)$' }
        if ($wantSoils) {
            Say 'Downloading the soils database (~40 MB)'
            $zip = Join-Path $tmp 'soils.zip'
            if (-not (Get-FromMirror $SoilsZip $zip 'zip')) { throw "Could not download $SoilsZip from any mirror." }
            Expand-Archive -Path $zip -DestinationPath (Join-Path $tmp 'soils') -Force
            $src = Get-ChildItem (Join-Path $tmp 'soils') -Recurse -Filter 'swatplus_soils.sqlite' | Select-Object -First 1
            if (-not $src) { throw 'The soils archive has no swatplus_soils.sqlite.' }
            Move-Item -Force $src.FullName $soilsDb
            Say "Soils database -> $soilsDb"
        } else {
            Say 'Skipping the soils database (add it later with -WithSoils).'
        }
    }
    if ($enabled) { Say 'Start QGIS; QSWAT+ is enabled and in the Plugins menu.' }
    else { Say 'Start QGIS, then enable QSWATPlus in Plugins > Manage and Install Plugins.' }
} finally {
    Remove-Item -Recurse -Force $tmp -ErrorAction SilentlyContinue
}
}