param(
    [string]$ModelsSourceDir = "",
    [string]$ModelBaseUrl = "https://huggingface.co/sauropod/parakeet-tdt-0.6b-v2/resolve/053aaaf310acc6682816d188f0762c4aa78c5acb",
    [string]$ChecksumsUrl = "",
    [string]$DestinationDir = ""
)

Set-StrictMode -Version Latest
$ErrorActionPreference = "Stop"
[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12

function Write-Step {
    param([string]$Message)
    Write-Host "[download-models] $Message"
}

function Ensure-FileExists {
    param([string]$Path, [string]$Label)
    if (-not (Test-Path -Path $Path -PathType Leaf)) {
        throw "Missing $Label at '$Path'."
    }
}

function Get-Sha256Hex {
    param([string]$Path)

    if (-not (Test-Path -Path $Path -PathType Leaf)) {
        throw "Cannot hash missing file '$Path'."
    }

    return (Get-FileHash -Path $Path -Algorithm SHA256).Hash.ToLowerInvariant()
}

function Assert-NotLfsPointer {
    param([string]$Path, [string]$Label)

    if (-not (Test-Path -Path $Path -PathType Leaf)) {
        throw "Cannot inspect missing $Label at '$Path'."
    }

    $stream = [System.IO.File]::OpenRead($Path)
    try {
        $buffer = New-Object byte[] 256
        $read = $stream.Read($buffer, 0, $buffer.Length)
        if ($read -le 0) {
            throw "$Label at '$Path' is empty."
        }

        $head = [System.Text.Encoding]::UTF8.GetString($buffer, 0, $read)
        if ($head -match "git-lfs.github.com/spec/v1") {
            throw "$Label at '$Path' appears to be a Git LFS pointer, not a real model artifact."
        }
    }
    finally {
        $stream.Dispose()
    }
}

function Assert-ModelLayout {
    param([string]$DestinationDir)

    $encoder = Join-Path $DestinationDir "encoder-parakeet-tdt-0.6b-v2.onnx"
    $encoderData = Join-Path $DestinationDir "encoder-parakeet-tdt-0.6b-v2.onnx.data"
    $decoder = Join-Path $DestinationDir "decoder_joint-parakeet-tdt-0.6b-v2.onnx"
    $decoderData = Join-Path $DestinationDir "decoder_joint-parakeet-tdt-0.6b-v2.onnx.data"

    Ensure-FileExists -Path $encoder -Label "encoder model"
    Ensure-FileExists -Path $encoderData -Label "encoder sidecar"
    Ensure-FileExists -Path $decoder -Label "decoder model"

    Assert-NotLfsPointer -Path $encoder -Label "encoder model"
    Assert-NotLfsPointer -Path $decoder -Label "decoder model"

    $encoderLen = (Get-Item $encoder).Length
    $encoderDataLen = (Get-Item $encoderData).Length
    $decoderLen = (Get-Item $decoder).Length
    $decoderDataExists = Test-Path -Path $decoderData -PathType Leaf
    $decoderDataLen = if ($decoderDataExists) { (Get-Item $decoderData).Length } else { 0 }

    Write-Step ("Model size summary: encoder.onnx={0} encoder.onnx.data={1} decoder.onnx={2} decoder.onnx.data_exists={3} decoder.onnx.data={4}" -f `
        $encoderLen, $encoderDataLen, $decoderLen, $decoderDataExists, $decoderDataLen)

    # Small decoder files are valid only when the decoder sidecar is present.
    # This fails early with a clear message instead of later ONNX runtime errors.
    if ($decoderLen -lt 1MB -and -not $decoderDataExists) {
        throw "Decoder ONNX appears to require external data (small file) but decoder sidecar is missing: '$decoderData'."
    }

    if ($encoderDataLen -lt 128MB) {
        throw "Encoder sidecar appears truncated: '$encoderData' ($encoderDataLen bytes)."
    }
}

function Assert-ChecksumIfPresent {
    param(
        [string]$Path,
        [string]$Name,
        [hashtable]$ChecksumMap
    )

    if (-not $ChecksumMap.ContainsKey($Name)) {
        return
    }

    $expected = ($ChecksumMap[$Name] | Out-String).Trim().ToLowerInvariant()
    if ([string]::IsNullOrWhiteSpace($expected)) {
        return
    }

    $actual = Get-Sha256Hex -Path $Path
    if ($actual -ne $expected) {
        throw "Checksum mismatch for '$Name'. Expected=$expected Actual=$actual"
    }
}

function Get-PinnedCanonicalChecksums {
    return @{
        "encoder-parakeet-tdt-0.6b-v2.onnx" = "00e6e049507e402a4c9aeaa28aeaf0044e627171bd9622cb4af49094133c25a9"
        "encoder-parakeet-tdt-0.6b-v2.onnx.data" = "6120304740e831a6869205265d751df83e35d944e53f511da3084705f87ad9e7"
        "decoder_joint-parakeet-tdt-0.6b-v2.onnx" = "0e58750ba2f289eb0ade1d4cba27753d5e0800c2708c0fa77cbab40db787e8e2"
    }
}

function Use-PinnedChecksumsForSource {
    param(
        [string]$ChecksumsUrl,
        [string]$ModelBaseUrl
    )

    if (-not [string]::IsNullOrWhiteSpace($ChecksumsUrl)) {
        return $false
    }

    return $ModelBaseUrl -like "https://huggingface.co/sauropod/parakeet-tdt-0.6b-v2/resolve/*"
}

function Try-ResolveHuggingFaceChecksums {
    param([string]$ModelBaseUrl)

    $pattern = "^https://huggingface\\.co/(?<repo>[^/]+/[^/]+)/resolve/(?<revision>[^/?#]+)"
    if ($ModelBaseUrl -notmatch $pattern) {
        return @{}
    }

    $repo = $Matches["repo"]
    $revision = $Matches["revision"]
    $apiUrl = "https://huggingface.co/api/models/$repo/revision/$revision?blobs=true"

    try {
        Write-Step "Fetching checksum metadata from Hugging Face API: $apiUrl"
        $payload = Invoke-WebRequest -Uri $apiUrl -UseBasicParsing
        $json = $payload.Content | ConvertFrom-Json
    }
    catch {
        Write-Step "Unable to read Hugging Face API metadata for checksums."
        return @{}
    }

    $map = @{}
    foreach ($entry in $json.siblings) {
        if ($null -eq $entry) {
            continue
        }

        $name = $entry.rfilename
        $sha = $null
        if ($null -ne $entry.lfs) {
            $sha = $entry.lfs.sha256
        }

        if (-not [string]::IsNullOrWhiteSpace($name) -and -not [string]::IsNullOrWhiteSpace($sha)) {
            $map[$name] = ($sha | Out-String).Trim().ToLowerInvariant()
        }
    }

    if ($map.Count -gt 0) {
        Write-Step "Resolved $($map.Count) checksums from Hugging Face API."
    }

    return $map
}

function Resolve-Checksums {
    param(
        [string]$ChecksumsUrl,
        [string]$ModelBaseUrl
    )

    $resolvedUrl = $ChecksumsUrl
    if ([string]::IsNullOrWhiteSpace($resolvedUrl)) {
        $resolvedUrl = "$ModelBaseUrl/checksums.json"
    }

    try {
        Write-Step "Fetching checksums from $resolvedUrl"
        $payload = Invoke-WebRequest -Uri $resolvedUrl -UseBasicParsing
        $json = $payload.Content | ConvertFrom-Json
    }
    catch {
        if (Use-PinnedChecksumsForSource -ChecksumsUrl $ChecksumsUrl -ModelBaseUrl $ModelBaseUrl) {
            $apiMap = Try-ResolveHuggingFaceChecksums -ModelBaseUrl $ModelBaseUrl
            if ($apiMap.Count -gt 0) {
                return $apiMap
            }

            Write-Step "Checksums unavailable; using pinned canonical checksum map."
            return Get-PinnedCanonicalChecksums
        }

        throw "Cannot proceed without checksum verification for source: $resolvedUrl"
    }

    $map = @{}
    if ($json -is [System.Collections.IEnumerable] -and -not ($json -is [string])) {
        foreach ($entry in $json) {
            if ($null -ne $entry.filename -and $null -ne $entry.sha256) {
                $map[$entry.filename] = ($entry.sha256 | Out-String).Trim()
            }
        }
        if ($map.Count -gt 0) {
            return $map
        }
    }

    foreach ($property in $json.PSObject.Properties) {
        $map[$property.Name] = ($property.Value | Out-String).Trim()
    }

    if ($map.Count -gt 0) {
        return $map
    }

    if (Use-PinnedChecksumsForSource -ChecksumsUrl $ChecksumsUrl -ModelBaseUrl $ModelBaseUrl) {
        $apiMap = Try-ResolveHuggingFaceChecksums -ModelBaseUrl $ModelBaseUrl
        if ($apiMap.Count -gt 0) {
            return $apiMap
        }

        Write-Step "Checksum payload was empty; using pinned canonical checksum map."
        return Get-PinnedCanonicalChecksums
    }

    throw "Cannot proceed without checksum verification for source: $resolvedUrl"
}

function Copy-SharedAssets {
    param(
        [string]$RepoRoot,
        [string]$DestinationDir
    )

    $vocab = Join-Path $RepoRoot "Resources\parakeet_vocab_1024.json"
    $vad = Join-Path $RepoRoot "Resources\silero_vad.onnx"

    Ensure-FileExists -Path $vocab -Label "vocabulary"
    Ensure-FileExists -Path $vad -Label "VAD model"

    Copy-Item -Path $vocab -Destination $DestinationDir -Force
    Copy-Item -Path $vad -Destination $DestinationDir -Force
}

function Copy-CanonicalExports {
    param(
        [string]$SourceDir,
        [string]$DestinationDir
    )

    $encoder = Join-Path $SourceDir "encoder-parakeet-tdt-0.6b-v2.onnx"
    $decoder = Join-Path $SourceDir "decoder_joint-parakeet-tdt-0.6b-v2.onnx"
    $encoderData = Join-Path $SourceDir "encoder-parakeet-tdt-0.6b-v2.onnx.data"
    $decoderData = Join-Path $SourceDir "decoder_joint-parakeet-tdt-0.6b-v2.onnx.data"

    Ensure-FileExists -Path $encoder -Label "encoder model"
    Ensure-FileExists -Path $decoder -Label "decoder model"

    Copy-Item -Path $encoder -Destination $DestinationDir -Force
    Copy-Item -Path $decoder -Destination $DestinationDir -Force

    if (Test-Path -Path $encoderData -PathType Leaf) {
        Copy-Item -Path $encoderData -Destination $DestinationDir -Force
    }
    if (Test-Path -Path $decoderData -PathType Leaf) {
        Copy-Item -Path $decoderData -Destination $DestinationDir -Force
    }
}

function Download-File {
    param(
        [string]$Url,
        [string]$DestinationPath,
        [string]$ExpectedSha256 = ""
    )

    $name = [System.IO.Path]::GetFileName($DestinationPath)
    if (Test-Path -Path $DestinationPath -PathType Leaf) {
        if (-not [string]::IsNullOrWhiteSpace($ExpectedSha256)) {
            try {
                $actual = Get-Sha256Hex -Path $DestinationPath
                if ($actual -eq $ExpectedSha256.ToLowerInvariant()) {
                    Write-Step "Using existing verified file: $name"
                    return
                }

                Write-Step "Existing file checksum mismatch for $name. Re-downloading."
                Remove-Item -Path $DestinationPath -Force
            }
            catch {
                Write-Step "Could not verify existing $name. Re-downloading."
                Remove-Item -Path $DestinationPath -Force -ErrorAction SilentlyContinue
            }
        }
        else {
            Write-Step "Using existing file (no checksum available): $name"
            return
        }
    }

    $tmpPath = "$DestinationPath.download"
    if (Test-Path -Path $tmpPath -PathType Leaf) {
        Remove-Item -Path $tmpPath -Force
    }

    Write-Step "Downloading $name"

    $isCi = $env:GITHUB_ACTIONS -eq "true"
    if ($isCi) {
        # Prefer deterministic HTTP download in CI.
        Invoke-WebRequest -Uri $Url -OutFile $tmpPath -UseBasicParsing
    }
    else {
        try {
            Start-BitsTransfer -Source $Url -Destination $tmpPath -Description "IceWhisper model download: $name"
        }
        catch {
            Write-Step "BITS failed for $name; retrying with Invoke-WebRequest."
            Invoke-WebRequest -Uri $Url -OutFile $tmpPath -UseBasicParsing
        }
    }

    if (-not (Test-Path -Path $tmpPath -PathType Leaf)) {
        throw "Download did not produce a file for '$name'."
    }
    if ((Get-Item $tmpPath).Length -le 0) {
        throw "Downloaded file for '$name' is empty."
    }

    Move-Item -Path $tmpPath -Destination $DestinationPath -Force

    if (-not [string]::IsNullOrWhiteSpace($ExpectedSha256)) {
        $actual = Get-Sha256Hex -Path $DestinationPath
        if ($actual -ne $ExpectedSha256.ToLowerInvariant()) {
            Remove-Item -Path $DestinationPath -Force -ErrorAction SilentlyContinue
            throw "Checksum mismatch after download for '$name'."
        }
    }
}

$scriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path
$repoRoot = Split-Path -Parent $scriptDir

if ([string]::IsNullOrWhiteSpace($DestinationDir)) {
    $destinationDir = Join-Path $env:LOCALAPPDATA "IceWhisper\Models"
}
else {
    $destinationDir = [System.IO.Path]::GetFullPath($DestinationDir)
}

New-Item -ItemType Directory -Path $destinationDir -Force | Out-Null
Write-Step "Model destination: $destinationDir"

if (-not [string]::IsNullOrWhiteSpace($ModelsSourceDir)) {
    $resolvedSource = (Resolve-Path -Path $ModelsSourceDir).Path
    Write-Step "Using provided model export directory: $resolvedSource"
    Copy-CanonicalExports -SourceDir $resolvedSource -DestinationDir $destinationDir
    Copy-SharedAssets -RepoRoot $repoRoot -DestinationDir $destinationDir
}
else {
    $checksums = Resolve-Checksums -ChecksumsUrl $ChecksumsUrl -ModelBaseUrl $ModelBaseUrl
    $downloads = @(
        @{ Name = "encoder-parakeet-tdt-0.6b-v2.onnx"; Url = "$ModelBaseUrl/encoder-parakeet-tdt-0.6b-v2.onnx"; Required = $true },
        @{ Name = "decoder_joint-parakeet-tdt-0.6b-v2.onnx"; Url = "$ModelBaseUrl/decoder_joint-parakeet-tdt-0.6b-v2.onnx"; Required = $true },
        @{ Name = "encoder-parakeet-tdt-0.6b-v2.onnx.data"; Url = "$ModelBaseUrl/encoder-parakeet-tdt-0.6b-v2.onnx.data"; Required = $true },
        @{ Name = "decoder_joint-parakeet-tdt-0.6b-v2.onnx.data"; Url = "$ModelBaseUrl/decoder_joint-parakeet-tdt-0.6b-v2.onnx.data"; Required = $false }
    )

    foreach ($entry in $downloads) {
        $targetPath = Join-Path $destinationDir $entry.Name
        $expectedSha = ""
        if ($checksums.ContainsKey($entry.Name)) {
            $expectedSha = ($checksums[$entry.Name] | Out-String).Trim()
        }
        if ([string]::IsNullOrWhiteSpace($expectedSha) -and $entry.Required) {
            throw "Cannot proceed without checksum verification for $($entry.Name)"
        }
        try {
            Download-File -Url $entry.Url -DestinationPath $targetPath -ExpectedSha256 $expectedSha
        }
        catch {
            if ($entry.Required) {
                throw
            }
            Write-Step "Optional file unavailable ($($entry.Name)); continuing."
        }
    }

    Copy-SharedAssets -RepoRoot $repoRoot -DestinationDir $destinationDir
}

$required = @(
    "encoder-parakeet-tdt-0.6b-v2.onnx",
    "decoder_joint-parakeet-tdt-0.6b-v2.onnx",
    "encoder-parakeet-tdt-0.6b-v2.onnx.data",
    "parakeet_vocab_1024.json",
    "silero_vad.onnx"
)

foreach ($file in $required) {
    Ensure-FileExists -Path (Join-Path $destinationDir $file) -Label $file
}

Assert-ModelLayout -DestinationDir $destinationDir

if (-not [string]::IsNullOrWhiteSpace($ModelsSourceDir)) {
    Write-Step "Imported models from local source. Skipping checksum verification for imported files."
}
else {
    $checksums = Resolve-Checksums -ChecksumsUrl $ChecksumsUrl -ModelBaseUrl $ModelBaseUrl
    foreach ($file in $required + "decoder_joint-parakeet-tdt-0.6b-v2.onnx.data") {
        $path = Join-Path $destinationDir $file
        if (Test-Path -Path $path -PathType Leaf) {
            Assert-ChecksumIfPresent -Path $path -Name $file -ChecksumMap $checksums
        }
    }
}

Write-Step "Installed models and assets:"
Get-ChildItem -Path $destinationDir |
    Where-Object { $required -contains $_.Name -or $_.Name -eq "decoder_joint-parakeet-tdt-0.6b-v2.onnx.data" } |
    Select-Object Name, Length, LastWriteTime
