#!/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"

# Check for system-wide GOMEMLIMIT and set if not present
$currentGoMemLimit = [System.Environment]::GetEnvironmentVariable("GOMEMLIMIT", [System.EnvironmentVariableTarget]::Machine)

if ($currentGoMemLimit) {
    $env:GOMEMLIMIT = $currentGoMemLimit
    Log-Info "System-wide GOMEMLIMIT is already set to $currentGoMemLimit. Reusing it."
} else {
    # Compute GOMEMLIMIT based on OTEL_PROCESSOR_MEMORY_LIMIT_MB or default to 40
    if ($env:OTEL_PROCESSOR_MEMORY_LIMIT_MB) {
        $memLimit = [int]$env:OTEL_PROCESSOR_MEMORY_LIMIT_MB
    } else {
        # Inspired by otel-pipeline-common.yml
        # processors:
        #  memory_limiter:
        #    check_interval: 1s
        #    limit_mib: 40
        $memLimit = 40
    }
    $goMemLimit = [math]::Floor($memLimit * 0.67)
    $gomemlimitValue = "${goMemLimit}MiB"

    # Set as system-wide (machine) environment variable
    [System.Environment]::SetEnvironmentVariable("GOMEMLIMIT", $gomemlimitValue, [System.EnvironmentVariableTarget]::Machine)
    $env:GOMEMLIMIT = $gomemlimitValue
    Log-Info "System-wide GOMEMLIMIT set to $gomemlimitValue"
}

