[CmdletBinding()]
param(
    [int[]] $Milestones = @(100, 300, 500),
    [int] $PollSeconds = 20
)

$ErrorActionPreference = 'Stop'

$root = $PSScriptRoot
$trainingDir = Join-Path $root 'training'
$checkpointRoot = Join-Path $trainingDir 'checkpoints'
$configSource = Join-Path $trainingDir 'aletheia_ru.onnx.json'
$python = Join-Path $root '.venv\Scripts\python.exe'
$outputRoot = Join-Path $trainingDir 'milestones'
$monitorLog = Join-Path $outputRoot 'monitor.log'
New-Item -ItemType Directory -Force -Path $outputRoot | Out-Null

function Write-MonitorLog([string] $Message) {
    $line = '{0} {1}' -f (Get-Date).ToUniversalTime().ToString('o'), $Message
    Add-Content -LiteralPath $monitorLog -Value $line -Encoding utf8
}

function Invoke-LoggedProcess(
    [string] $FilePath,
    [string[]] $Arguments,
    [string] $StdoutPath,
    [string] $StderrPath
) {
    $process = Start-Process -FilePath $FilePath `
        -ArgumentList $Arguments `
        -RedirectStandardOutput $StdoutPath `
        -RedirectStandardError $StderrPath `
        -WindowStyle Hidden `
        -Wait `
        -PassThru
    if ($process.ExitCode -ne 0) {
        throw "Process failed with exit code $($process.ExitCode): $FilePath"
    }
}

function Get-LatestCheckpoint {
    Get-ChildItem -LiteralPath $checkpointRoot -Filter '*.ckpt' -File -Recurse -ErrorAction SilentlyContinue |
        Where-Object { $_.Name -match '^epoch=(\d+)-step=(\d+)\.ckpt$' } |
        Sort-Object LastWriteTimeUtc -Descending |
        Select-Object -First 1
}

function Show-TrainingStatus {
    $trainingLog = Join-Path $trainingDir 'scheduled-training.stdout.log'
    $trainingTask = Get-ScheduledTask -TaskName 'AletheiaTTS-Training' -ErrorAction SilentlyContinue
    $progressLine = ''
    if (Test-Path -LiteralPath $trainingLog -PathType Leaf) {
        $progressLine = Get-Content -LiteralPath $trainingLog -Tail 30 |
            Where-Object { $_ -match '^Epoch\s+\d+:' } |
            Select-Object -Last 1
    }
    $checkpoint = Get-LatestCheckpoint
    $gpu = (& nvidia-smi --query-gpu=name,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw --format=csv,noheader,nounits 2>$null) -join [Environment]::NewLine

    try {
        $Host.UI.RawUI.WindowTitle = 'Aletheia TTS - training monitor'
        Clear-Host
    } catch {
        # The monitor also works when no interactive console is attached.
    }
    Write-Host 'Aletheia Russian TTS training'
    Write-Host ('Updated: {0}' -f (Get-Date).ToString('yyyy-MM-dd HH:mm:ss'))
    Write-Host ('Training task: {0}' -f $(if ($trainingTask) { $trainingTask.State } else { 'Not found' }))
    Write-Host ('Progress: {0}' -f $(if ($progressLine) { $progressLine.Trim() } else { 'Waiting for log data' }))
    Write-Host ('GPU: {0}' -f $(if ($gpu) { $gpu.Trim() } else { 'nvidia-smi unavailable' }))
    if ($checkpoint) {
        Write-Host ('Latest checkpoint: {0} ({1:N0} bytes)' -f $checkpoint.Name, $checkpoint.Length)
    } else {
        Write-Host 'Latest checkpoint: waiting'
    }
    Write-Host ''
    Write-Host 'Evaluation milestones:'
    foreach ($milestone in $Milestones) {
        $completePath = Join-Path $outputRoot ('epoch-{0:d4}\complete.json' -f $milestone)
        $state = if (Test-Path -LiteralPath $completePath -PathType Leaf) { 'READY' } else { 'waiting' }
        Write-Host ('  {0,4} epochs: {1}' -f $milestone, $state)
    }
    Write-Host ''
    Write-Host 'This window may be minimized. Closing it stops milestone exports only.'
}

