#!/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"
    Write-Output "-m,       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 {
        $HomeDir = (Get-Item -Path (Join-Path -Path (Split-Path -Path $MyInvocation.MyCommand.Definition -Parent) -ChildPath "../..")).FullName
    }

    # Source the serviceHelper.ps1
    . (Join-Path -Path $HomeDir -ChildPath "am\bin\tools\serviceHelper.ps1")
}

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")

Convert-AllYamlPaths -PipelinesDir $env:AM_CONFIG_PIPELINES | Out-Null

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

                    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)"
                    } else {
                        Log-Info "Failed to start standalone otel-collector."
                        Log-Info "Check the logs at: $env:AM_HOME\logs\otel-collector.log and $env:AM_HOME\logs\otel-collector.err"
                        exit 1
                    }

                } 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

                    if (Is-Otel-Collector-Running-As-A-Service) {
                        Log-Info "Service am-otel-collector started."
                    } else {
                        Log-Info "Failed to start the service am-otel-collector."
                        Log-Info "Check the logs at: $env:AM_HOME\logs\otel-collector.log and $env:AM_HOME\logs\otel-collector.err"
                        exit 1
                    }
                } catch {
                    Log-Info "Error: Failed to setup the service am-otel-collector."
                    exit 1
                }
            }

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

            # 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"
                # Start-Process -FilePath "java" -ArgumentList "-D`"muleHome=$HomeDir`" -cp `"$classpath`" com.mulesoft.dias.mule.agent.DirectoryWatcher" -RedirectStandardError (Join-Path -Path $env:AM_HOME -ChildPath "logs\otel-collector.err") -NoNewWindow
                try {
                    $dwProcess = Start-Process -FilePath "java" `
                            -ArgumentList "-D`"muleHome=$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
    }
}