$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) {

            # Check if Java is available
            if (-not (Get-Command java -ErrorAction SilentlyContinue)) {
                Log-Info "Error: Java is not installed or not in PATH"
                exit 1
            }

            # Construct the classpath
            $classpath = Join-Path -Path $env:AM_HOME -ChildPath "lib\*"
            Log-Info "Using classpath: $classpath"

            # Invoke FileProcessingUtils to process the config files.
            Log-Info "Invoking FileProcessingUtils to process $env:AM_CONFIG_PIPELINES"
            # For FileProcessingUtils
            $fpuProcess = Start-Process -FilePath "java" `
                                        -ArgumentList @(
                                                    "-DmuleHome=$HomeDir"
                                                    "-cp"
                                                    "$classpath"
                                                    "com.mulesoft.dias.util.FileProcessingUtils"
                                                    "$env:AM_CONFIG_PIPELINES"
                                                ) `
                                        -RedirectStandardError (Join-Path -Path $env:AM_HOME -ChildPath "logs\fileprocessingutils.err") `
                                        -RedirectStandardOutput (Join-Path -Path $env:AM_HOME -ChildPath "logs\fileprocessingutils.log") `
                                        -NoNewWindow `
                                        -Wait `
                                        -PassThru

            if ($fpuProcess.ExitCode -ne 0) {
                $errLog = Join-Path -Path $env:AM_HOME -ChildPath "logs\fileprocessingutils.err"
                Log-Info "FileProcessingUtils failed with exit code: $($fpuProcess.ExitCode)"
                Log-Info "Halting otel collector setup."
                if (Test-Path $errLog) {
                    Log-Info "Error log contents:"
                    Get-Content $errLog | Write-Output
                }
                exit $fpuProcess.ExitCode
            } else {
                Log-Info "FileProcessingUtils completed successfully with exit code: 0"
            }

            Log-Info $env:AM_CONFIG_READY

            $mode = Get-InstallMode
            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

                    # Retry checking is otel-collector started as a process
                    $maxRetries = 3
                    $retryCount = 0
                    $retryInterval = 2 # seconds
                    $collectorRunning = $false

                    while ($retryCount -lt $maxRetries) {
                        if (Is-Otel-Collector-Running-As-A-Process) {
                            $process = Get-Process -Name "otel-collector" -ErrorAction SilentlyContinue
                            Log-Info "Standalone otel-collector started with PID $($process.Id)"
                            $collectorRunning = $true
                            break # Exit the retry loop on success
                        } else {
                            if ($retryCount -lt $maxRetries) {
                                Log-Info "Attempt $($retryCount+1)/$($maxRetries): otel-collector process not detected. Retrying in $retryInterval seconds..."
                                Start-Sleep -Seconds $retryInterval
                            }
                            $retryCount++
                        }
                    }

                    if (-not $collectorRunning) {
                        Log-Info "Couldn't identify if otel-collector process is 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 and $env:AM_HOME\logs\otel-collector.err"
                    }


                } catch {
                    Log-Info "Error: Failed to setup the standalone am-otel-collector."
                    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`" restart -m `"$HomeDir`"" `
                            -NoNewWindow `
                            -PassThru `
                            -RedirectStandardOutput $amServiceLog `
                            -RedirectStandardError $amServiceErr
                    Log-Info "Waiting 5 seconds..."
                    Start-Sleep -Seconds 5

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

                    while ($retryCount -lt $maxRetries) {
                        if (Is-Otel-Collector-Running-As-A-Service) {
                            Log-Info "Service am-otel-collector started."
                            $collectorRunning = $true
                            break # Exit the retry loop on success
                        } else {
                            if ($retryCount -lt $maxRetries) {
                                Log-Info "Attempt $($retryCount+1)/$($maxRetries): am-otel-collector service not detected. Retrying in $retryInterval seconds..."
                                Start-Sleep -Seconds $retryInterval
                            }
                            $retryCount++
                        }
                    }

                    if (-not $collectorRunning) {
                        Log-Info "Couldn't identify if am-otel-collector is 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 and $env:AM_HOME\logs\otel-collector.err"
                    }

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

            # Setting up DirectoryWatcher

            # Find and stop all DirectoryWatcher Java processes
            Get-Process java -ErrorAction SilentlyContinue | Where-Object {
                $_.Path -like "*java.exe" -and
                        (Get-CimInstance Win32_Process -Filter "ProcessId=$($_.Id)").CommandLine -match "com\.mulesoft\.dias\.mule\.agent\.DirectoryWatcher"
            } | ForEach-Object {
                try {
                    Log-Info "Stopping existing DirectoryWatcher process with PID $($_.Id)"
                    Stop-Process -Id $_.Id -Force
                    Log-Info "Successfully stopped DirectoryWatcher process with PID $($_.Id)"
                } catch {
                    Log-Info "Failed to stop DirectoryWatcher process with PID $($_.Id): $_"
                }
            }

            # Start DirectoryWatcher
            try {
                Log-Info "Starting DirectoryWatcher with muleHome: $HomeDir"
                try {
                    $dwProcess = Start-Process -FilePath "java" `
                                                -ArgumentList @(
                                                                    "-DmuleHome=$HomeDir"
                                                                    "-cp"
                                                                    "$classpath"
                                                                    "com.mulesoft.dias.mule.agent.DirectoryWatcher"
                                                                ) `
                                                -RedirectStandardError (Join-Path -Path $env:AM_HOME -ChildPath "logs\directorywatcher.err") `
                                                -RedirectStandardOutput (Join-Path -Path $env:AM_HOME -ChildPath "logs\directorywatcher.log") `
                                                -NoNewWindow `
                                                -PassThru

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

                    if ($dwProcess -and (Get-Process -Id $dwProcess.Id -ErrorAction SilentlyContinue)) {
                        $dwPidFile = Join-Path -Path $env:AM_HOME -ChildPath "directorywatcher.pid"
                        Set-Content -Path $dwPidFile -Value $dwProcess.Id
                        Log-Info "DirectoryWatcher started in background with PID $($dwProcess.Id) (PID also written to $dwPidFile)"
                    } else {
                        Log-Info "Error: DirectoryWatcher process did not start or exited immediately."
                        $errLog = Join-Path -Path $env:AM_HOME -ChildPath "logs\directorywatcher.err"
                        if (Test-Path $errLog) {
                            Log-Info "Error log contents:"
                            Get-Content $errLog | Write-Output
                        }
                        exit 1
                    }


                } catch {
                    Log-Info "Exception occurred while starting DirectoryWatcher: $_"
                    $errLog = Join-Path -Path $env:AM_HOME -ChildPath "logs\directorywatcher.err"
                    if (Test-Path $errLog) {
                        Log-Info "Error log contents:"
                        Get-Content $errLog | Write-Output
                    }
                    exit 1
                }

            } catch {
                Log-Info "Error: Failed to start DirectoryWatcher: $_"
                exit 1
            }

            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
    }
}