function Save-Milestone([int] $Milestone, [IO.FileInfo] $SourceCheckpoint) {
    $milestoneName = 'epoch-{0:d4}' -f $Milestone
    $milestoneDir = Join-Path $outputRoot $milestoneName
    $completePath = Join-Path $milestoneDir 'complete.json'
    if (Test-Path -LiteralPath $completePath -PathType Leaf) {
        return
    }

    New-Item -ItemType Directory -Force -Path $milestoneDir | Out-Null
    $checkpointPath = Join-Path $milestoneDir "aletheia_ru_$milestoneName.ckpt"
    if (-not (Test-Path -LiteralPath $checkpointPath -PathType Leaf)) {
        $firstLength = $SourceCheckpoint.Length
        Start-Sleep -Seconds 10
        $refreshed = Get-Item -LiteralPath $SourceCheckpoint.FullName
        if ($firstLength -ne $refreshed.Length -or $refreshed.Length -le 0) {
            throw "Checkpoint is not stable yet: $($SourceCheckpoint.FullName)"
        }
        $partialPath = "$checkpointPath.partial"
        Copy-Item -LiteralPath $refreshed.FullName -Destination $partialPath -Force
        Move-Item -LiteralPath $partialPath -Destination $checkpointPath -Force
    }

    $modelPath = Join-Path $milestoneDir "aletheia_ru_$milestoneName.onnx"
    $configPath = "$modelPath.json"
    if (-not (Test-Path -LiteralPath $modelPath -PathType Leaf)) {
        Invoke-LoggedProcess $python @(
            '-m', 'piper.train.export_onnx',
            '--checkpoint', $checkpointPath,
            '--output-file', $modelPath
        ) (Join-Path $milestoneDir 'export.stdout.log') (Join-Path $milestoneDir 'export.stderr.log')
    }
    Copy-Item -LiteralPath $configSource -Destination $configPath -Force

    $inputPath = Join-Path $milestoneDir 'test_sentences.txt'
    $testText = @(
        '0JIg0YLQuNGI0LjQvdC1INCy0LXRh9C10YDQvdC10Lkg0LHQuNCx0LvQuNC+0YLQtdC60Lgg0YjQtdC70LXRgdGC0LXQu9C4INGB0YLRgNCw0L3QuNGG0YssINC4INC60LDQttC00LDRjyDQvdC+0LLQsNGPINCz0LvQsNCy0LAg0L7RgtC60YDRi9Cy0LDQu9CwINGD0LTQuNCy0LjRgtC10LvRjNC90YvQuSDQvNC40YAu',
        '0JrQvtCz0LTQsCDRh9C10LvQvtCy0LXQuiDQtNC10LvQsNC10YIg0LLRi9Cx0L7RgCwg0L7QvSDQvdC1INCy0YHQtdCz0LTQsCDQt9Cw0YDQsNC90LXQtSDQt9C90LDQtdGCLCDQuiDQutCw0LrQuNC8INC/0L7RgdC70LXQtNGB0YLQstC40Y/QvCDQv9GA0LjQstC10LTRkdGCINC10LPQviDRgNC10YjQtdC90LjQtS4=',
        '0JfQsCDQvtC60L3QvtC8INC80LXQtNC70LXQvdC90L4g0L3QsNGH0LjQvdCw0LvRgdGPINC00L7QttC00YwsINC90L4g0L/Rg9GC0LXRiNC10YHRgtCy0LjQtSDQs9C10YDQvtC10LIg0YLQvtC70YzQutC+INC90LDQsdC40YDQsNC70L4g0YHQuNC70YMu'
    ) | ForEach-Object { [Text.Encoding]::UTF8.GetString([Convert]::FromBase64String($_)) }
    $testText = $testText -join [Environment]::NewLine
    [IO.File]::WriteAllText($inputPath, $testText, [Text.UTF8Encoding]::new($false))

    $samples = @(
        @{ Name = 'neutral'; Length = '1.0'; Noise = '0.667'; Width = '0.8'; Silence = '0.0' },
        @{ Name = 'book'; Length = '1.5'; Noise = '0.75'; Width = '0.95'; Silence = '0.25' },
        @{ Name = 'fast'; Length = '0.82'; Noise = '0.70'; Width = '0.85'; Silence = '0.0' }
    )
    foreach ($sample in $samples) {
        $wavPath = Join-Path $milestoneDir ("sample_{0}.wav" -f $sample.Name)
        if (Test-Path -LiteralPath $wavPath -PathType Leaf) {
            continue
        }
        Invoke-LoggedProcess $python @(
            '-m', 'piper',
            '--model', $modelPath,
            '--config', $configPath,
            '--input-file', $inputPath,
            '--output-file', $wavPath,
            '--length-scale', $sample.Length,
            '--noise-scale', $sample.Noise,
            '--noise-w-scale', $sample.Width,
            '--sentence-silence', $sample.Silence
        ) (Join-Path $milestoneDir "$($sample.Name).stdout.log") (Join-Path $milestoneDir "$($sample.Name).stderr.log")
    }

    $sourceMatch = [regex]::Match($SourceCheckpoint.Name, '^epoch=(\d+)-step=(\d+)\.ckpt$')
    $artifacts = Get-ChildItem -LiteralPath $milestoneDir -File |
        Where-Object { $_.Extension -in @('.ckpt', '.onnx', '.json', '.wav') } |
        ForEach-Object {
            [ordered]@{
                name = $_.Name
                bytes = $_.Length
                sha256 = (Get-FileHash -LiteralPath $_.FullName -Algorithm SHA256).Hash.ToLowerInvariant()
            }
        }
    [ordered]@{
        requested_completed_epochs = $Milestone
        source_epoch_zero_based = [int]$sourceMatch.Groups[1].Value
        source_step = [int]$sourceMatch.Groups[2].Value
        created_utc = (Get-Date).ToUniversalTime().ToString('o')
        artifacts = @($artifacts)
    } | ConvertTo-Json -Depth 5 | Set-Content -LiteralPath $completePath -Encoding utf8
    Write-MonitorLog "Completed milestone $Milestone from $($SourceCheckpoint.Name)"
}

Write-MonitorLog "Monitor started for milestones: $($Milestones -join ', ')"
while ($true) {
    Show-TrainingStatus
    $pending = @($Milestones | Where-Object {
        -not (Test-Path -LiteralPath (Join-Path $outputRoot ('epoch-{0:d4}\complete.json' -f $_)) -PathType Leaf)
    })
    if ($pending.Count -eq 0) {
        Write-MonitorLog 'All milestones completed'
        exit 0
    }

    try {
        $checkpoint = Get-LatestCheckpoint
        if ($checkpoint -and $checkpoint.Name -match '^epoch=(\d+)-step=(\d+)\.ckpt$') {
            $completedEpochs = [int]$Matches[1] + 1
            foreach ($milestone in $pending) {
                if ($completedEpochs -ge $milestone) {
                    Save-Milestone $milestone $checkpoint
                }
            }
        }
    } catch {
        Write-MonitorLog ("Retryable error: " + ($_ | Out-String).Trim())
    }

    Start-Sleep -Seconds $PollSeconds
}