[CmdletBinding()]
param(
    [Parameter(Mandatory = $true)]
    [string] $DatasetDir,
    [Parameter(Mandatory = $true)]
    [string] $OutputDir,
    [Parameter(Mandatory = $true)]
    [string] $PythonExe,
    [string] $CacheDir,
    [string] $CheckpointPath,
    [int] $BatchSize = 4,
    [int] $NumWorkers = 2,
    [int] $MaxEpochs = 2000,
    [switch] $SmokeTest
)

$ErrorActionPreference = 'Stop'

$dataset = [IO.Path]::GetFullPath($DatasetDir)
$output = [IO.Path]::GetFullPath($OutputDir)
$python = [IO.Path]::GetFullPath($PythonExe)
$metadata = Join-Path $dataset 'metadata.csv'
$audioDir = Join-Path $dataset 'wav'
$configPath = Join-Path $output 'aletheia_ru.onnx.json'
$cacheDir = if ($CacheDir) {
    [IO.Path]::GetFullPath($CacheDir)
} else {
    Join-Path $output 'cache'
}
$checkpointDir = Join-Path $output 'checkpoints'

if (-not (Test-Path -LiteralPath $python -PathType Leaf)) {
    throw "Python executable not found: $python"
}
if (-not (Test-Path -LiteralPath $metadata -PathType Leaf)) {
    throw "Piper metadata not found: $metadata"
}
if (-not (Test-Path -LiteralPath $audioDir -PathType Container)) {
    throw "Piper audio directory not found: $audioDir"
}

New-Item -ItemType Directory -Force -Path $output, $cacheDir, $checkpointDir | Out-Null

$fitArgs = @(
    '-m', 'piper.train', 'fit',
    '--data.voice_name', 'aletheia_ru',
    '--data.csv_path', $metadata,
    '--data.audio_dir', $audioDir,
    '--data.espeak_voice', 'ru',
    '--data.cache_dir', $cacheDir,
    '--data.config_path', $configPath,
    '--data.batch_size', $BatchSize,
    '--data.num_workers', $NumWorkers,
    '--model.sample_rate', '22050',
    '--trainer.accelerator', 'gpu',
    '--trainer.devices', '1',
    '--trainer.precision', '16-mixed',
    '--trainer.max_epochs', $MaxEpochs,
    '--trainer.default_root_dir', $checkpointDir
)

if ($SmokeTest) {
    $fitArgs += @('--trainer.fast_dev_run', 'true', '--trainer.num_sanity_val_steps', '0')
}
if ($CheckpointPath) {
    $checkpointToResume = [IO.Path]::GetFullPath($CheckpointPath)
    if (-not (Test-Path -LiteralPath $checkpointToResume -PathType Leaf)) {
        throw "Resume checkpoint not found: $checkpointToResume"
    }
    $fitArgs += @('--ckpt_path', $checkpointToResume)
}

& $python @fitArgs
if ($LASTEXITCODE -ne 0) {
    throw "Piper training exited with code $LASTEXITCODE"
}

if ($SmokeTest) {
    Write-Output 'PIPER_CUDA_SMOKE_TEST_OK'
    exit 0
}

$checkpoint = Get-ChildItem -LiteralPath $checkpointDir -Filter '*.ckpt' -File -Recurse |
    Sort-Object LastWriteTimeUtc -Descending |
    Select-Object -First 1
if (-not $checkpoint) {
    throw "Training finished without a checkpoint under $checkpointDir"
}

$onnxPath = Join-Path $output 'aletheia_ru.onnx'
& $python -m piper.train.export_onnx --checkpoint $checkpoint.FullName --output-file $onnxPath
if ($LASTEXITCODE -ne 0) {
    throw "Piper ONNX export exited with code $LASTEXITCODE"
}

$artifacts = [ordered]@{}
foreach ($artifactPath in @($onnxPath, $configPath)) {
    if (-not (Test-Path -LiteralPath $artifactPath -PathType Leaf)) {
        throw "Expected artifact not found: $artifactPath"
    }
    $item = Get-Item -LiteralPath $artifactPath
    $artifacts[$item.Name] = [ordered]@{
        bytes = $item.Length
        sha256 = (Get-FileHash -LiteralPath $item.FullName -Algorithm SHA256).Hash.ToLowerInvariant()
    }
}

$manifestPath = Join-Path $output 'artifacts.json'
$artifacts | ConvertTo-Json -Depth 4 | Set-Content -LiteralPath $manifestPath -Encoding utf8
$artifacts | ConvertTo-Json -Depth 4