# Import OpenAI API Key Powershell Script
# Simple script to pull the defined key from CS metadata and set in a prescribed location
# Need to run 'Set-ExecutionPolicy Unrestricted' in an Admin Powershell prompt before scripts are allowed to run
# Run with 'powershell.exe -ExecutionPolicy Bypass -WindowStyle Hidden -File <path to ps file> -Silent' if using scheduled task
# Make sure it is set to only run when user is logs on

param(
    [switch]$Silent
)

$logfile = Join-Path $PSScriptRoot "OpenAI_Import.log"

function Log {
    param([string]$msg)
    Write-Host $msg
    Add-Content -Path $logfile -Value $msg
}

Log "`r`n===== Starting OpenAI Import (Metadata Mode) -- $(Get-Date) ====="

# ------------------------------------------------------------
# STEP 1: Check if installed as a scheduled task and if not...
# ------------------------------------------------------------

$taskFolder = '\Learning Tree\'
$taskName   = 'OpenAI Import'
$existingTask = Get-ScheduledTask -TaskName $taskName -TaskPath $taskFolder -ErrorAction SilentlyContinue

if (-not $existingTask) {
    $taskAction = New-ScheduledTaskAction `
        -Execute 'powershell.exe' `
        -Argument "-ExecutionPolicy Bypass -WindowStyle Hidden -File `"$PSCommandPath`" -AutoRename"
    $taskTrigger   = New-ScheduledTaskTrigger -AtLogOn
    $taskPrincipal = New-ScheduledTaskPrincipal `
        -UserId "$env:USERDOMAIN\$env:USERNAME" `
        -LogonType Interactive
    Register-ScheduledTask `
        -TaskName $taskName `
        -TaskPath $taskFolder `
        -Action $taskAction `
        -Trigger $taskTrigger `
        -Principal $taskPrincipal `
        -Description 'OpenAI Key Import at logon'
    (New-Object -ComObject WScript.Shell).Popup(
        "Scheduled task '$taskName' created for user $env:USERDOMAIN\$env:USERNAME.",
        0,
        "Task Created",
        64
    ) | Out-Null
    Log "Scheduled task '$taskName' was added to the scheduler."
    [System.Windows.Forms.MessageBox]::Show(
        "Scheduled task was installed so will run a logon",
        "OpenAI Key"
    )
    Exit 0
}
else {
    Log "Scheduled task '$taskName' already exists - skipping."
}

# ------------------------------------------------------------
# STEP 2: Get metadata directly from CS
# ------------------------------------------------------------

try {
    $metadata = Invoke-RestMethod -Uri "https://metadata.cloudshare.com/api/v3/unauthenticated/metadata"

    $OPENAI = $metadata."custom-properties".openAIKey
    $CLASSID = $metadata."class-id"
    # $CLASSID not used, but useful to see if updated on this VM. Could use to bypass if second boot
} catch {
    Log "Failed to retrieve metadata"
    exit 1
}

if (-not $OPENAI -or $OPENAI -eq "null") {
    Log "No OpenAI key found in metadata"
    exit 1
}

Log "Class ID: $CLASSID"
Log "OpenAI Key: $OPENAI"

# ------------------------------------------------------------
# STEP 3: Optional GUI prompt
# ------------------------------------------------------------

if (-not $Silent) {
    Add-Type -AssemblyName System.Windows.Forms
    $result = [System.Windows.Forms.MessageBox]::Show(
        "Install OpenAI key from metadata?",
        "OpenAI Key",
        "YesNo"
    )

    if ($result -ne "Yes") {
        Log "User cancelled"
        exit
    }
}

# ------------------------------------------------------------
# STEP 4: Write key to standard locations
# ------------------------------------------------------------

$keyFiles = @(
    "C:\course\openai_key.txt",
    "C:\Users\Student\keys\key.txt"
)

foreach ($file in $keyFiles) {
    $dir = Split-Path $file
    if (!(Test-Path $dir)) {
        New-Item -ItemType Directory -Path $dir | Out-Null
    }
    "OPENAI_API_KEY=$OPENAI" | Set-Content -Path $file
    Log "Wrote key to: $file"
}

# ------------------------------------------------------------
# STEP 5: Update any config files (if present)
# ------------------------------------------------------------

$configFiles = @(
    "C:\Users\Student\.config\shell_gpt\.sgptrc"
)
# Possible future use, but not for 1264
foreach ($file in $configFiles) {
    if (Test-Path $file) {
        (Get-Content $file) -replace "OPENAI_API_KEY=.*", "OPENAI_API_KEY=$OPENAI" |
            Set-Content $file

        Log "Updated config file: $file"
    }
}

# ------------------------------------------------------------
# STEP 6: Optional confirmation
# ------------------------------------------------------------

#if (-not $Silent) {
    [System.Windows.Forms.MessageBox]::Show(
        "The new OpenAI key was installed successfully",
        "OpenAI Key"
    )
#}

Log "Completed successfully"
