# Downloads prachuryyaIITG/SampurNER_Tamil_IndicBERTv2 ONNX assets to models/ta/ for the PII Redaction POC.
param(
    [string]$Python = "python"
)

$ErrorActionPreference = "Stop"
$repoRoot = Split-Path -Parent $PSScriptRoot
$modelsDir = Join-Path (Join-Path $repoRoot "models") "ta"
$scriptPath = Join-Path $PSScriptRoot "download-tamil-ner-model.py"
$modelId = "prachuryyaIITG/SampurNER_Tamil_IndicBERTv2"
$baseUrl = "https://huggingface.co/$modelId/resolve/main"

function Ensure-ModelsDirectory {
    New-Item -ItemType Directory -Force -Path $modelsDir | Out-Null
}

function Download-HuggingFaceAsset {
    param(
        [string]$RelativePath,
        [string]$Destination
    )

    $url = "$baseUrl/$RelativePath"
    Write-Host "Downloading $url"
    Invoke-WebRequest -Uri $url -OutFile $Destination -UseBasicParsing
}

function Export-LabelsFromConfig {
    param([string]$ConfigPath, [string]$LabelsPath)

    $config = Get-Content $ConfigPath -Raw | ConvertFrom-Json
    $labelMap = @{}
    foreach ($property in $config.id2label.PSObject.Properties) {
        $labelMap[[int]$property.Name] = [string]$property.Value
    }

    $labels = for ($index = 0; $index -lt $labelMap.Count; $index++) {
        $labelMap[$index]
    }

    $labels | Set-Content -Path $LabelsPath -Encoding utf8
}

function Resolve-PythonExecutable {
    param([string]$Preferred = "python")

    if ($Preferred -ne "python") {
        if ((Get-Command $Preferred -ErrorAction SilentlyContinue) -and
            -not ((Get-Command $Preferred).Source -like "*WindowsApps*")) {
            return $Preferred
        }
    }

    $candidates = @(
        (Join-Path $env:LOCALAPPDATA "Programs\Python\Python313\python.exe"),
        (Join-Path $env:LOCALAPPDATA "Programs\Python\Python312\python.exe"),
        (Join-Path $env:LOCALAPPDATA "Programs\Python\Python311\python.exe"),
        "C:\Program Files\Python312\python.exe",
        "C:\Program Files\Python313\python.exe"
    )

    foreach ($candidate in $candidates) {
        if (Test-Path $candidate) {
            return $candidate
        }
    }

    $pythonCommand = Get-Command python -ErrorAction SilentlyContinue
    if ($pythonCommand -and $pythonCommand.Source -notlike "*WindowsApps*") {
        return $pythonCommand.Source
    }

    return $null
}

function Export-VocabFromTokenizerJson {
    param([string]$TokenizerJsonPath, [string]$VocabPath)

    $tokenizer = Get-Content $TokenizerJsonPath -Raw | ConvertFrom-Json
    $vocab = $tokenizer.model.vocab
    if (-not $vocab) {
        return $false
    }

    $orderedTokens = $vocab.PSObject.Properties |
        Sort-Object { [int]$_.Value } |
        ForEach-Object { $_.Name }

    $orderedTokens | Set-Content -Path $VocabPath -Encoding utf8
    return $true
}

function Copy-TokenizerAssets {
    param([string]$SourceDir)

    foreach ($name in @("sentencepiece.bpe.model", "spiece.model", "tokenizer.model")) {
        $source = Join-Path $SourceDir $name
        if (Test-Path $source) {
            $destination = Join-Path $modelsDir $name
            Copy-Item $source $destination -Force
            return @($destination)
        }
    }

    $tokenizerJsonSource = Join-Path $SourceDir "tokenizer.json"
    if (Test-Path $tokenizerJsonSource) {
        $tokenizerJsonDestination = Join-Path $modelsDir "tokenizer.json"
        Copy-Item $tokenizerJsonSource $tokenizerJsonDestination -Force
        $saved = @($tokenizerJsonDestination)

        $vocabPath = Join-Path $modelsDir "vocab.txt"
        if (Export-VocabFromTokenizerJson -TokenizerJsonPath $tokenizerJsonDestination -VocabPath $vocabPath) {
            $saved += $vocabPath
        }

        Write-Warning (
            "No SentencePiece model on Hugging Face; saved WordPiece assets ($(
                ($saved | ForEach-Object { Split-Path $_ -Leaf }) -join ', '
            )). TamilOnnxNerRunner uses vocab.txt (WordPiece) when present, otherwise SentencePiece model files."
        )
        return $saved
    }

    return @()
}

