Param(
    [switch]$gpu,
    [string]$CONFIG_FILE = 'config.production.yaml',
    [switch]$buildFrontend
)

Write-Host "Starting RWKV local setup for Windows..." -ForegroundColor Green

if (-not (Get-Command python -ErrorAction SilentlyContinue)) {
    Write-Host "Python not found. Please install Python 3.10+ and add it to PATH." -ForegroundColor Red
    exit 1
}

Write-Host "Creating virtual environment (./.venv) ..."
python -m venv .venv
.\.venv\Scripts\Activate.ps1
pip install --upgrade pip setuptools wheel

if ($gpu) {
    Write-Host "GPU support requested. Installing GPU dependencies (cu124) ..." -ForegroundColor Yellow
    pip install -e .[cu124]
} else {
    Write-Host "Installing CPU-only dependencies ..." -ForegroundColor Yellow
    pip install -e .[cpu]
}

Write-Host "Installing extra tooling (huggingface_hub for downloads) ..."
pip install huggingface-hub
pip install beautifulsoup4

Write-Host "Downloading models from config: $CONFIG_FILE" -ForegroundColor Green
python .\download_models.py --config $CONFIG_FILE --token $env:HF_TOKEN

# Ensure the config file is available as config.local.yaml that the app's default reads
if (Test-Path $CONFIG_FILE) {
    Write-Host "Copying $CONFIG_FILE to config.local.yaml so the application uses it by default..." -ForegroundColor Green
    Copy-Item -Force $CONFIG_FILE config.local.yaml

    # If GPU not requested, set STRATEGY to CPU in config.local.yaml
    if (-not $gpu) {
        Write-Host "GPU not requested. Setting STRATEGY to 'cpu fp16' in config.local.yaml..." -ForegroundColor Yellow
        try {
            $yaml = (Get-Content config.local.yaml -Raw)
            # Replace the STRATEGY line with CPU + fp16 precision
            $yaml = $yaml -replace '(?m)^STRATEGY:.*$', 'STRATEGY: "cpu fp16"'
            # Also ensure RWKV_CUDA_ON is disabled when GPU is not requested
            if ($yaml -match 'RWKV_CUDA_ON:') {
                $yaml = $yaml -replace '(?m)^RWKV_CUDA_ON:.*$', 'RWKV_CUDA_ON: False'
            } else {
                $yaml = $yaml + "`nRWKV_CUDA_ON: False`n"
            }
            $yaml | Out-File -Encoding utf8 config.local.yaml -Force
        } catch {
            Write-Host "Warning: failed to modify config.local.yaml; please set STRATEGY manually" -ForegroundColor Yellow
        }
    }
}

if ($buildFrontend) {
    if (-not (Get-Command pnpm -ErrorAction SilentlyContinue)) {
        Write-Host "pnpm not found. Installing pnpm globally via npm..." -ForegroundColor Yellow
        npm install -g pnpm
    }

    if (-not (Test-Path .\web-frontend)) {
        Write-Host "Cloning web frontend repo..."
        git clone https://github.com/SolomonLeon/web-rwkv-realweb.git web-frontend
    }

    Push-Location web-frontend
    pnpm install
    if ($env:MODELSCOPE_ENVIRONMENT -eq "studio") {
        pnpm run build --mode target-rwkv-modelscope-space
    } else {
        pnpm run build --mode target-rwkv-hf-space
    }
    Pop-Location

    # Copy dist to the project's dist-frontend
    if (Test-Path .\web-frontend\dist) {
        Remove-Item -Recurse -Force .\dist-frontend -ErrorAction SilentlyContinue
        Copy-Item -Recurse .\web-frontend\dist .\dist-frontend
        Write-Host "Frontend built and copied to ./dist-frontend" -ForegroundColor Green
    }
}

Write-Host "Setup complete. Run the app with: \n$env:CONFIG_FILE='$CONFIG_FILE'\npython app.py" -ForegroundColor Cyan