# Variables
$ffmpegPath = (Get-Command ffmpeg -ErrorAction SilentlyContinue).Path
$ffprobePath = (Get-Command ffprobe -ErrorAction SilentlyContinue).Path
$inputDir = "./input"
$outputDir = "./output"

# Functions
# Function to get the bitrate of a video file using ffprobe
function Get-Bitrate {
    param (
        [string]$filePath
    )

    # Get bitrate from ffprobe
    $bitrateOutput = & $ffprobePath -v error -select_streams v:0 -show_entries stream=bit_rate -of default=noprint_wrappers=1:nokey=1 $filePath

    if ([string]::IsNullOrEmpty($bitrateOutput) -or $bitrateOutput -eq "N/A") {
        # If direct bitrate is not available, calculate bitrate based on file size and duration
        $fileSize = (Get-Item $filePath).Length
        $durationOutput = & $ffprobePath -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 $filePath

        if ($fileSize -gt 0 -and $durationOutput -ne "N/A" -and $durationOutput -ne "") {
            $duration = [float]$durationOutput
            $bitrate = [math]::Round(($fileSize * 8) / $duration / 1000) # in kbps
            return $bitrate
        } else {
            Write-Host "Unable to determine bitrate for $filePath."
            return 1000
        }
    }

    return [int]$bitrateOutput
}

# Function to get the duration of a video file using ffprobe
function Get-Duration {
    param (
        [string]$filePath
    )

    $durationOutput = & $ffprobePath -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 $filePath

    if ($durationOutput -ne "N/A" -and $durationOutput -ne "") {
        return [float]$durationOutput
    } else {
        Write-Host "Unable to determine duration for $filePath."
        return $null
    }
}

# Function to compress the video file until it is below the specified size
function Compress-File {
    param (
        [string]$inputFilePath,
        [string]$outputFilePath,
        [int]$initialBitrate,
        [int]$maxSizeKB
    )

    $bitrate = $initialBitrate
    $originalDuration = Get-Duration -filePath $file.FullName
    $speedFactor = $originalDuration / 3
    $outputFilePathTemp = "$outputFilePath.webm"

    if ($speedFactor -le 1) {
        $speedFactor = 1
    }

    do {
        Write-Host Trying bitrate: ${bitrate}k

        # Compress with the current bitrate
        & $ffmpegPath -loglevel quiet -i $inputFilePath -t 3 -r 25 -filter:v "setpts=PTS/$speedFactor,scale=-1:512" -c:v libvpx-vp9 -b:v ${bitrate}k $outputFilePathTemp -y

        # Check the size of the compressed file
        $fileSize = (Get-Item $outputFilePathTemp).Length
        $fileSizeKB = [math]::Round($fileSize / 1KB)

        if ($fileSizeKB -le $maxSizeKB) {
            # Rename the temporary file to the final output file name
            if (Test-Path $outputFilePath) { Remove-Item -Path $outputFilePath -Force }
            Rename-Item -Path $outputFilePathTemp -NewName (Split-Path $outputFilePath -Leaf)
            return
        }

        $bitrateOffset = 100
        if (($fileSizeKB / $maxSizeKB) -gt 2) {
            $bitrateOffset = 100 * ($fileSizeKB / $maxSizeKB)
        }

        # Reduce the bitrate
        $bitrate = [math]::Max([int]($bitrate - $bitrateOffset), 100) # Ensuring bitrate doesn't go below 100k

    } while ($bitrate -gt 100)

    # If bitrate falls too low, retain the last generated file
    if (Test-Path $outputFilePathTemp) {
        if (Test-Path $outputFilePath) { Remove-Item -Path $outputFilePath -Force }
        Rename-Item -Path $outputFilePathTemp -NewName (Split-Path $outputFilePath -Leaf)
    }
}

# Check if ffmpeg and ffprobe executables exist
if (-not $ffmpegPath -or -not $ffprobePath) {
    Write-Host "ffmpeg or ffprobe is not installed on this system. Exiting..."
    exit 1
}

# Check if the output directory exists, if not, create it
if (-not (Test-Path $outputDir)) {
    New-Item -Path $outputDir -ItemType Directory -Force | Out-Null
}

# Process each file in the input directory
foreach ($file in Get-ChildItem -Path $inputDir -File) {
    $outputFile = "$outputDir/$($file.BaseName).webm"

    # Check if the file is a PNG or GIF
    if ($file.Extension -notin ".png", ".gif") {
        continue
    }

    Write-Host "Processing $($file.FullName)..."

    # Convert the file to WebM format
    if ($file.Extension -eq ".png") {
        & $ffmpegPath -loglevel quiet -i $file.FullName -r 25 -vf "scale=-1:512" -c:v libvpx-vp9 -pix_fmt rgba $outputFile -y
    }

    if ($file.Extension -eq ".gif") {
        # Get the original duration of the GIF
        $originalDuration = Get-Duration -filePath $file.FullName

        if ($originalDuration -gt 0) {
            $speedFactor = $originalDuration / 3

            if ($speedFactor -le 1) {
                $speedFactor = 1
            }

            # Convert the file to WebM format
            & $ffmpegPath -loglevel quiet -i $file.FullName -t 3 -r 25 -filter:v "setpts=PTS/$speedFactor,scale=-1:512" -c:v libvpx-vp9 $outputFile -y
        } else {
            Write-Host "Invalid duration for GIF $($file.FullName). Skipping..."
            continue
        }

        # Check if the file size exceeds 256 KB (262144 bytes)
        $fileSize = (Get-Item $outputFile).Length

        if ($fileSize -gt 262144) {
            Write-Host "File $($outputFile) exceeds 256 KB, compressing..."

            # Get the initial bitrate of the WebM file
            $initialBitrate = Get-Bitrate -filePath $outputFile

            # Compress the file until it is below 256 KB
            Compress-File -inputFilePath $file.FullName -outputFilePath $outputFile -initialBitrate $initialBitrate -maxSizeKB 256
        }
    }
}