[CmdletBinding()]
param(
    [string]$TaskName = "TradeBot Windows Training Agent",
    [string]$ApiBaseUrl = "https://tb.kusoft.xyz",
    [string]$ApiAuth = "",
    [int]$PollSeconds = 10,
    [int]$WatchdogMinutes = 5,
    [string]$RepoRoot = "",
    [switch]$StartNow,
    [switch]$KeepLegacyRetrainer
)

$ErrorActionPreference = "Stop"

if (-not $RepoRoot) {
    $RepoRoot = (Resolve-Path (Join-Path $PSScriptRoot "..")).Path
}
$Agent = Join-Path $RepoRoot "tools\windows_training_agent.py"
if (-not (Test-Path $Agent)) {
    throw "Windows training agent not found: $Agent"
}

function Resolve-Python {
    $venvPython = Join-Path $RepoRoot ".venv\Scripts\python.exe"
    if (Test-Path $venvPython) {
        return $venvPython
    }

    $userPython = Join-Path $env:LOCALAPPDATA "Programs\TradeBotPython312\python.exe"
    if (Test-Path $userPython) {
        return $userPython
    }

    foreach ($candidate in @("python.exe", "python")) {
        $command = Get-Command $candidate -ErrorAction SilentlyContinue
        if ($command) {
            return $command.Source
        }
    }
    throw "Python was not found. Create .venv or install Python 3.12."
}

function Resolve-WindowlessPython {
    $python = Resolve-Python
    $pythonw = Join-Path (Split-Path -Parent $python) "pythonw.exe"
    if (Test-Path $pythonw) {
        return $pythonw
    }
    return $python
}

if ($ApiAuth) {
    [Environment]::SetEnvironmentVariable("TRADEBOT_API_AUTH", $ApiAuth, "User")
    $env:TRADEBOT_API_AUTH = $ApiAuth
}
[Environment]::SetEnvironmentVariable("TRADEBOT_API_BASE_URL", $ApiBaseUrl, "User")
[Environment]::SetEnvironmentVariable("TRADEBOT_TRAINING_WORKER_NAME", $env:COMPUTERNAME, "User")
$env:TRADEBOT_API_BASE_URL = $ApiBaseUrl
$env:TRADEBOT_TRAINING_WORKER_NAME = $env:COMPUTERNAME

if (-not $KeepLegacyRetrainer) {
    foreach ($legacyName in @("TradeBot PyTorch Forecaster Retrainer", "TradeBot LSTM Retrainer")) {
        $legacyTask = Get-ScheduledTask -TaskName $legacyName -ErrorAction SilentlyContinue
        if ($legacyTask) {
            Unregister-ScheduledTask -TaskName $legacyName -Confirm:$false
            Write-Host "Removed legacy scheduled task '$legacyName'."
        }
    }
}

$python = Resolve-WindowlessPython
$currentUser = [System.Security.Principal.WindowsIdentity]::GetCurrent().Name
$arguments = @(
    "-u",
    "`"$Agent`"",
    "--repo-root", "`"$RepoRoot`"",
    "--api-base-url", "`"$ApiBaseUrl`"",
    "--poll-seconds", $PollSeconds.ToString()
) -join " "

$action = New-ScheduledTaskAction -Execute $python -Argument $arguments -WorkingDirectory $RepoRoot
$trigger = @(
    New-ScheduledTaskTrigger -AtLogOn -User $currentUser
    New-ScheduledTaskTrigger -AtStartup
    New-ScheduledTaskTrigger `
        -Once `
        -At (Get-Date).AddMinutes(1) `
        -RepetitionInterval (New-TimeSpan -Minutes $WatchdogMinutes) `
        -RepetitionDuration (New-TimeSpan -Days 3650)
)
$principal = New-ScheduledTaskPrincipal `
    -UserId $currentUser `
    -LogonType Interactive `
    -RunLevel Limited
$settings = New-ScheduledTaskSettingsSet `
    -StartWhenAvailable `
    -MultipleInstances IgnoreNew `
    -AllowStartIfOnBatteries `
    -DontStopIfGoingOnBatteries `
    -RestartCount 999 `
    -RestartInterval (New-TimeSpan -Minutes 1) `
    -ExecutionTimeLimit (New-TimeSpan -Days 30)

Register-ScheduledTask `
    -TaskName $TaskName `
    -Action $action `
    -Trigger $trigger `
    -Principal $principal `
    -Settings $settings `
    -Description "Keeps the TradeBot Windows training agent online and polls the public bot API for retrain jobs." `
    -Force | Out-Null

if ($StartNow) {
    Start-ScheduledTask -TaskName $TaskName
}

Write-Host "Registered scheduled task '$TaskName' for Windows startup, logon, and watchdog restarts."
Write-Host "Agent API: $ApiBaseUrl"
Write-Host "Agent script: $Agent"