# Downloads dslim/bert-base-NER ONNX assets to models/ for the PII Redaction POC.
param(
    [string]$Python = "python"
)

$ErrorActionPreference = "Stop"
$repoRoot = Split-Path -Parent $PSScriptRoot
$modelsDir = Join-Path $repoRoot "models"
$scriptPath = Join-Path $PSScriptRoot "download-ner-model.py"
$baseUrl = "https://huggingface.co/dslim/bert-base-NER/resolve/main/onnx"

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 Download-WithPowerShell {
    Ensure-ModelsDirectory

    $modelPath = Join-Path $modelsDir "ner-model.onnx"
    $vocabPath = Join-Path $modelsDir "vocab.txt"
    $configPath = Join-Path $modelsDir "config.json"
    $labelsPath = Join-Path $modelsDir "ner-labels.txt"

    Download-HuggingFaceAsset -RelativePath "model.onnx" -Destination $modelPath
    Download-HuggingFaceAsset -RelativePath "vocab.txt" -Destination $vocabPath
    Download-HuggingFaceAsset -RelativePath "config.json" -Destination $configPath
    Export-LabelsFromConfig -ConfigPath $configPath -LabelsPath $labelsPath
    Remove-Item $configPath -Force

    Write-Host ""
    Write-Host "NER model assets saved:"
    Write-Host "  $modelPath"
    Write-Host "  $vocabPath"
    Write-Host "  $labelsPath"
    Write-Host ""
    Write-Host "Run from repository root:"
    Write-Host "  dotnet run --project src/PiiRedaction.ConsoleApp -- --name CustomerNameOnly"
}

function Download-WithPython {
    if (-not (Get-Command $Python -ErrorAction SilentlyContinue)) {
        return $false
    }

    $pythonCommand = Get-Command $Python
    if ($pythonCommand.Source -like "*WindowsApps*") {
        return $false
    }

    Write-Host "Using Python: $($pythonCommand.Source)"
    Write-Host "Repository root: $repoRoot"
    Write-Host ""

    Push-Location $repoRoot
    try {
        & $Python $scriptPath
        if ($LASTEXITCODE -ne 0) {
            throw "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
}