param(
    [string]$Root = 'C:\Users\seven\AletheiaBookTTS',
    [int]$MaxSentences = 100000,
    [int]$TeacherBatchSize = 16,
    [ValidateSet('first', 'hash')]
    [string]$Sampling = 'hash'
)

$ErrorActionPreference = 'Stop'
$ProgressPreference = 'SilentlyContinue'
$env:PYTHONUTF8 = '1'
$env:PYTHONIOENCODING = 'utf-8'

$displayName = 'AletheiaBookTTS-RuWikisource'
$corpusDir = Join-Path $Root 'corpus'
$workDir = Join-Path $Root 'work'
$teacherDir = Join-Path $Root 'teacher'
$python = Join-Path $Root '.venv\Scripts\python.exe'
$dump = Join-Path $corpusDir 'ruwikisource-latest-pages-articles-multistream.xml.bz2'
$sentences = Join-Path $corpusDir "ruwikisource-booktts-$MaxSentences-$Sampling.txt"
$metadata = Join-Path $corpusDir "ruwikisource-booktts-$MaxSentences-$Sampling.json"
$dataset = Join-Path $workDir "student-dataset-$MaxSentences-$Sampling"
$datasetPartial = "$dataset.partial"
$expectedBytes = 2104816879L
$checksumUrl = 'https://dumps.wikimedia.org/ruwikisource/latest/ruwikisource-latest-sha1sums.txt'

Import-Module BitsTransfer
while (-not (Test-Path -LiteralPath $dump)) {
    $job = Get-BitsTransfer -ErrorAction SilentlyContinue |
        Where-Object DisplayName -eq $displayName |
        Select-Object -First 1
    if (-not $job) {
        throw "BITS job $displayName is absent and dump is not downloaded"
    }
    if ($job.JobState -eq 'Transferred') {
        Complete-BitsTransfer -BitsJob $job
        break
    }
    if ($job.JobState -eq 'Error') {
        throw "BITS download failed: $($job.ErrorDescription)"
    }
    $percent = if ($job.BytesTotal -gt 0 -and $job.BytesTotal -lt [uint64]::MaxValue) {
        [math]::Round(100 * $job.BytesTransferred / $job.BytesTotal, 2)
    } else { 0 }
    Write-Output "download state=$($job.JobState) bytes=$($job.BytesTransferred)/$($job.BytesTotal) percent=$percent"
    Start-Sleep -Seconds 20
}

$dumpItem = Get-Item -LiteralPath $dump
if ($dumpItem.Length -ne $expectedBytes) {
    throw "Unexpected dump size: $($dumpItem.Length), expected $expectedBytes"
}
$checksumText = (Invoke-WebRequest -UseBasicParsing -Uri $checksumUrl).Content
$checksumLine = ($checksumText -split "`n" | Where-Object { $_ -match 'pages-articles-multistream\.xml\.bz2\s*$' } | Select-Object -First 1).Trim()
if (-not $checksumLine) {
    throw 'Cannot find multistream dump in official SHA-1 list'
}
$expectedSha1 = ($checksumLine -split '\s+')[0].ToUpperInvariant()
$actualSha1 = (Get-FileHash -LiteralPath $dump -Algorithm SHA1).Hash
if ($actualSha1 -ne $expectedSha1) {
    throw "SHA-1 mismatch: actual=$actualSha1 expected=$expectedSha1"
}
Write-Output "dump verified bytes=$($dumpItem.Length) sha1=$actualSha1"

if (-not (Test-Path -LiteralPath $sentences)) {
    & $python (Join-Path $workDir 'extract_wikisource_corpus.py') `
        --dump $dump `
        --output $sentences `
        --metadata $metadata `
        --max-sentences $MaxSentences `
        --sampling $Sampling
    if ($LASTEXITCODE -ne 0) { throw "Corpus extraction failed with exit code $LASTEXITCODE" }
}

if (-not (Test-Path -LiteralPath $dataset)) {
    if (Test-Path -LiteralPath $datasetPartial) {
        throw "Partial dataset already exists: $datasetPartial"
    }
    & $python (Join-Path $workDir 'booktts_student.py') dataset `
        --assets $teacherDir `
        --corpus $sentences `
        --output $datasetPartial `
        --teacher-batch-size $TeacherBatchSize `
        --teacher-dtype float16
    if ($LASTEXITCODE -ne 0) { throw "Teacher dataset failed with exit code $LASTEXITCODE" }
    Move-Item -LiteralPath $datasetPartial -Destination $dataset
}

Write-Output "pipeline complete corpus=$sentences dataset=$dataset"