param (
	[Parameter(Mandatory=$true)]
	[string]$outputFilePath,
    [string]$rentryUrl = "https://rentry.org/ponyxl_loras_n_stuff"
)

function Get-RemoteFileSize {
    param (
        [string]$url
    )

    try {
		$ProgressPreference = 'SilentlyContinue'
        $response = Invoke-WebRequest -Uri $url -Method Head 
		$ProgressPreference = 'Continue'
        $fileSize = [int]$response.Headers['Content-Length']
        return $fileSize;
    } catch {
        Write-Error "Failed to retrieve file size. $_"
    }
}

function Create-Entry {
    param (
		[string]$name,
        [string]$url,
		[int]$size,
		[string] $comment,
		[int]$status
    )
    return [pscustomobject]@{Name=$name;Url=$url;Size=$size;Status=$status;Comment=$comment;Selected=0}
}

Clear-Host

#The path to the history file
$historyFile = Join-Path $outputFilePath "_history.txt"

# Initialize an empty dictionary for tracking the history
$historyDict = @{}

if (Test-Path $historyFile -PathType Leaf) {
	# Read the file line by line and populate the dictionary
    Get-Content -Path $historyFile | ForEach-Object {
        # Split the line into key and value based on whitespace
        $lineParts = $_ -split '\s+'

        # Ensure there are at least two parts on the line
        if ($lineParts.Length -ge 2) {
            $key = $lineParts[0]
			
			#$index = $_.IndexOf($lineParts[1])
			#$comment = $_.Substring($index + $lineParts[1].Length).Trim()
            #$value = @($lineParts[1], $comment)
			
            # Add an entry to the dictionary
            $historyDict[$key] = $lineParts[1]
        }
    }
}

Write-Host "Getting details about the Loras"

# Define a regular expression pattern to match URLs within <a> HTML elements ending with "safetensors"
$downloadLinkPattern = '<a [^>]*href=["'']?(https?://(?:www\.)?[^\s]+safetensors\b[^"''\s>]*)[^>]*>([^<]*)</a>([^<]*)'

# Download the content of the rentry.org page
$pageContent = Invoke-RestMethod -Uri $rentryUrl

# Find all matches using the regular expression pattern
$matches = $pageContent | Select-String -Pattern $downloadLinkPattern -AllMatches | ForEach-Object { $_.Matches }

# Initialize an empty list for tracking the available loras, each entry is a tuple with name, url, file size, integer indicating status (0 is new, 1 is updated), and integer indicating if user wants to download
$availableLoras = @()

# Process each match and extract download link, the first string inside <a> element, and the HTML content after </a>
foreach ($match in $matches) {
    $url = $match.Groups[1].Value
    $htmlAfterA = $match.Groups[3].Value
	
	# Trim the string and split it into two variables
	$trimmedString = $htmlAfterA.Trim()
	$firstWord = $trimmedString.Split(' ', 2)[0]
	$remainingText = $trimmedString.Split(' ', 2)[1]
	
	$outputFileWithPath = Join-Path "$outputFilePath" "ponyxl_$firstWord.safetensors"

	$comment = if($remainingText -ne $null) { $remainingText.Trim() } else { "" }
	
	if($historyDict[$firstWord] -eq $null -or -Not (Test-Path $outputFileWithPath -PathType Leaf)) {
		$fileSize = Get-RemoteFileSize -url $url
		
		#We don't have the lora yet, add it to the list of options
		$availableLoras += Create-Entry -name $firstWord -url $url -size $fileSize -comment $comment -status 0 
	} elseif(-Not ($historyDict[$firstWord] -eq $url)) {
		#The Lora was updated
		$fileSize = Get-RemoteFileSize -url $url
		
		$availableLoras += Create-Entry -name $firstWord -url $url -size $fileSize -comment $comment -status 1
	}
}


if($availableLoras.Count -eq 0) {
	Write-Host "There are no available updates based on your history file."
	Exit
}

Do {
	Clear-Host
	Write-Host "Below are the available LoRAs, asterisk indicates a LoRA is queued."
	$index = 1
	foreach ($availableLora in $availableLoras) {
		$name = $availableLora.Name
		$queued = if ($availableLora.Selected -eq 1) { "(*)" } else { "" }
		$size = $availableLora.Size / 1048576
		$status = if ($availableLora.Status -eq 0) { "New" } else { "Updated" }
		
		$size = $size.ToString("F2")
		
		Write-Host "[$index]) $queued $name - $status ($size MB)"
		$index++;
	}
	
	Write-Host "Enter a number to toggle queued status, a to queue all, or c to continue to download the queued LoRAs."
	$input = Read-Host "Command"
	
	if($input -eq "a") {
		foreach ($availableLora in $availableLoras) {
			$availableLora.Selected = 1
		}
	} elseif($input -ne "c") {
		$parsedInt = -1
		$success = [int]::TryParse($input, [ref]$parsedInt)
		if($success) {
			$availableLoras[$parsedInt - 1].Selected = ($availableLoras[$parsedInt - 1].Selected + 1) % 2
		}
	}
} while($input -ne "c")

if (-not (Test-Path $outputFilePath -PathType Container)) {
    # If the folder doesn't exist, create it
    New-Item -Path $outputFilePath -ItemType Directory
}

#Download the selected LoRAs
$downloadedLoras = 0
foreach ($availableLora in $availableLoras) {
	if($availableLora.Selected -eq 1) {
		try {
			$name = $availableLora.Name
			$url = $availableLora.Url.Trim()
			$comment = $availableLora.Comment
			$outputFileWithPath = Join-Path $outputFilePath "ponyxl_$name.safetensors"
			
			#if the file already exists, rename the old one before downloading the new one
			if (Test-Path $outputFileWithPath) {
				$counter = 1
				$newName = "ponyxl_$($name)_old_$counter.safetensors"
				$newPath = Join-Path $outputFilePath $newName

				while (Test-Path $newPath) {
					$newName = "ponyxl_$($name)_old_$counter.safetensors"
					$newPath = Join-Path $outputFilePath $newName
					$counter++
				}

				Rename-Item -LiteralPath $outputFileWithPath -NewName $newName
				Write-Host "Renamed old $name LoRA to: $newName"
			}
	
			Write-Host "Downloading lora $name with url $url"
			Invoke-WebRequest $url -OutFile $outputFileWithPath
		
			++$downloadedLoras
			$historyDict[$name] = $url
			
			if ($comment -ne $null -and $comment -ne "") {
				$commentFileWithPath = Join-Path "$outputFilePath" "ponyxl_$name.txt"
				# Clear the existing content of the output text file
				 if (Test-Path $commentFileWithPath -PathType Leaf) {
					Clear-Content -Path $commentFileWithPath

					# Write the new string to the text file
					Add-Content -Path $commentFileWithPath -Value $comment
				}
			}
		} catch {
			Write-Host "Error downloading lora $name, skipping, error was $_"
		}
	}
}

if (Test-Path $historyFile -PathType Leaf) {
	# Clear the existing content of the history file
	Clear-Content -Path $historyFile
}

#Write the new history
foreach ($entry in $historyDict.GetEnumerator()) {
	$entryString = "{0} {1}" -f $entry.Key, $entry.Value
	Add-Content -Path $historyFile -Value $entryString
}

Write-Host "Downloaded $downloadedLoras LoRAs."
Write-Host "Wrote the most recent urls for each model to _history.txt."