# Trigger Picasso Evaluation GitHub Action
# Usage: .\trigger_picasso_eval.ps1

[CmdletBinding()]
param(
    [string]$BraintrustOrg = "non-pme",
    [string]$PicassoEnv = "runner",
    [string]$ConversationMode = "DefaultMode",
    [string]$TaskId = "",
    [string]$EvalTags = "",
    [string]$FeatureFlags = "",
    [bool]$ShouldOverrideFeatureFlags = $true,
    [bool]$RunBaselineAndTreatment = $false,
    [string]$PicassoBranchUrl = "",
    [string]$EvalRunnerCmdArgs = "",
    [string]$Branch = "",
    [string]$Model = "",
    [switch]$WatchLogs,
    [switch]$DownloadLogs,
    [string]$LogOutputPath = ""
)

$ErrorActionPreference = "Stop"

# Repository for GitHub Actions
$Repo = "infinity-microsoft/picasso"

Write-Host "Triggering Picasso Evaluation workflow..." -ForegroundColor Cyan
Write-Host "  Branch: $Branch" -ForegroundColor Gray
Write-Host "  Braintrust Org: $BraintrustOrg" -ForegroundColor Gray
Write-Host "  Picasso Env: $PicassoEnv" -ForegroundColor Gray
Write-Host "  Conversation Mode: $ConversationMode" -ForegroundColor Gray
Write-Host "  Task ID: $TaskId" -ForegroundColor Gray
Write-Host "  Model: $Model" -ForegroundColor Gray

$args = @(
    "workflow", "run", "aipa-evals-picasso-eval.yml",
    "--repo", $Repo,
    "--ref", $Branch,
    "-f", "braintrust_org=$BraintrustOrg",
    "-f", "picasso_env=$PicassoEnv",
    "-f", "conversation_mode=$ConversationMode",
    "-f", "task_id=$TaskId",
    "-f", "model=$Model",
    "-f", "should_override_feature_flags=$($ShouldOverrideFeatureFlags.ToString().ToLower())",
    "-f", "run_baseline_and_treatment=$($RunBaselineAndTreatment.ToString().ToLower())"
)

if ($EvalTags) {
    $args += @("-f", "eval_tags=$EvalTags")
    Write-Host "  Eval Tags: $EvalTags" -ForegroundColor Gray
}

if ($FeatureFlags) {
    $args += @("-f", "feature_flags=$FeatureFlags")
    Write-Host "  Feature Flags: $FeatureFlags" -ForegroundColor Gray
}

if ($EvalRunnerCmdArgs) {
    $args += @("-f", "eval_runner_cmd_args=$EvalRunnerCmdArgs")
    Write-Host "  Additional Args: $EvalRunnerCmdArgs" -ForegroundColor Gray
}

$maxRetries = 3
$retryCount = 0
$experimentIdFound = $false
$linkFound = $false
$runId = $null

