#!/usr/bin/pwsh

# Define parameters
param(
    [Alias('h')][switch]$Help,
    [Alias('m')][string]$HomeDir=""
)

# Display help if requested
if ($Help) {
    Write-Output "DESCRIPTION: Setup Anypoint Monitoring OpenTelemetry Collector"
    Write-Output "Usage: .\setup.ps1 [options]"
    Write-Output " "
    Write-Output "options:"
    Write-Output "-h          show help and exit"
    Write-Output "-m [DIR]    set home directory - defaults to MULE_HOME"
    Write-Output " "
    exit 0
}

if (-not $HomeDir -or -not (Test-Path $HomeDir)) {
    if ($env:MULE_HOME -and (Test-Path $env:MULE_HOME)) {
        $HomeDir = $env:MULE_HOME
    } else {
        $scriptDir = Split-Path -Parent $MyInvocation.MyCommand.Definition
        $HomeDir = (Resolve-Path -Path (Join-Path -Path $scriptDir -ChildPath "..\..")).Path
        Write-Host "Determined the MULE_HOME"
    }
}

# Validate that serviceHelper.ps1 exists
$serviceHelperPath = Join-Path -Path $HomeDir -ChildPath "am\bin\tools\serviceHelper.ps1"

if (-not (Test-Path $serviceHelperPath)) {
    Write-Host "Error: Invalid MULE_HOME directory. Could not find $serviceHelperPath"
    Write-Host "Please specify a valid MULE_HOME directory with the -m parameter."
    exit 1
}

# Source the serviceHelper.ps1
. $serviceHelperPath

