#!/usr/bin/pwsh

param (
    [ValidateSet("install", "start", "stop", "restart", "enable", "disable", "remove", "status")]
    [string]$Action = "status",
    [string]$ServiceName = "am-otel-collector"
)

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"
# Set environment variable globally (machine scope)
[System.Environment]::SetEnvironmentVariable("OTEL_COLLECTOR_DATA", $env:OTEL_COLLECTOR_DATA, [System.EnvironmentVariableTarget]::Machine)

$env:OTEL_COLLECTOR_LOGS = Join-Path -Path $env:AM_HOME -ChildPath "logs"

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

# Source the message.ps1
. (Join-Path -Path $env:AM_HOME -ChildPath "bin\tools\message.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 Service-Exists {
    return Get-Service -Name $ServiceName -ErrorAction SilentlyContinue
}

# Function: Install-OtelService
# Purpose: Installs the OpenTelemetry Collector as a Windows service, including configuration and recovery options.
function Install-OtelService {
    # Check if the otel-collector executable exists
    if (-Not (Test-Path $env:OTEL_COLLECTOR_EXE)) {
        Write-Error "Executable not found at $OTEL_COLLECTOR_EXE"
        return $false
    }
    # Check if the configuration directory exists
    if (-Not (Test-Path $env:AM_CONFIG_PIPELINES)) {
        Write-Error "otelcol config not found at $AM_CONFIG_PIPELINES"
        return $false
    }

    # Get the configuration arguments (all config file paths)
    $new_config = Get-ConfigArgs
    if ([string]::IsNullOrWhiteSpace($new_config)) {
        Write-Error "No config paths found. Cannot continue installation."
        return $false
    }

    # If the service already exists, remove it before creating a new one
    if (Get-Service -Name $ServiceName -ErrorAction SilentlyContinue) {
        Log-Info "Service '$ServiceName' already exists. Removing..."
        if (-not (Remove-OtelService)) {
            Write-Error "Failed to remove existing otel-collector service."
            return $false
        }
    }

    # Build the binPath argument for the service, including the executable and config arguments
    $binaryPath = "`"$env:OTEL_COLLECTOR_EXE`" $new_config"
    # Log-Info "Binary path: "
    # Log-Info "$binaryPath"
    # Create the service with the specified binary path, display name, and set it to start automatically
    $result = sc.exe create $ServiceName binPath= $binaryPath DisplayName= "OpenTelemetry Collector" start= auto
    if ($LASTEXITCODE -eq 0) {
        Log-Info "Service '$ServiceName' created successfully."

        # Set service recovery options to restart on failure after 60 seconds
        $failureResult = sc.exe failure $ServiceName reset= 60 actions= restart/60000
        if ($LASTEXITCODE -eq 0) {
            Log-Info "Recovery options set successfully for '$ServiceName'."
            return $true
        } else {
            Write-Warning "Service was created, but failed to set recovery options. sc.exe exit code: $LASTEXITCODE"
            return $false
        }
    } else {
        Write-Error "Failed to create service. sc.exe exit code: $LASTEXITCODE"
        return $false
    }
}

# Function: Start-OtelService
# Purpose: Starts the OpenTelemetry Collector Windows service.
function Start-OtelService {
    try {
        Start-Service -Name $ServiceName -ErrorAction Stop
        Log-Info "Service '$ServiceName' started."
        return $true
    } catch {
        Log-Info "Failed to start service: $_"
        return $false
    }
}

# Function: Stop-OtelService
# Purpose: Stops the OpenTelemetry Collector Windows service if it is running.
function Stop-OtelService {
    if (-Not (Service-Exists)) {
        Write-Warning "Service '$ServiceName' does not exist."
        return $false
    }
    try {
        Stop-Service -Name $ServiceName -Force -ErrorAction Stop
        Log-Info "Service '$ServiceName' stopped."
        return $true
    } catch {
        Write-Error "Failed to stop service: $_"
        return $false
    }
}

# Function: Restart-OtelService
# Purpose: Restarts the OpenTelemetry Collector Windows service.
function Restart-OtelService {
    if (-Not (Service-Exists)) {
        Write-Warning "Service '$ServiceName' does not exist."
        return $false
    }
    try {
        Restart-Service -Name $ServiceName -Force -ErrorAction Stop
        Log-Info "Service '$ServiceName' restarted."
        return $true
    } catch {
        Write-Error "Failed to restart service: $_"
        return $false
    }
}

# Function: Enable-OtelService
# Purpose: Sets the OpenTelemetry Collector Windows service to start automatically on boot.
function Enable-OtelService {
    if (-Not (Service-Exists)) {
        Write-Warning "Service '$ServiceName' does not exist."
        return $false
    }
    try {
        Set-Service -Name $ServiceName -StartupType Automatic -ErrorAction Stop
        Log-Info "Service '$ServiceName' enabled to start automatically."
        return $true
    } catch {
        Write-Error "Failed to enable service: $_"
        return $false
    }
}

# Function: Disable-OtelService
# Purpose: Sets the OpenTelemetry Collector Windows service to require manual start (disables auto-start).
function Disable-OtelService {
    if (-Not (Service-Exists)) {
        Write-Warning "Service '$ServiceName' does not exist."
        return $false
    }
    try {
        Set-Service -Name $ServiceName -StartupType Disabled -ErrorAction Stop
        Log-Info "Service '$ServiceName' disabled from starting automatically."
        return $true
    } catch {
        Write-Error "Failed to disable service: $_"
        return $false
    }
}

# Function: Remove-OtelService
# Purpose: Stops and deletes the OpenTelemetry Collector Windows service if it exists.
function Remove-OtelService {
    if (-Not (Service-Exists)) {
        Write-Warning "Service '$ServiceName' does not exist or already removed."
        return $false
    }

    try {
        Stop-Service -Name $ServiceName -Force -ErrorAction SilentlyContinue
        sc.exe delete $ServiceName | Out-Null
        if ($LASTEXITCODE -eq 0) {
            Log-Info "Service '$ServiceName' removed successfully."
            return $true
        } else {
            Write-Error "Failed to remove service. sc.exe exit code: $LASTEXITCODE"
            return $false
        }
    } catch {
        Write-Error "Error while removing service: $_"
        return $false
    }
}

function Get-OtelServiceStatus {
    $svc = Get-Service -Name $ServiceName -ErrorAction SilentlyContinue
    if ($null -ne $svc) {
        Write-Output "Service '$ServiceName' is $($svc.Status) with start type: $((Get-WmiObject -Class Win32_Service -Filter "Name='$ServiceName'").StartMode)"
        return $true
    } else {
        Write-Warning "Service '$ServiceName' not found."
        return $false
    }
}

# Main Dispatcher
switch ($Action) {
    "install" { if (Install-OtelService) { exit 0 } else { exit 1 } }
    "start"   {
        Install-OtelService
        Start-OtelService
    }
    "stop"    { if (Stop-OtelService)  { exit 0 } else { exit 1 } }
    "restart" {
        Install-OtelService
        Restart-OtelService
    }
    "enable"  { if (Enable-OtelService)  { exit 0 } else { exit 1 } }
    "disable" { if (Disable-OtelService) { exit 0 } else { exit 1 } }
    "remove"  { if (Remove-OtelService)  { 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-Warning "Unsupported action: $Action"; exit 1 }
}

exit 0