while ($retryCount -lt $maxRetries -and (-not ($experimentIdFound -and $linkFound) -or $retryCount -eq 0)) {
    if ($retryCount -gt 0) {
        Write-Host "`nRetry attempt $retryCount of $($maxRetries - 1)..." -ForegroundColor Yellow
        Start-Sleep -Seconds 10
    }
    
    $startTime = Get-Date
    
    gh @args

    if ($LASTEXITCODE -eq 0) {
        Write-Host "`nWorkflow triggered successfully!" -ForegroundColor Green
        
        # Wait for GitHub to register the run
        Write-Host "Waiting for run to be registered..." -ForegroundColor Yellow
        Start-Sleep -Seconds 5
        
        # Get the most recent run ID for this branch
        $runId = gh run list --workflow=aipa-evals-picasso-eval.yml --repo $Repo --branch=$Branch --limit 1 --json databaseId --jq '.[0].databaseId' 2>$null
        
        if ($runId) {
            Write-Host "Run ID: $runId" -ForegroundColor Green
            Write-Host "View run: gh run view $runId --repo $Repo" -ForegroundColor Yellow
            Write-Host "Watch run: gh run watch $runId --repo $Repo" -ForegroundColor Yellow
            Write-Host "View in browser: gh run view $runId --repo $Repo --web" -ForegroundColor Yellow
            
            # Watch logs if requested
            if ($WatchLogs) {
                Write-Host "`nWatching logs in real-time..." -ForegroundColor Cyan
                gh run watch $runId --repo $Repo
            }

            # Download logs if requested
            if ($DownloadLogs) {
                $logFile = if ($LogOutputPath) { $LogOutputPath } else { "run_${runId}_logs.txt" }
                Write-Host "`nDownloading logs to $logFile..." -ForegroundColor Cyan
                
                # Wait for the run to complete (no timeout, wait until it finishes)
                Write-Host "Waiting for run to complete..." -ForegroundColor Yellow
                $status = ""
                $conclusion = ""
                $startWaitTime = Get-Date
                
                while ($status -notin @("completed", "failure", "cancelled", "timed_out")) {
                    Start-Sleep -Seconds 30
                    $runInfo = gh run view $runId --repo $Repo --json status,conclusion 2>$null | ConvertFrom-Json
                    if ($runInfo) {
                        $status = $runInfo.status
                        $conclusion = $runInfo.conclusion
                        $elapsed = [math]::Round(((Get-Date) - $startWaitTime).TotalSeconds)
                        Write-Host "  Status: $status" -ForegroundColor Gray
                    }
                }

                # If run was cancelled, exit immediately without retrying - no logs will be available
                if ($conclusion -eq "cancelled") {
                    Write-Host "Run $runId was cancelled. Exiting without retry." -ForegroundColor Yellow
                    break
                }
                
                # Calculate total run time
                $totalRunTime = ((Get-Date) - $startWaitTime).TotalMinutes
                Write-Host "Run completed in $([math]::Round($totalRunTime, 2)) minutes" -ForegroundColor Cyan
                
                # Download the logs - handle errors for cancelled/incomplete jobs
                Write-Host "Downloading logs..." -ForegroundColor Gray
                $logDownloaded = $false
                $logError = $null
                
                # Temporarily change error action to continue so we can handle gh errors
                $oldErrorAction = $ErrorActionPreference
                $ErrorActionPreference = "Continue"
                
                try {
                    # Try to download full run logs
                    $logOutput = & gh run view $runId --repo $Repo --log 2>&1
                    $ghExitCode = $LASTEXITCODE
                    
                    if ($ghExitCode -eq 0 -and $logOutput -and $logOutput.Trim()) {
                        $logOutput | Out-File -FilePath $logFile -Encoding utf8
                        $logDownloaded = $true
                        Write-Host "Logs saved to: $logFile" -ForegroundColor Green
                    } else {
                        $logError = "gh run view --log returned empty or failed with exit code $ghExitCode"
                        Write-Host "Warning: $logError" -ForegroundColor Yellow
                        
                        # Fallback: Download logs via API (returns zip format)
                        Write-Host "Attempting to download logs via GitHub API (zip format)..." -ForegroundColor Yellow
                        $tempZip = "run_${runId}_logs.zip"
                        $tempDir = "run_${runId}_logs_extracted"
                        
                        try {
                            $repoInfo = gh repo view $Repo --json nameWithOwner --jq '.nameWithOwner' 2>$null
                            if (-not $repoInfo) {
                                $repoInfo = $Repo
                            }
                            $token = gh auth token
                            $apiUrl = "https://api.github.com/repos/$repoInfo/actions/runs/$runId/logs"
                            
                            Write-Host "  Downloading from: $apiUrl" -ForegroundColor Gray
                            Invoke-WebRequest -Uri $apiUrl -Headers @{Authorization="Bearer $token"} -OutFile $tempZip
                            
                            if (Test-Path $tempZip) {
                                Write-Host "  Extracting zip archive..." -ForegroundColor Gray
                                Expand-Archive -Path $tempZip -DestinationPath $tempDir -Force
                                
                                # Merge all log files
                                $logFiles = Get-ChildItem -Path $tempDir -Recurse -Filter "*.txt"
                                if ($logFiles) {
                                    $logFiles | Get-Content | Out-File -FilePath $logFile -Encoding utf8
                                    $logDownloaded = $true
                                    Write-Host "Logs extracted and saved to: $logFile" -ForegroundColor Green
                                }
                            }
                        } catch {
                            Write-Host "  API download failed: $($_.Exception.Message)" -ForegroundColor Yellow
                        } finally {
                            # Cleanup temp files
                            if (Test-Path $tempZip) { Remove-Item $tempZip -Force }
                            if (Test-Path $tempDir) { Remove-Item $tempDir -Recurse -Force }
                        }
                        
                        # If API download also failed, try downloading logs from individual completed jobs
                        if (-not $logDownloaded) {
                            Write-Host "Attempting to download logs from individual jobs..." -ForegroundColor Gray
                            $jobsJson = & gh run view $runId --repo $Repo --json jobs 2>$null
                            if ($jobsJson) {
                                $jobsData = $jobsJson | ConvertFrom-Json
                                $completedJobs = $jobsData.jobs | Where-Object { $_.conclusion -ne $null -and $_.conclusion -ne "" }
                                
                                $allJobLogs = @()
                                foreach ($job in $completedJobs) {
                                    Write-Host "  Fetching logs for job: $($job.name)..." -ForegroundColor Gray
                                    $jobLog = & gh run view --repo $Repo --job $($job.databaseId) --log 2>$null
                                    if ($LASTEXITCODE -eq 0 -and $jobLog) {
                                        $allJobLogs += "=== Job: $($job.name) ==="
                                        $allJobLogs += $jobLog
                                    }
                                }
                                
                                if ($allJobLogs.Count -gt 0) {
                                    $allJobLogs | Out-File -FilePath $logFile -Encoding utf8
                                    $logDownloaded = $true
                                    Write-Host "Job logs saved to: $logFile" -ForegroundColor Green
                                }
                            }
                        }
                    }
                } catch {
                    $logError = $_.Exception.Message
                    Write-Host "Warning: Error downloading logs: $logError" -ForegroundColor Yellow
                } finally {
                    $ErrorActionPreference = $oldErrorAction
                }
                
                # If log download failed completely, trigger retry
                if (-not $logDownloaded) {
                    Write-Host "Failed to download any logs for run $runId" -ForegroundColor Red
                    $retryCount++
                    
                    if ($retryCount -lt $maxRetries) {
                        Write-Host "Will retry workflow trigger..." -ForegroundColor Yellow
                        continue
                    } else {
                        Write-Host "Max retries reached. Could not get logs." -ForegroundColor Red
                        break
                    }
                }
                
                # Extract and display Run Evals section
                Write-Host "`nExtracting 'Run Evals' step output..." -ForegroundColor Cyan
                $evalLogs = Select-String -Path $logFile -Pattern "Run Evals" -Context 0,500 -ErrorAction SilentlyContinue
                if ($evalLogs) {
                    Write-Host "`n=== Run Evals Output ===" -ForegroundColor Green
                    $evalLogs | ForEach-Object { $_.Context.PostContext } | Select-Object -First 100
                }
                
                # Search for Experiment ID and Braintrust links
                Write-Host "`nSearching for Experiment ID and Braintrust links..." -ForegroundColor Cyan
                $responseBodyPattern = 'Response body:\s*(\{.*\})'
                $braintrustLinkPattern = "https://www\.braintrust\.dev/app/[a-zA-Z0-9_%/\-]+"
                
                $responseBodyMatches = Select-String -Path $logFile -Pattern $responseBodyPattern -AllMatches -ErrorAction SilentlyContinue
                $linkMatches = Select-String -Path $logFile -Pattern $braintrustLinkPattern -AllMatches -ErrorAction SilentlyContinue
                
                # Search for Machine Evals Dashboard URL (when EvalTags is set)
                $machineEvalsDashboardUrl = $null
                if ($EvalTags) {
                    $machineEvalsPattern = "Machine Evals Dashboard URL:\s*(https://[^\s]+)"
                    $machineEvalsMatches = Select-String -Path $logFile -Pattern $machineEvalsPattern -AllMatches -ErrorAction SilentlyContinue
                    if ($machineEvalsMatches) {
                        foreach ($match in $machineEvalsMatches) {
                            if ($match.Matches) {
                                foreach ($m in $match.Matches) {
                                    $machineEvalsDashboardUrl = $m.Groups[1].Value
                                    break
                                }
                            }
                            if ($machineEvalsDashboardUrl) { break }
                        }
                    }
                }
                
                $experimentIds = @()
                if ($responseBodyMatches) {
                    foreach ($match in $responseBodyMatches) {
                        if ($match.Matches) {
                            foreach ($m in $match.Matches) {
                                $jsonString = $m.Groups[1].Value
                                try {
                                    $jsonObj = $jsonString | ConvertFrom-Json
                                    if ($jsonObj.successful_experiment_ids) {
                                        $experimentIds += $jsonObj.successful_experiment_ids
                                    }
                                } catch {
                                    # JSON parsing failed, skip this match
                                }
                            }
                        }
                    }
                }
                if ($experimentIds.Count -gt 0) {
                    Write-Host "`n=== Experiment IDs Found ===" -ForegroundColor Green
                    Write-Host "Experiment ID:$($experimentIds -join ',')" -ForegroundColor Yellow
                    $experimentIdFound = $true
                }
                if ($linkMatches) {
                    Write-Host "`n=== Braintrust Links Found ===" -ForegroundColor Green
                    foreach ($match in $linkMatches) {
                        Write-Host $match.Line -ForegroundColor Yellow
                    }
                    $linkFound = $true
                }
                
                # Display Machine Evals Dashboard URL if found
                if ($machineEvalsDashboardUrl) {
                    Write-Host "`n=== Machine Evals Dashboard URL ===" -ForegroundColor Green
                    Write-Host "Machine Evals Dashboard URL:$machineEvalsDashboardUrl" -ForegroundColor Yellow
                } elseif ($EvalTags) {
                    Write-Host "`nWarning: EvalTags was set but Machine Evals Dashboard URL not found in logs" -ForegroundColor Yellow
                }
                
                # Check if both were found
                if (-not ($experimentIdFound -and $linkFound)) {
                    Write-Host "`nWarning: Failed to extract Experiment ID or Braintrust links" -ForegroundColor Yellow
                    
                    # Never retry a cancelled run - it was intentionally stopped
                    if ($conclusion -eq "cancelled") {
                        Write-Host "Run was cancelled - skipping retry" -ForegroundColor Yellow
                        break
                    }
                    
                    # If job completed within 15 minutes AND extraction failed, consider it a failed run
                    if ($totalRunTime -lt 15) {
                        Write-Host "Job completed too quickly ($([math]::Round($totalRunTime, 2)) min) without expected output - likely failed" -ForegroundColor Yellow
                        $retryCount++
                        
                        if ($retryCount -lt $maxRetries) {
                            Write-Host "Will retry workflow trigger..." -ForegroundColor Yellow
                        } else {
                            Write-Host "Max retries reached. Final logs saved to: $logFile" -ForegroundColor Red
                        }
                    } else {
                        Write-Host "Job ran for $([math]::Round($totalRunTime, 2)) minutes but extraction failed - may be extraction issue" -ForegroundColor Yellow
                        Write-Host "Logs saved to: $logFile for investigation" -ForegroundColor Yellow
                        break
                    }
                } else {
                    Write-Host "`nSuccessfully extracted Experiment ID and Braintrust links!" -ForegroundColor Green
                    
                    # Clean up log file after successful extraction
                    if (Test-Path $logFile) {
                        Remove-Item -Path $logFile -Force
                        Write-Host "Log file deleted: $logFile" -ForegroundColor Gray
                    }
                    
                    break
                }
            } else {
                # If not downloading logs, just return the run ID
                break
            }
        } else {
            Write-Host "Check status with: gh run list --workflow=aipa-evals-picasso-eval.yml --repo $Repo --limit 5" -ForegroundColor Yellow
            $retryCount++
        }
    } else {
        Write-Host "`nFailed to trigger workflow. Exit code: $LASTEXITCODE" -ForegroundColor Red
        $retryCount++
    }
}

if ($runId) {
    return $runId
}
