<#
.SYNOPSIS
    Windows script to download, setup, validate and enforce Skyrim Together
.DESCRIPTION
    - Downloads and Validates the Skyrim Together package
    - Creates a shortcut on the desktop to play Skyrim Together
.EXAMPLE
    # Run script locally
    C:\PS> .\skyrim_together.ps1
.EXAMPLE
    # Run script remotely
    C:\PS> Set-ExecutionPolicy -ExecutionPolicy Bypass
    C:\PS> iex "& { $(irm https://wobig.tech/downloads/scripts/skyrim_together.ps1) }"
.EXAMPLE
    # Run script remotely (using older versions of powershell)
    C:\PS> Set-ExecutionPolicy -ExecutionPolicy Bypass
    C:\PS> Invoke-WebRequest -Uri "https://wobig.tech/downloads/scripts/skyrim_together.ps1" | Invoke-Expression
.EXAMPLE
    # Download the script locally then run
    C:\PS> Set-ExecutionPolicy -ExecutionPolicy Bypass
    C:\PS> Invoke-WebRequest -Uri "https://wobig.tech/downloads/scripts/skyrim_together.ps1" -OutFile "skyrim_together.ps1"
    C:\PS> powershell .\skyrim_together.ps1
.NOTES
    Author: Rick Wobig
    Source: https://wobig.tech/downloads/scripts/skyrim_together.ps1
#>

[CmdletBinding()]
param
(
    [Parameter(Mandatory = $false, HelpMessage = "Removes extra files or directories that don't match the expected manifest")]
    [Switch]
    [bool]$Clean = $false,

    [Parameter(Mandatory = $false, HelpMessage = "Whether to ignore free drive space where Skyrim is installed")]
    [Switch]
    [bool]$SkipDriveCheck = $false,

    [Parameter(Mandatory = $false, HelpMessage = "Displays a 'nice' animation instead of the 'standard' animation while doing things")]
    [Switch]
    [bool]$Nice = $false
)

# region Globals

