Stock-trading-programming/scripts/restore_after_format.ps1

param(
    [switch]$SkipWinget,
    [switch]$SkipCliInstall,
    [switch]$SkipScheduler,
    [switch]$OnlinePip
)

$ErrorActionPreference = "Stop"

[Console]::OutputEncoding = [System.Text.Encoding]::UTF8
$OutputEncoding = [System.Text.Encoding]::UTF8
$env:PYTHONUTF8 = "1"
$env:PYTHONIOENCODING = "utf-8"

$Project = Split-Path -Parent $PSScriptRoot
$VenvPython = Join-Path $Project ".venv\Scripts\python.exe"
$Requirements = Join-Path $Project "requirements.txt"
$Wheelhouse = Join-Path $Project "vendor\wheels"

function Write-Step {
    param([string]$Message)
    Write-Host ""
    Write-Host "== $Message ==" -ForegroundColor Cyan
}

function Test-Admin {
    $Identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $Principal = New-Object Security.Principal.WindowsPrincipal($Identity)
    return $Principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Ensure-Directory {
    param([string]$Path)
    New-Item -ItemType Directory -Force -Path $Path | Out-Null
}

function Install-WingetPackage {
    param(
        [string]$Id,
        [string]$CommandName
    )

    if ($SkipWinget) {
        Write-Host "[skip] winget install disabled for $Id"
        return
    }

    if ($CommandName -and (Get-Command $CommandName -ErrorAction SilentlyContinue)) {
        Write-Host "[ok] $CommandName already available"
        return
    }

    if (-not (Get-Command winget -ErrorAction SilentlyContinue)) {
        Write-Host "[warn] winget not found. Install $Id manually if this step is needed." -ForegroundColor Yellow
        return
    }

    Write-Host "[install] $Id"
    winget install --id $Id -e --accept-package-agreements --accept-source-agreements
    Refresh-ProcessPath
}

function Refresh-ProcessPath {
    $MachinePath = [Environment]::GetEnvironmentVariable("Path", "Machine")
    $UserPath = [Environment]::GetEnvironmentVariable("Path", "User")
    $env:Path = "$MachinePath;$UserPath"
}

function Resolve-BasePython {
    Refresh-ProcessPath

    if (Get-Command py -ErrorAction SilentlyContinue) {
        $Path = (& py -3.11 -c "import sys; print(sys.executable)" 2>$null)
        if ($LASTEXITCODE -eq 0 -and $Path) {
            return $Path.Trim()
        }
    }

    if (Get-Command python -ErrorAction SilentlyContinue) {
        $Version = (& python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')" 2>$null)
        if ($LASTEXITCODE -eq 0 -and $Version.Trim() -eq "3.11") {
            $Path = (& python -c "import sys; print(sys.executable)")
            return $Path.Trim()
        }
    }

    $KnownPaths = @(
        (Join-Path $env:LOCALAPPDATA "Programs\Python\Python311\python.exe"),
        "C:\Program Files\Python311\python.exe",
        "C:\Program Files (x86)\Python311\python.exe"
    )
    foreach ($KnownPath in $KnownPaths) {
        if (Test-Path $KnownPath) {
            return $KnownPath
        }
    }

    throw "Python 3.11 was not found. Re-run without -SkipWinget, or install Python 3.11 manually."
}

function Install-NodeCli {
    if ($SkipCliInstall) {
        Write-Host "[skip] CLI install disabled"
        return
    }

    Refresh-ProcessPath

    if (-not (Get-Command npm -ErrorAction SilentlyContinue)) {
        Write-Host "[warn] npm not found. Claude/Codex CLI install skipped." -ForegroundColor Yellow
        return
    }

    $HasClaude = (Get-Command claude.cmd -ErrorAction SilentlyContinue) -or (Test-Path (Join-Path $env:APPDATA "npm\claude.cmd"))
    $HasCodex = (Get-Command codex.cmd -ErrorAction SilentlyContinue) -or (Test-Path (Join-Path $env:APPDATA "npm\codex.cmd"))

    if (-not $HasClaude) {
        Write-Host "[install] Claude Code CLI"
        npm install -g @anthropic-ai/claude-code
    } else {
        Write-Host "[ok] Claude CLI already available"
    }

    if (-not $HasCodex) {
        Write-Host "[install] Codex CLI"
        npm install -g @openai/codex
    } else {
        Write-Host "[ok] Codex CLI already available"
    }
}

if (-not (Test-Admin)) {
    Write-Host "[warn] Not running as Administrator. Scheduler registration may fail." -ForegroundColor Yellow
}

Set-Location $Project

Write-Step "Project folders"
Ensure-Directory (Join-Path $Project "logs")
Ensure-Directory (Join-Path $Project "data")
Ensure-Directory (Join-Path $Project "models")
Ensure-Directory (Join-Path $Project "reports\daily")
Ensure-Directory (Join-Path $Project "reports\proposals")

if (-not (Test-Path (Join-Path $Project ".env"))) {
    Write-Host "[warn] .env not found. Restore it from your backup before live API use." -ForegroundColor Yellow
}

Write-Step "System tools"
Refresh-ProcessPath
Install-WingetPackage -Id "Git.Git" -CommandName "git"
Install-WingetPackage -Id "OpenJS.NodeJS.LTS" -CommandName "node"
Install-WingetPackage -Id "Python.Python.3.11" -CommandName $null

Write-Step "Python virtual environment"
$BasePython = Resolve-BasePython
Write-Host "[ok] Python base: $BasePython"

if (-not (Test-Path $VenvPython)) {
    & $BasePython -m venv (Join-Path $Project ".venv")
}

& $VenvPython -m ensurepip --upgrade

Write-Step "Python libraries"
if (-not (Test-Path $Requirements)) {
    throw "requirements.txt not found: $Requirements"
}

if ((Test-Path $Wheelhouse) -and -not $OnlinePip) {
    Write-Host "[install] using local wheelhouse: $Wheelhouse"
    & $VenvPython -m pip install --no-index --find-links $Wheelhouse -r $Requirements
} else {
    Write-Host "[install] using online package index"
    & $VenvPython -m pip install -r $Requirements
}

Write-Step "Claude/Codex CLI"
Install-NodeCli

Write-Step "Sanity checks"
& $VenvPython -c "import aiohttp, pandas, sklearn, joblib, holidays; print('python deps ok')"

if (-not $SkipScheduler) {
    Write-Step "Windows Task Scheduler"
    & powershell.exe -NoProfile -ExecutionPolicy Bypass -File (Join-Path $Project "scripts\setup_scheduler.ps1")
} else {
    Write-Host "[skip] scheduler registration disabled"
}

Write-Step "Done"
Write-Host "Restore complete. Reboot once if newly installed winget packages are not visible in new terminals." -ForegroundColor Green