# SMART PRE-FLIGHT ENCODER - VERBOSE CONSOLE OUTPUT
# Usage: .\Encode-TVShows.ps1

param(
    [string]$TvDir = "Z:\tv",
    [string]$ContentDir = "Z:\content",
    [int]$MaxJobs = 2,               
    [int]$Av1Q = 34,                 
    [int]$HevcQ = 28,                
    [string]$EncoderAV1 = "av1_amf", 
    [string]$EncoderHEVC = "hevc_amf",
    [switch]$SkipAV1 = $true
)

# --- CONFIGURATION ---
$Global:FFMPEG = "ffmpeg"
$Global:FFPROBE = "ffprobe"
$Global:TEMP_DIR = "C:\Users\bnair\Videos\encodes"
$Global:LockDir = "C:\Users\bnair\Videos\encodes\locks"
$LogFileTV = "C:\Users\bnair\Videos\encodes\encoding-log-tv.csv"
$LogFileContent = "C:\Users\bnair\Videos\encodes\encoding-log-content.csv"

if (-not (Test-Path $Global:TEMP_DIR)) { New-Item -ItemType Directory -Path $Global:TEMP_DIR -Force | Out-Null }
if (-not (Test-Path $Global:LockDir)) { New-Item -ItemType Directory -Path $Global:LockDir -Force | Out-Null }

function Init-LogFile {
    param([string]$Path)
    if (-not (Test-Path $Path)) { "Timestamp,File,InputSize,OutputSize,CodecUsed,Status,Savings" | Out-File -FilePath $Path -Encoding UTF8 }
}
Init-LogFile $LogFileTV
Init-LogFile $LogFileContent

function Test-Tools {
    Write-Host "Setup: Checking required tools..." -ForegroundColor Cyan
    if (-not (Get-Command $Global:FFMPEG -ErrorAction SilentlyContinue)) { Write-Host "ERROR: ffmpeg not found!" -ForegroundColor Red; exit 1 }
    Write-Host "Setup: Tools found." -ForegroundColor Green
}

$SharedFunctions = {
    function Get-LockId {
        param([string]$FilePath)
        try {
            $pathBytes = [System.Text.Encoding]::UTF8.GetBytes($FilePath)
            $hash = [System.Security.Cryptography.SHA256]::Create().ComputeHash($pathBytes)
            return [BitConverter]::ToString($hash).Replace("-", "").Substring(0, 16)
        } catch { return "unknown_lock" }
    }

    function Run-FFmpeg {
        param($In, $Out, $Enc, $Q, $Seek=$null, $Duration=$null)
        $argsList = "-hide_banner -loglevel error -y"
        if ($Seek) { $argsList += " -ss $Seek" }
        $argsList += " -i `"$In`""
        if ($Duration) { $argsList += " -t $Duration" }
        $argsList += " -c:v $Enc -usage transcoding -quality quality -rc cqp -qp_i $Q -qp_p $Q -qp_b $Q -c:a copy `"$Out`""
        return Start-Process -FilePath "ffmpeg" -ArgumentList $argsList -NoNewWindow -Wait -PassThru
    }

    function Process-VideoFile {
        param($InputFile, $CurrentLogFile, $LockDir, $TempDir, $Av1Q, $HevcQ, $EncAV1, $EncHEVC, $SkipAV1)

        $pidStr = $PID.ToString()
        $lockId = Get-LockId -FilePath $InputFile
        $lockFile = Join-Path $LockDir "$lockId.lock"
        
        try {
            if (Test-Path $lockFile) { return } 
            $pidStr | Out-File -FilePath $lockFile -Force
            
            $fileName = Split-Path $InputFile -Leaf
            Write-Host "[$pidStr] Found: $fileName" -ForegroundColor White

            # Skip Logic
            $currentCodec = (& "ffprobe" -v error -select_streams v:0 -show_entries stream=codec_name -of csv=p=0 "$InputFile" 2>&1)
            if ($SkipAV1 -and ($currentCodec -match "av1" -or $currentCodec -match "hevc")) {
                Write-Host "[$pidStr] SKIP: Already optimized ($currentCodec)" -ForegroundColor DarkGray
                return
            }

            $inputSize = (Get-Item $InputFile).Length
            
            # --- PHASE 1: PRE-FLIGHT SAMPLE ---
            Write-Host "[$pidStr] Testing: Generating 60s sample..." -ForegroundColor Yellow
            $tempSample = Join-Path $TempDir "$fileName.sample.mkv"
            $procSample = Run-FFmpeg $InputFile $tempSample $EncAV1 $Av1Q "00:05:00" "60"
            
            $doFullAV1 = $true
            
            if ($procSample.ExitCode -eq 0 -and (Test-Path $tempSample)) {
                $sampleSize = (Get-Item $tempSample).Length
                # Threshold: 150MB for 60s is ~20Mbps (Likely bloat)
                if ($sampleSize -gt 150MB) {
                    Write-Host "[$pidStr] Test Result: FAIL. Sample was $([math]::Round($sampleSize/1MB))MB. Too big for AV1." -ForegroundColor Red
                    $doFullAV1 = $false
                } else {
                    Write-Host "[$pidStr] Test Result: PASS. Sample was $([math]::Round($sampleSize/1MB))MB. Proceeding with AV1." -ForegroundColor Green
                }
                Remove-Item $tempSample -Force
            }

            $finalStatus = "Failed"
            $finalCodec = "None"
            $finalSize = 0
            $finalSavings = 0.00
            
            # --- PHASE 2: FULL AV1 ---
            if ($doFullAV1) {
                Write-Host "[$pidStr] Action: Starting Full AV1 Encode..." -ForegroundColor Cyan
                $tempAV1 = Join-Path $TempDir "$fileName.av1.mkv"
                $procAV1 = Run-FFmpeg $InputFile $tempAV1 $EncAV1 $Av1Q
                
                if ($procAV1.ExitCode -eq 0 -and (Test-Path $tempAV1)) {
                    $sizeAV1 = (Get-Item $tempAV1).Length
                    
                    if ($sizeAV1 -lt $inputSize) {
                        $finalSavings = [math]::Round((1 - ($sizeAV1 / $inputSize)) * 100, 2)
                        Write-Host "[$pidStr] AV1 ACCEPTED: Saved ${finalSavings}%" -ForegroundColor Green
                        
                        $finalOut = $InputFile -replace '\.mkv$', '_av1.mkv' -replace '\.mp4$', '_av1.mp4'
                        Move-Item $tempAV1 -Destination $finalOut -Force
                        Remove-Item $InputFile -Force
                        Rename-Item $finalOut -NewName $fileName -Force
                        
                        "$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss'),`"$InputFile`",$inputSize,$sizeAV1,AV1,`"Replaced (AV1)`",$finalSavings%" | Out-File -FilePath $CurrentLogFile -Append -Encoding UTF8
                        return
                    } else {
                        Write-Host "[$pidStr] AV1 REJECTED: Larger than source ($([math]::Round($sizeAV1/1GB,2)) GB). Deleting..." -ForegroundColor Red
                        Remove-Item $tempAV1 -Force
                    }
                }
            }

            # --- PHASE 3: HEVC FALLBACK ---
            Write-Host "[$pidStr] Action: Trying HEVC Fallback..." -ForegroundColor Cyan
            $tempHEVC = Join-Path $TempDir "$fileName.hevc.mkv"
            $procHEVC = Run-FFmpeg $InputFile $tempHEVC $EncHEVC $HevcQ

            if ($procHEVC.ExitCode -eq 0 -and (Test-Path $tempHEVC)) {
                $sizeHEVC = (Get-Item $tempHEVC).Length
                
                if ($sizeHEVC -lt $inputSize) {
                    $finalSavings = [math]::Round((1 - ($sizeHEVC / $inputSize)) * 100, 2)
                    Write-Host "[$pidStr] HEVC ACCEPTED: Saved ${finalSavings}%" -ForegroundColor Green
                    
                    $finalOut = $InputFile -replace '\.mkv$', '_hevc.mkv' -replace '\.mp4$', '_hevc.mp4'
                    Move-Item $tempHEVC -Destination $finalOut -Force
                    Remove-Item $InputFile -Force
                    Rename-Item $finalOut -NewName $fileName -Force
                    
                    $finalStatus = "Replaced (HEVC)"
                    $finalCodec = "HEVC"
                    $finalSize = $sizeHEVC
                } else {
                    Write-Host "[$pidStr] HEVC REJECTED: Also larger. Keeping original." -ForegroundColor Red
                    Remove-Item $tempHEVC -Force
                    $finalStatus = "Rejected (Both Larger)"
                    $finalSize = $sizeHEVC
                    $finalCodec = "HEVC"
                }
            }

            "$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss'),`"$InputFile`",$inputSize,$finalSize,$finalCodec,`"$finalStatus`",$finalSavings%" | Out-File -FilePath $CurrentLogFile -Append -Encoding UTF8

        } finally {
            Remove-Item $lockFile -Force -ErrorAction SilentlyContinue
        }
    }
}

