#!/usr/bin/pwsh

# Implemented am.ps1 in a comprehensive way (because it's also invoked from DirectoryWatcher).
# Comprehensive because am.ps1 can determine if otel-collector is running as a service or process
# and then route the appropriate action for start / stop / restart action modes
# The default is process mode. The key to this is in Get-InstallMode function defined in serviceHelper.ps1

param(
    [Parameter(Position=0)]
    [ValidateSet('start', 'stop', 'restart', 'status')]
    [string]$Action,

    [Parameter()]
    [Alias('m')]
    [string]$HomeDir
)

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

# 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

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"
$PIDFILE = Join-Path $env:AM_HOME "am-agent.pid"

#Log-Info "AM_HOME is set to $env:AM_HOME"
#Log-Info "AM_CONFIG is set to $env:AM_CONFIG"
#Log-Info "AM_CONFIG_PIPELINES is set to $env:AM_CONFIG_PIPELINES"
#Log-Info "OTEL_COLLECTOR_EXE is set to $env:OTEL_COLLECTOR_EXE"
#Log-Info "OTEL_COLLECTOR_DATA is set to $env:OTEL_COLLECTOR_DATA"
#Log-Info "OTEL_COLLECTOR_LOGS is set to $env:OTEL_COLLECTOR_LOGS"

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


function Get-ConfigArgs {
    $configFiles = Get-ChildItem -Path $env:AM_CONFIG_PIPELINES -Filter "*.yml"
    if ($configFiles.Count -eq 0) {
        Log-Info "WARN: No configuration files found in $env:AM_CONFIG_PIPELINES"
    }
    $configArgs = ($configFiles | ForEach-Object { "--config=`"$($_.FullName)`"" }) -join " "
    Log-Info "DEBUG: Configuration files found: $($configFiles.Count)"
    Log-Info "DEBUG: Configuration arguments: $configArgs"
    return $configArgs
}

function Start-OtelCollector-Process {
    if (Get-Process -Name "otel-collector" -ErrorAction SilentlyContinue) {
        Log-Info "otel-collector is already running"
        return
    }

    $new_config = Get-ConfigArgs
    $process = Start-Process -FilePath $env:OTEL_COLLECTOR_EXE `
        -ArgumentList $new_config `
        -WindowStyle Hidden `
        -PassThru `
        -RedirectStandardError "$env:OTEL_COLLECTOR_LOGS\otel-collector.err" `
        -RedirectStandardOutput "$env:OTEL_COLLECTOR_LOGS\otel-collector.log"
    $process.Id | Out-File -FilePath $PIDFILE -Force
    Start-Sleep -Seconds 1
    if (Get-Process -Id $process.Id -ErrorAction SilentlyContinue) {
        Log-Info "otel-collector started with PID $($process.Id)"
        return $true
    } else {
        Log-Info "Failed to start otel-collector"
        return $false
    }
}

function Start-OtelCollector-Service {
    Log-Info "Starting service using: $serviceScript"
    Start-Process -FilePath "powershell.exe" -ArgumentList "-File `"$serviceScript`" start" -NoNewWindow -Wait
}

function Start-OtelCollector {
    # Check if am is installed as process or service
    $mode = Get-InstallMode
    if ($mode -eq "process") {
        # handle as a process
        return Start-OtelCollector-Process
    } else {
        # handle as a service
        return Start-OtelCollector-Service
    }
}



function Stop-OtelCollector-Process {
    $success = $true
    # 1. Try to gracefully stop all otel-collector processes by name (like shutdown_otel)
    $allCollectors = Get-Process -Name "otel-collector" -ErrorAction SilentlyContinue
    if ($allCollectors) {
        $allCollectors | ForEach-Object {
            Log-Info "Attempting to stop otel-collector process with PID $($_.Id)"
            try {
                $_ | Stop-Process -Force
                Log-Info "Gracefully stopped otel-collector (PID $($_.Id))"
            } catch {
                Log-Info "Failed to stop otel-collector (PID $($_.Id)): $_"
                $success = $false
            }
        }
    }

    # 2. Now check the PID file and force-stop the specific process if it still exists
    if (Test-Path $PIDFILE) {
        Log-Info "PID file found at: $PIDFILE"
        $collectorPid = Get-Content $PIDFILE | Select-Object -First 1

        # First check if the PID is numeric
        if (-not [int]::TryParse($collectorPid, [ref]$null)) {
            Log-Info "Invalid PID found in file: $collectorPid - removing PID file"
            Remove-Item $PIDFILE -Force
            return $success
        }

        # Get the process safely without throwing errors
        $process = Get-Process -Id $collectorPid -ErrorAction SilentlyContinue

        if ($process) {
            Log-Info "Attempting to force-stop otel-collector process with PID $collectorPid (from PID file)"
            try {
                Stop-Process -Id $collectorPid -Force
                Log-Info "Force-stopped otel-collector (PID $collectorPid) from PID file"
            } catch {
                Log-Info "Failed to force-stop otel-collector (PID $collectorPid): $_"
                $success = $false
            }
        } else {
            Log-Info "No running otel-collector process found with PID $collectorPid"
        }
        Remove-Item $PIDFILE -Force
    } else {
        Log-Info "PID file not found, nothing to stop."
    }
    return $success
}

function Stop-OtelCollector-Service {
    Log-Info "Stopping service using: $serviceScript"
    Start-Process -FilePath "powershell.exe" -ArgumentList "-File `"$serviceScript`" stop" -NoNewWindow -Wait
    Start-Process -FilePath "powershell.exe" -ArgumentList "-File `"$serviceScript`" remove" -NoNewWindow -Wait
}

function Stop-OtelCollector {
    # Check if am is installed as process or service
    $mode = Get-InstallMode
    if ($mode -eq "process") {
        # handle as a process
        return Stop-OtelCollector-Process
    } else {
        # handle as a service
        return Stop-OtelCollector-Service
    }
}

function Restart-OtelCollector-Process {
    $stopResult = Stop-OtelCollector
    if (-not $stopResult) {
        Log-Info "Failed to stop otel-collector during restart."
        return $false
    }
    Start-Sleep -Seconds 2
    $startResult = Start-OtelCollector
    if ($startResult) {
        Log-Info "otel-collector restarted successfully."
        return $true
    } else {
        Log-Info "Failed to start otel-collector during restart."
        return $false
    }
}

function Restart-OtelCollector-Service {
    Log-Info "Restarting service using: $serviceScript"
    Start-Process -FilePath "powershell.exe" -ArgumentList "-File `"$serviceScript`" restart" -NoNewWindow -Wait
}

function Restart-OtelCollector {
    # Check if am is installed as process or service
    $mode = Get-InstallMode
    if ($mode -eq "process") {
        # handle as a process
        return Restart-OtelCollector-Process
    } else {
        # handle as a service
        return Restart-OtelCollector-Service
    }
}

switch ($Action) {
    "start"   { if (Start-OtelCollector) { exit 0 } else { exit 1 } }
    "stop"    { if (Stop-OtelCollector) { exit 0 } else { exit 1 } }
    "restart" { if (Restart-OtelCollector) { exit 0 } else { exit 1 } }
    "status"   {
        try {
            if(Is-Otel-Collector-Running) {
                Write-Host "OpenTelemetry Collector is running."
                exit 0
            } else {
                Write-Host "OpenTelemetry Collector is not running."
                exit 1
            }
        } catch {
            Write-Error "Error checking status: $_"
            exit 2
        }
    }
    default   { Write-Host "Usage: .\am.ps1 -Action [start|stop|restart|status] -m <HomeDir>" }
}