# Java program running on Windows treats paths differently when the trailing character is \
# If MULE_HOME env is not set, HomeDir is autodetected and things are merry.
# But if HomeDir is supplied as an option, with a valid path but ending in \, then AgentConfiguratorApp complains:
# ERROR: The provided Mule Home doesn't exist, please check if the provided directory is correct.
# Hence the TrimEnd
$HomeDir = $HomeDir.TrimEnd('\','/')
$env:MULE_HOME = $HomeDir

Log-Info "MULE_HOME is set to $HomeDir"

# Set AM_HOME based on the determined HomeDir
$env:AM_HOME = Join-Path -Path $HomeDir -ChildPath "am"
$env:AM_CONFIG = Join-Path -Path $env:AM_HOME -ChildPath "config"
$env:AM_CONFIG_PIPELINES = Join-Path -Path $env:AM_HOME -ChildPath "config\pipelines"

$env:OTEL_COLLECTOR_EXE = Join-Path $env:AM_HOME "otel-collector\windows\otel-collector.exe"
$env:OTEL_COLLECTOR_DATA = Join-Path -Path $HomeDir -ChildPath ".mule\.am\otel-collector\data"
$env:OTEL_COLLECTOR_LOGS = Join-Path -Path $env:AM_HOME -ChildPath "logs"

# Configure GOMEMLIMIT using shared function
Set-GoMemLimit

$serviceScript = Join-Path -Path $env:AM_HOME -ChildPath "bin\am-service.ps1"
$processScript = Join-Path -Path $env:AM_HOME -ChildPath "bin\am.ps1"

Log-Info "AM_HOME is set to $env:AM_HOME"

# Source the message.ps1
. (Join-Path -Path $env:AM_HOME -ChildPath "bin\tools\message.ps1")

# Number of attempts to check for the configuration file
$attempts = 30

Log-Info $env:AM_CONFIG_WAIT

while ($true) {
    $attempts--

    if ($attempts -gt 0) {
        $configPath = Join-Path -Path $env:AM_HOME -ChildPath "config/pipelines/otel-config.yml"
        Log-Info "Checking for configuration file at: $configPath"

        if (Test-Path $configPath) {

            Log-Info $env:AM_CONFIG_READY

            $mode = Get-InstallMode
            if ($mode -eq "not_installed") {
                Log-Info "Error: Installation mode is not selected. Please run the install.ps1 script first."
                exit 1
            }
            
            Log-Info "Configuration files will be processed automatically when OTEL collector starts"
            
            if ($mode -eq "process") {
                # handle as a process
                try {
                    Log-Info "Starting standalone otel-collector using: $processScript"
                    if (-not (Test-Path $processScript)) {
                        Log-Info "Error: AM script not found at: $processScript"
                        exit 1
                    }

                    $amProcessLog = Join-Path $env:AM_HOME "logs\otel-collector.log"
                    $amProcessErr = Join-Path $env:AM_HOME "logs\otel-collector.err"
                    $process = Start-Process -FilePath "powershell.exe" `
                            -ArgumentList "-File `"$processScript`" restart -m `"$HomeDir`"" `
                            -NoNewWindow `
                            -PassThru `
                            -RedirectStandardOutput $amProcessLog `
                            -RedirectStandardError $amProcessErr

                    Log-Info "Waiting 5 seconds..."
                    Start-Sleep -Seconds 5

                    # Check if both components started successfully
                    $maxRetries = 3
                    $retryCount = 0
                    $retryInterval = 2 # seconds
                    $bothRunning = $false

                    while ($retryCount -lt $maxRetries) {
                        $otelRunning = Is-Otel-Collector-Running-As-A-Process
                        $dwRunning = Is-DirectoryWatcher-Running-As-A-Process
                        
                        if ($otelRunning -and $dwRunning) {
                            $otelProcess = Get-Process -Name "otel-collector" -ErrorAction SilentlyContinue
                            Log-Info "Both components started successfully - OTEL PID: $($otelProcess.Id)"
                            $bothRunning = $true
                            break # Exit the retry loop on success
                        } else {
                            if ($retryCount -lt $maxRetries) {
                                $missingComponents = @()
                                if (-not $otelRunning) { $missingComponents += "OTEL Collector" }
                                if (-not $dwRunning) { $missingComponents += "DirectoryWatcher" }
                                Log-Info "Attempt $($retryCount+1)/$($maxRetries): Missing components: $($missingComponents -join ', '). Retrying in $retryInterval seconds..."
                                Start-Sleep -Seconds $retryInterval
                            }
                            $retryCount++
                        }
                    }

                    if (-not $bothRunning) {
                        Log-Info "Couldn't identify if both components are running within the specified timeout and after $maxRetries attempts, please check the status by running: am.ps1 status"
                        Log-Info "Additionally please check the logs at: $env:AM_HOME\logs\otel-collector.log, $env:AM_HOME\logs\otel-collector.err, $env:AM_HOME\logs\directorywatcher.log, and $env:AM_HOME\logs\directorywatcher.err"
                    }


                } catch {
                    Log-Info "Error: Failed to setup the monitoring components."
                    exit 1
                }

            } else {
                # handle as a service
                try {
                    Log-Info "Starting service using: $serviceScript"
                    if (-not (Test-Path $serviceScript)) {
                        Log-Info "Error: am-service script not found at: $serviceScript"
                        exit 1
                    }

                    $amServiceLog = Join-Path $env:AM_HOME "logs\otel-collector.log"
                    $amServiceErr = Join-Path $env:AM_HOME "logs\otel-collector.err"
                    $process = Start-Process -FilePath "powershell.exe" `
                            -ArgumentList "-File `"$serviceScript`" -Action restart" `
                            -NoNewWindow `
                            -Wait `
                            -RedirectStandardOutput $amServiceLog `
                            -RedirectStandardError $amServiceErr
                    Log-Info "Service restart command completed. Waiting 2 seconds for services to stabilize..."
                    Start-Sleep -Seconds 2

                    $maxRetries = 3
                    $retryCount = 0
                    $retryInterval = 2 # seconds
                    $bothServicesRunning = $false

                    while ($retryCount -lt $maxRetries) {
                        $otelRunning = Is-Otel-Collector-Running-As-A-Service
                        $dwRunning = Is-DirectoryWatcher-Running-As-A-Service
                        
                        if ($otelRunning -and $dwRunning) {
                            Log-Info "Both services started successfully."
                            $bothServicesRunning = $true
                            break # Exit the retry loop on success
                        } else {
                            if ($retryCount -lt $maxRetries) {
                                $missingServices = @()
                                if (-not $otelRunning) { $missingServices += "OTEL Collector" }
                                if (-not $dwRunning) { $missingServices += "DirectoryWatcher" }
                                Log-Info "Attempt $($retryCount+1)/$($maxRetries): Missing services: $($missingServices -join ', '). Retrying in $retryInterval seconds..."
                                Start-Sleep -Seconds $retryInterval
                            }
                            $retryCount++
                        }
                    }

                    if (-not $bothServicesRunning) {
                        Log-Info "Couldn't confirm both services are running within the specified timeout and after $maxRetries attempts, please check the status by running: am-service.ps1 status"
                        Log-Info "Additionally please check the logs at: $env:AM_HOME\logs\otel-collector.log and $env:AM_HOME\logs\otel-collector.err"
                    }

                } catch {
                    Log-Info "Error: Failed to setup the service am-otel-collector."
                    exit 1
                }
            }

            # DirectoryWatcher is now started together with OTEL Collector via am.ps1 restart
            # No separate DirectoryWatcher setup needed - both components are managed together

            break
        } else {
            Log-Info "Configuration file not found, waiting 10 seconds... (Attempts remaining: $attempts)"
            Start-Sleep -Seconds 10
        }
    } else {
        Log-Info $env:AM_CONFIG_NOT_FOUND
        break
    }
}