function Download-WithPowerShell {
    Ensure-ModelsDirectory

    $modelPath = Join-Path $modelsDir "model.onnx"
    $configPath = Join-Path $modelsDir "config.json"
    $labelsPath = Join-Path $modelsDir "ner-labels.txt"
    $tempDir = Join-Path $modelsDir "_download_temp"
    New-Item -ItemType Directory -Force -Path $tempDir | Out-Null

    try {
        Download-HuggingFaceAsset -RelativePath "onnx/model.onnx" -Destination $modelPath
    }
    catch {
        Write-Host "Pre-exported ONNX not found; downloading config and tokenizer for manual export..."
        Download-HuggingFaceAsset -RelativePath "config.json" -Destination $configPath
        Export-LabelsFromConfig -ConfigPath $configPath -LabelsPath $labelsPath

        $tokenizerJsonPath = Join-Path $modelsDir "tokenizer.json"
        try {
            Download-HuggingFaceAsset -RelativePath "tokenizer.json" -Destination $tokenizerJsonPath
            $vocabPath = Join-Path $modelsDir "vocab.txt"
            Export-VocabFromTokenizerJson -TokenizerJsonPath $tokenizerJsonPath -VocabPath $vocabPath | Out-Null
        }
        catch {
            Write-Host "tokenizer.json not available from Hugging Face."
        }

        foreach ($name in @("sentencepiece.bpe.model", "spiece.model", "tokenizer.model")) {
            try {
                Download-HuggingFaceAsset -RelativePath $name -Destination (Join-Path $modelsDir $name)
                break
            }
            catch {
                continue
            }
        }

        $pythonExe = Resolve-PythonExecutable -Preferred $Python
        if ($pythonExe) {
            throw (
                "Tamil ONNX model is not published on Hugging Face. Re-run with Python export:`n" +
                "  .\scripts\download-tamil-ner-model.ps1 -Python `"$pythonExe`""
            )
        }

        throw @"
Tamil ONNX model is not published on Hugging Face (onnx/model.onnx returns 404).
Install Python 3.12+ and re-run this script:
  winget install Python.Python.3.12 --accept-package-agreements --accept-source-agreements
  .\scripts\download-tamil-ner-model.ps1 -Python `"`$env:LOCALAPPDATA\Programs\Python\Python312\python.exe`"
Only config.json, ner-labels.txt, and tokenizer.json were saved under models/ta/.
"@
    }

    Download-HuggingFaceAsset -RelativePath "config.json" -Destination $configPath
    Export-LabelsFromConfig -ConfigPath $configPath -LabelsPath $labelsPath
    Remove-Item $configPath -Force

    foreach ($name in @("sentencepiece.bpe.model", "spiece.model")) {
        try {
            Download-HuggingFaceAsset -RelativePath $name -Destination (Join-Path $modelsDir $name)
            break
        }
        catch {
            continue
        }
    }

    Write-Host ""
    Write-Host "Tamil NER model assets saved to $modelsDir"
}

function Download-WithPython {
    $pythonExe = Resolve-PythonExecutable -Preferred $Python
    if (-not $pythonExe) {
        return $false
    }

    Write-Host "Using Python: $pythonExe"
    Write-Host "Repository root: $repoRoot"
    Write-Host ""

    Push-Location $repoRoot
    try {
        & $pythonExe $scriptPath
        if ($LASTEXITCODE -ne 0) {
            throw "Tamil model download script failed with exit code $LASTEXITCODE."
        }
    }
    finally {
        Pop-Location
    }

    return $true
}

Write-Host "Repository root: $repoRoot"

if (-not (Download-WithPython)) {
    Write-Host "Python export unavailable; downloading pre-exported ONNX assets from Hugging Face..."
    Write-Host ""
    Download-WithPowerShell
}