$global:AnimationFrames = @("|","/","-","\","|")
if ($Nice){
    $global:AnimationFrames = @("          O","D         O","=D        O","==D       O","===D      O","====D     O","B====D    O"," B====D   O","  B====D  O","   B====D O","    B====DO","     B====O","      B===O","       B==O","        B=O","         BO","        B=O","       B==O","      B===O","     B====O","    B====DO","   B====D O","  B====D  O"," B====D   O","B====D    O","====D     O","===D      O","==D       O","=D        O","D         O","          O")
}

# endregion

# region Functions


function Get-SkyrimDirectory {
    param (
        [string]$CurrentDirectory
    )

    $skyrimPath = $null
    $exportedPath = Join-Path -Path $CurrentDirectory -ChildPath "skyrim_path"
    $defaultSkyrimPath = "C:\Program Files (x86)\Steam\SteamApps\common\Skyrim Special Edition\"
    $exampleSkyrimPath = "D:\Steam\SteamApps\common\Skyrim Special Edition\"

    if (Test-Path -Path (Join-Path -Path $defaultSkyrimPath -ChildPath "SkyrimSE.exe")){
        return $defaultSkyrimPath
    }

    if (Test-Path -Path $exportedPath){
        $skyrimPath = (Get-Content -Path $exportedPath).Trim()
    }

    if (-not $skyrimPath -and -not (Test-Path -Path $defaultSkyrimPath)){
        Write-Host -ForegroundColor Yellow "Skyrim wasn't found at the default path, please provide the path Skyrim exists at"
        Write-Host -ForegroundColor Yellow "  (The full path to the folder where SkyrimSE.exe exists)"
        Write-Host -ForegroundColor Yellow "  Example: $exampleSkyrimPath"
        $providedPath = Read-Host -Prompt "Skyrim Directory"

        if (-not $providedPath){
            Write-Host -ForegroundColor Red "You didn't provide a path, please try again"
            exit
        }

        if (-not (Test-Path -Path $providedPath)){
            Write-Host -ForegroundColor Red "The path you entered doesn't exist, please try again"
            Write-Host -ForegroundColor Red "  Provided Path: $providedPath"
            exit
        }

        $skyrimBinaryPath = Join-Path -Path $providedPath -ChildPath "SkyrimSE.exe"
        if (-not (Test-Path -Path $skyrimBinaryPath)){
            Write-Host -ForegroundColor Red "Provided Path: $providedPath"
            Write-Host -ForegroundColor Red "The path you provided is not a valid Skyrim install path, please try again"
            return $null
        }

        if (-not $providedPath.EndsWith("\")){
            $providedPath += "\"
        }

        $providedPath | Set-Content $exportedPath
        $skyrimPath = $providedPath
    }

    return $skyrimPath
}

function Get-AllPaths {
    param (
        [string]$path
    )
    $paths = Get-ChildItem -Path $path -Recurse -Force | Where-Object { $_.FullName -ne $captureFilePath } | ForEach-Object { $_.FullName }
    return $paths
}

function Invoke-ExpectedFiles {
    param (
        [string]$FilePath,
        [bool]$Capture = $true,
        [bool]$TakeAction = $false,
        [string]$expectedFilesUrl = $true
    )

    # Define the output file path for the dumped capture list
    $captureFilePath = Join-Path -Path $FilePath -ChildPath "skyrim_together_expected.txt"

    if ($Capture) {
        # Capture mode
        Write-Host "Capturing file and directory paths..."

        # Enumerate all files and directories and write to the capture file
        $allPaths = Get-AllPaths -path $FilePath
        $allPaths | Set-Content -Path $captureFilePath

        Write-Host "Capture completed. Paths written to $captureFilePath"
    }
    else {
        $expectedPaths = Get-FileContentFromUrl -Url $expectedFilesUrl

        if ($null -eq $expectedPaths) {
            Write-Host -ForegroundColor Red "Failed to get expected paths, please verify you have internet connectivity"
            exit
        }

        Write-Host -ForegroundColor Cyan "Ensuring expected Skyrim Together files and directories, this may take anywhere from 1 minute to 30 minutes based on your PC resources"

        # Enumerate current paths
        $currentPaths = Get-AllPaths -path $FilePath

        # Convert captured paths to absolute paths
        $absoluteCapturedPaths = (($expectedPaths.Split("`n")) | ForEach-Object { Join-Path -Path $FilePath -ChildPath $_ } | Where-Object { $_.Trim() -ne "" }) -join "`n"

        # Find paths that are not in the captured list
        $pathsToDelete = $currentPaths | Where-Object { -not $absoluteCapturedPaths.Contains($_) }

        # Initialize counters
        $filesDeleted = 0
        $directoriesDeleted = 0

        # Delete the paths that are not in the captured list
        foreach ($path in $pathsToDelete) {
            if (Test-Path -Path $path) {
                if ((Get-Item -Path $path).PSIsContainer) {
                    if ($TakeAction){
                        Remove-Item -Path $path -Recurse -Force
                    }
                    else {
                        Write-Host "Would have deleted directory: $path"
                    }
                    $directoriesDeleted++
                } else {
                    if ($TakeAction){
                        Remove-Item -Path $path -Force
                    }
                    else {
                        Write-Host "Would have deleted file: $path"
                    }
                    $filesDeleted++
                }
            }
        }

        Write-Host -ForegroundColor DarkGray "  $filesDeleted erroneous files cleaned up"
        Write-Host -ForegroundColor DarkGray "  $directoriesDeleted erroneous directories cleaned up"
    }
}

function Remove-FileIfExist {
    param (
        [string]$FilePath
    )

    if (Test-Path -Path $FilePath){
        Remove-Item -Path $FilePath -Force
        Write-Host -ForegroundColor DarkGray "Existing file at $FilePath has been deleted."
    }
}

function Invoke-FileDownloadFast {
    param (
        [string]$Url,
        [string]$SavePath
    )

    # NOTE: This function doesn't show a progress bar but downloads about 70x faster, definitely seems worth it to me!
    $ProgressPreference = 'SilentlyContinue'
    (New-Object Net.WebClient).DownloadFile($Url, $SavePath)
}

function Get-FileIntegrityHash {
    param (
        [string]$FilePath
    )

    # Check if the file exists
    if (-not (Test-Path -Path $FilePath)) {
        Write-Host "The specified file does not exist." -ForegroundColor Red
        return $null
    }

    try {
        # Create a SHA256 hash object
        $sha256 = [System.Security.Cryptography.SHA256]::Create()

        # Open the file and compute the hash
        $fileStream = [System.IO.File]::OpenRead($FilePath)
        $hashBytes = $sha256.ComputeHash($fileStream)
        $fileStream.Close()

        # Convert the hash bytes to a hexadecimal string
        $hashString = -join ($hashBytes | ForEach-Object { $_.ToString("x2") })

        return $hashString
    } catch {
        Write-Error "An error occurred while calculating the hash: $_"
        return $null
    }
}

function Get-FileContentFromUrl {
    param (
        [string]$Url
    )

    # Create a web client
    $webClient = New-Object System.Net.WebClient

    try {
        # Download the file content as a string
        $content = $webClient.DownloadString($Url)
        return $content.Trim()
    } catch {
        Write-Error "An error occurred while downloading the file: $_"
        return $null
    } finally {
        # Dispose of the web client
        $webClient.Dispose()
    }
}

function Invoke-ArchiveExtract {
    param (
        [string]$ZipFilePath,
        [string]$DestinationDirectory,
        [string]$ExpectedZipName
    )

    # Check if the provided file path matches the zip file name
    if ([System.IO.Path]::GetFileName($ZipFilePath) -ne $ExpectedZipName) {
        Write-Host -ForegroundColor Red "Invalid file, the provided file path does not match the required zip file name"
        return
    }

    # Check if the zip file exists
    if (-not (Test-Path -Path $ZipFilePath)) {
        Write-Host -ForegroundColor Red "The specified zip file does not exist"
        return
    }

    # Create the destination directory if it does not exist
    if (-not (Test-Path -Path $DestinationDirectory)) {
        New-Item -Path $DestinationDirectory -ItemType Directory | Out-Null
        Write-Host -ForegroundColor DarkGray "Created missing extract directory: $DestinationDirectory"
    }

    # Extract the zip file to the provided directory and overwrite existing files
    try {
        Write-Host -ForegroundColor Cyan "Expanding archive, this will take a few minutes..."
        Add-Type -AssemblyName System.IO.Compression.FileSystem
        Expand-Archive -Path $ZipFilePath -DestinationPath $DestinationDirectory -Force
        Write-Host -ForegroundColor Green "Archive has been successfully extracted to $DestinationDirectory"
    } catch {
        Write-Error "An error occurred while extracting the archive: $_"
    }
}

function Invoke-PackageValidation {
    param (
        [string]$ZipFilePath,
        [string]$IntegrityUrl
    )

    $validIntegrityHash = Get-FileContentFromUrl -Url $IntegrityUrl
    if ($null -eq $validIntegrityHash){
        Write-Host -ForegroundColor Red "Failed to get integrity hash, please verify you have internet connectivity"
        return $false
    }

    Write-Host -ForegroundColor DarkGray "Valid integrity hash: $validIntegrityHash"
    if (Test-Path -Path $ZipFilePath){
        $fileHash = Get-FileIntegrityHash -FilePath $ZipFilePath
        Write-Host -ForegroundColor DarkGray "File Hash:            $fileHash"
        if (-not ($validIntegrityHash -eq $fileHash)){
            return $false
        }
        else {
            return $true
        }
    }

    return $false
}

function Invoke-ShortcutCreate {
    param (
        [string]$ExecutablePath
    )

    # Get the logged-in user's desktop path
    $desktopPath = [System.Environment]::GetFolderPath("Desktop")

    # Create the shortcut path
    $shortcutPath = Join-Path -Path $desktopPath -ChildPath ([System.IO.Path]::GetFileNameWithoutExtension($ExecutablePath) + ".lnk")
    $startInPath = [System.IO.Path]::GetDirectoryName($ExecutablePath)

    Remove-FileIfExist -FilePath $shortcutPath

    # Use WSH to create the shortcut
    $wshShell = New-Object -ComObject WScript.Shell
    $shortcut = $wshShell.CreateShortcut($shortcutPath)
    $shortcut.TargetPath = $ExecutablePath
    $shortcut.WorkingDirectory = $startInPath
    $shortcut.Save()

    Write-Host -ForegroundColor Green "Shortcut created/updated at $shortcutPath"
}

function Get-FreeSpaceInMB {
    param (
        [string]$FilePath
    )

    # Get the drive letter from the file path
    $driveLetter = (Get-Item -Path $FilePath).PSDrive.Root

    # Get the drive info
    $driveInfo = Get-PSDrive -Name $driveLetter.Substring(0, 1)

    # Calculate free space in MB
    $freeSpaceMB = [math]::Round($driveInfo.Free / 1MB, 2)

    return $freeSpaceMB
}

function Get-DirectoryAndFileHashes {
    param (
        [string]$DirectoryPath
    )

    if (-not (Test-Path -Path $DirectoryPath)){
        Write-Host -ForegroundColor Red "Directory provided doesn't exist, unable to parse resources and hashes"
        Write-Host -ForegroundColor Red "  Directory: $DirectoryPath"
        return @()
    }

    $manifestHash = @{}

    # Get all files and directories recursively
    $entries = Get-ChildItem -Path $DirectoryPath -Recurse

    Write-Host ""
    $currentCount = 0
    foreach ($entry in $entries) {
        $currentCount += 1
        $animation = $($global:AnimationFrames)[$currentCount % $($global:AnimationFrames.Count)]
        Write-Host "`r $animation [$($currentCount)/$($entries.Count)]" -NoNewline -ForegroundColor Yellow

        $resourcePath = $entry.FullName.Replace($DirectoryPath, "")
        if ($entry.PSIsContainer) {
            # Add directory to the list
            $manifestHash[$resourcePath] = [PSCustomObject]@{
                Path = $resourcePath
                ItemType = 'Directory'
                Hash = $null
            }
        } else {
            # Compute the integrity hash for files
            $hash = Get-FileHash -Path $entry.FullName -Algorithm SHA256 | Select-Object -ExpandProperty Hash
            $manifestHash[$resourcePath] = [PSCustomObject]@{
                Path = $resourcePath
                ItemType = 'File'
                Hash = $hash
            }
        }
    }

    Write-Host ""
    Write-Host ""
    return $manifestHash
}

function Get-ResourcesToReacquire {
    param (
        [array]$ExpectedManifest,
        [hashtable]$CurrentManifest,
        [string]$SkyrimDirectory
    )

    if (-not (Test-Path -Path $SkyrimDirectory)){
        Write-Host -ForegroundColor Red "Skyrim directory doesn't exist, no way to validate expected manifest"
        Write-Host -ForegroundColor Red "  Skyrim Directory: $SkyrimDirectory"
        return @()
    }

    $reacquireResources = @()
    $currentCount = 0

    Write-Host ""
    foreach ($manifestResource in $ExpectedManifest) {
        $currentCount += 1
        $animation = $($global:AnimationFrames)[$currentCount % $($global:AnimationFrames.Count)]
        Write-Host "`r $animation [$($currentCount)/$($ExpectedManifest.Count)]" -NoNewline -ForegroundColor Yellow

        if ($CurrentManifest.ContainsKey($manifestResource.Path)){
            $expectedResource = $CurrentManifest[$manifestResource.Path]
        } else {
            $expectedResource = $null
        }

        if (-not $expectedResource -and $manifestResource.Hash -eq "original") {
            $resourceFullPath = Join-Path -Path $SkyrimDirectory -ChildPath $manifestResource.Path
            Write-Host -ForegroundColor Red "Skyrim native resource is missing, this likely means you need to verify the integrity of your game"
            Write-Host -ForegroundColor Red "  Resource: $resourceFullPath"
            continue
        }

        if ($manifestResource.Hash -eq "original"){
            continue
        }

        if (-not $expectedResource) {
            # If the item doesn't exist in the source directory
            $reacquireResources += [PSCustomObject]@{
                Path = $manifestResource.Path
                ItemType = $manifestResource.ItemType
                Hash = $manifestResource.Hash
            }
            continue
        }

        if ($manifestResource.ItemType -eq 'File' -and $manifestResource.Hash -ne $expectedResource.Hash) {
            # If it's a file and the hash is different
            $reacquireResources += [PSCustomObject]@{
                Path = $manifestResource.Path
                ItemType = $manifestResource.ItemType
                Hash = $manifestResource.Hash
            }
        }
    }

    Write-Host ""
    Write-Host ""
    return $reacquireResources
}

function Export-Results {
    param (
        [array]$Content,
        [string]$OutputFilePath
    )

    $jsonContent = $Content | ConvertTo-Json -Depth 3
    Set-Content -Path $OutputFilePath -Value $jsonContent
}

function Test-UrlExists {
    param (
        [string]$Url
    )

    try {
        # Send a HEAD request to the URL
        Invoke-WebRequest -Uri $Url -Method Head -ErrorAction Stop

        # If the request is successful, return true
        return $true
    }
    catch {
        # If an error occurs (e.g., 404 Not Found), return false
        return $false
    }
}

function Convert-ToUrlFriendly {
    param (
        [string]$ConvertContext
    )

    $urlPath = $ConvertContext -replace '\\', '/'
    $encodedUrlPath = [System.Web.HttpUtility]::UrlPathEncode($urlPath)

    return $encodedUrlPath
}

function Invoke-SkyrimResourcesFix {
    param (
        [array]$ResourcesToFix,
        [string]$RootUrl,
        [string]$SkyrimDirectory
    )

    if (-not (Test-Path -Path $SkyrimDirectory)){
        Write-Host -ForegroundColor Red "Skyrim directory doesn't exist, no way to fix resources for a non-existant directory"
        Write-Host -ForegroundColor Red "  Skyrim Directory: $SkyrimDirectory"
        return @()
    }

    $failedResourceFixes = @()

    $currentResourceIndex = 0
    foreach ($resource in $ResourcesToFix){
        try {
            $currentResourceIndex += 1
            $animation = $($global:AnimationFrames)[$currentResourceIndex % $($global:AnimationFrames.Count)]
            $resourceCountDisplay = "$($currentResourceIndex)/$($ResourcesToFix.Count)"
            Write-Host "`r $animation Fixing resource [$resourceCountDisplay]" -NoNewline -ForegroundColor Green

            $resourceFullPath = Join-Path -Path $SkyrimDirectory -ChildPath $resource.Path
            if ($resource.ItemType -eq "Directory"){
                New-Item -Path $resourceFullPath -ItemType Directory -Force | Out-Null
            } else {
                $parsedUriPath = Convert-ToUrlFriendly -ConvertContext $resource.Path
                $resourceUri = "$($RootUrl)$($parsedUriPath)"

                $fileDirectoryPath = Split-Path -Path $resourceFullPath
                if (-not (Test-Path -PathType Container -Path $fileDirectoryPath)){
                    New-Item -Path $fileDirectoryPath -ItemType Directory -Force | Out-Null
                }

                Invoke-FileDownloadFast -Url $resourceUri -SavePath $resourceFullPath
            }
        }
        catch {
            $failedResourceFixes += $resource
            Write-Host ""
            Write-Host ""
            Write-Host -ForegroundColor Red "Failure occurred fixing resource [$resourceCountDisplay]: $($PSItem.ToString())"
            Write-Host -ForegroundColor Red "  Resource: $($resource.Path)"
            Write-Host ""
        }
    }

    Write-Host ""
    Write-Host ""

    if ($failedResourceFixes.Count -ne 0){
        Write-Host -ForegroundColor Red "Failure to fix all corrupt or missing files, please check the output and try again"
        exit
    }
}

function Invoke-SkyrimPluginManifestEnforce {
    param (
        [string]$RootUrl
    )

    try {
        $skyrimPluginPath = Join-Path -Path $env:LOCALAPPDATA -ChildPath "Skyrim Special Edition\Plugins.txt"
        $skyrimLoadOrderPath = Join-Path -Path $env:LOCALAPPDATA -ChildPath "Skyrim Special Edition\loadorder.txt"

        $pluginUrl = "$($RootUrl)skyrim_plugins.txt"
        $loadOrderUrl = "$($RootUrl)skyrim_loadorder.txt"

        Invoke-FileDownloadFast -Url $pluginUrl -SavePath $skyrimPluginPath
        Invoke-FileDownloadFast -Url $loadOrderUrl -SavePath $skyrimLoadOrderPath

        Write-Host -ForegroundColor Green "Skyrim plugin manifest and order has been enforced successfully"
    }
    catch {
        Write-Host -ForegroundColor Red "Failure occurred enforcing mod plugin manifest: $($PSItem.ToString())"
        Write-Host -ForegroundColor Red "  Resource: $($env:LOCALAPPDATA)"
    }
}

# endregion

# region Script Execution


$currentDirectory = (Split-Path -Path $PSCommandPath)
$skyrimTogetherBinary = "Data\SkyrimTogetherReborn\SkyrimTogether.exe"
$rootUrl = "https://wobig.tech/downloads/package/"
$skyrimManifestUrl = "$($rootUrl)skyrim_together_manifest"
$neededSpaceMB = 52000

$skyrimDirectory = Get-SkyrimDirectory -CurrentDirectory $currentDirectory
if (-not $skyrimDirectory){
    Write-Host "Was unable to get a valid Skyrim Directory, please try again"
    exit
}

Write-Host -ForegroundColor Green "Found valid Skyrim directory @ $skyrimDirectory"

$freeDriveSpace = Get-FreeSpaceInMB -FilePath $skyrimDirectory
if (-not $SkipDriveCheck -and $freeDriveSpace -lt $neededSpaceMB){
    Write-Host -ForegroundColor Red "You don't have enough free space to install Skyrim Together, please clear some space on your drive and try again"
    Write-Host -ForegroundColor Red "  Free Space:   $($freeDriveSpace)MB"
    Write-Host -ForegroundColor Red "  Needed Space: $($neededSpaceMB)MB"
    Write-Host -ForegroundColor Red "  Skyrim Path:  $skyrimDirectory"
    exit
}

$skyrimManifest = Get-FileContentFromUrl -Url $skyrimManifestUrl
if (-not $skyrimManifest){
    Write-Host -ForegroundColor Red "Failure occurred getting the latest Skyrim Manifest, please check your internet connection and try again"
    exit
}

$expectedManifest = $skyrimManifest | ConvertFrom-Json

Write-Host -ForegroundColor Cyan "Parsing current local Skyrim resource manifest, this will take a minute or two depending on your PC hardware"
$currentSkyrimManifest = Get-DirectoryAndFileHashes -DirectoryPath $skyrimDirectory

Write-Host -ForegroundColor Cyan "Finished gathering local Skryim manifest, now validating missing and/or corrupt resources, this will take a few seconds"
# Compare directories and export the results
$reacquireResources = Get-ResourcesToReacquire -ExpectedManifest $expectedManifest -CurrentManifest $currentSkyrimManifest -SkyrimDirectory $skyrimDirectory

Invoke-SkyrimPluginManifestEnforce -RootUrl $rootUrl

if ($reacquireResources.Count -eq 0){
    Write-Host -ForegroundColor Green "All Skyrim resources exist and are not corrupt, all good to go!"
    exit
}

Write-Host -ForegroundColor Red "$($reacquireResources.Count) resources are corrupt or missing"
Write-Host -ForegroundColor Cyan "Acquiring and replacing the corrupt and/or missing files, this may take some time depending on how many resources need fixing"

Invoke-SkyrimResourcesFix -ResourcesToFix $reacquireResources -SkyrimDirectory $skyrimDirectory -RootUrl "$($rootUrl)skyrim/"

$skyrimTogetherPath = Join-Path -Path $skyrimDirectory -ChildPath $skyrimTogetherBinary
Invoke-ShortcutCreate -ExecutablePath $skyrimTogetherPath

Write-Host -ForegroundColor Green "All Skyrim resources exist and are not corrupt, all good to go!"


# endregion