function Process-Directory {
    param($TargetDirectory, $TargetLogFile, $PhaseName)
    if (-not (Test-Path $TargetDirectory)) { return }

    Write-Host "`n=== PHASE: $PhaseName ===" -ForegroundColor Magenta
    $files = Get-ChildItem -Path $TargetDirectory -Include *.mkv, *.mp4 -Recurse -File -ErrorAction SilentlyContinue
    $videoFiles = $files | Where-Object { $_.Name -notmatch "_av1" -and $_.Name -notmatch "_hevc" }
    
    Write-Host "Found $($videoFiles.Count) files." -ForegroundColor Cyan

    $processed = 0
    while ($processed -lt $videoFiles.Count) {
        $batchSize = [math]::Min($MaxJobs, ($videoFiles.Count - $processed))
        $currentBatch = $videoFiles[$processed..($processed + $batchSize - 1)]
        $jobs = @()

        foreach ($file in $currentBatch) {
            $jobs += Start-Job -InitializationScript $SharedFunctions -ScriptBlock { 
                param($f, $log, $lock, $temp, $av1q, $hevcq, $e1, $e2, $skip) 
                Process-VideoFile $f $log $lock $temp $av1q $hevcq $e1 $e2 $skip
            } -ArgumentList $file.FullName, $TargetLogFile, $Global:LockDir, $Global:TEMP_DIR, $Av1Q, $HevcQ, $EncoderAV1, $EncoderHEVC, $SkipAV1
        }

        while (($jobs | Where-Object { $_.State -eq 'Running' }).Count -gt 0) {
            $jobs | Receive-Job
            Start-Sleep -Seconds 2
        }
        $jobs | Receive-Job
        $jobs | Remove-Job -Force
        $processed += $batchSize
        Write-Host "Progress Phase ${PhaseName}: $processed / $($videoFiles.Count)" -ForegroundColor Yellow
    }
}

Test-Tools
Process-Directory $TvDir $LogFileTV "TV-Shows"
Process-Directory $ContentDir $LogFileContent "Content"
Write-Host "`nDone." -ForegroundColor Green