#!/usr/bin/pwsh

param (
    [ValidateSet("install", "start", "stop", "restart", "enable", "disable", "remove", "status")]
    [string]$Action = "status",
    [string]$ServiceName = "am-otel-collector",
    [switch]$OtelOnly,          # Only affects status display - both services always managed for operations
    [switch]$DirectoryWatcherOnly  # Only affects status display - both services always managed for operations
)

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

# Set environment variables
$env:MULE_HOME = $HomeDir
$env:AM_HOME = Join-Path -Path $HomeDir -ChildPath "am"

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

# DirectoryWatcher service settings
$DirectoryWatcherServiceName = "DirectoryWatcherService"
$DirectoryWatcherServiceScript = Join-Path -Path $env:AM_HOME -ChildPath "bin\dw-service.ps1"

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

#-----------------------------------------------------------------------------#
#                    OpenTelemetry Collector Functions                       #
#-----------------------------------------------------------------------------#

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 $env: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 $env:AM_CONFIG_PIPELINES"
        return $false
    }

    # Process configuration files with FileProcessingUtils before creating service
    Log-Info "Processing configuration files before service creation..."
    if (-not (Invoke-FileProcessingUtils -MuleHome $HomeDir -AmHome $env:AM_HOME)) {
        Write-Error "FileProcessingUtils failed - cannot create OTEL collector service"
        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"
    # 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
    }
}

#-----------------------------------------------------------------------------#
#                    DirectoryWatcher Service Functions                      #
#-----------------------------------------------------------------------------#

function Test-DirectoryWatcherServiceScript {
    if (-not (Test-Path $DirectoryWatcherServiceScript)) {
        Write-Error "DirectoryWatcher service script not found: $DirectoryWatcherServiceScript"
        Write-Error "Please ensure dw-service.ps1 is located in $env:AM_HOME\bin\"
        return $false
    }
    return $true
}

function Install-DirectoryWatcherService {
    Log-Info "Installing DirectoryWatcher service..."
    
    if (-not (Test-DirectoryWatcherServiceScript)) {
        return $false
    }
    
    try {
        Log-Info "Setting up DirectoryWatcher service with embedded paths..."
        Log-Info "  MULE_HOME: $env:MULE_HOME"
        Log-Info "  AM_HOME: $env:AM_HOME"
        Log-Info "  JAVA_HOME: $env:JAVA_HOME"
        
        # Execute the DirectoryWatcher service setup (using LocalSystem account)
        $result = & PowerShell.exe -ExecutionPolicy Bypass -File $DirectoryWatcherServiceScript -Setup -UserName "LocalSystem"
        
        if ($LASTEXITCODE -eq 0) {
            Log-Info "DirectoryWatcher service installed successfully"
            return $true
        } else {
            Write-Error "DirectoryWatcher service installation failed with exit code: $LASTEXITCODE"
            return $false
        }
    } catch {
        Write-Error "Error installing DirectoryWatcher service: $($_.Exception.Message)"
        return $false
    }
}

function Start-DirectoryWatcherService {
    if (-not (Test-DirectoryWatcherServiceScript)) {
        return $false
    }
    
    try {
        Log-Info "Starting DirectoryWatcher service..."
        $result = & PowerShell.exe -ExecutionPolicy Bypass -File $DirectoryWatcherServiceScript -Start
        
        if ($LASTEXITCODE -eq 0) {
            Log-Info "DirectoryWatcher service started successfully"
            return $true
        } else {
            Write-Error "DirectoryWatcher service start failed with exit code: $LASTEXITCODE"
            return $false
        }
    } catch {
        Write-Error "Error starting DirectoryWatcher service: $($_.Exception.Message)"
        return $false
    }
}

function Stop-DirectoryWatcherService {
    if (-not (Test-DirectoryWatcherServiceScript)) {
        return $false
    }
    
    try {
        Log-Info "Stopping DirectoryWatcher service..."
        $result = & PowerShell.exe -ExecutionPolicy Bypass -File $DirectoryWatcherServiceScript -Stop
        
        if ($LASTEXITCODE -eq 0) {
            Log-Info "DirectoryWatcher service stopped successfully"
            return $true
        } else {
            Write-Error "DirectoryWatcher service stop failed with exit code: $LASTEXITCODE"
            return $false
        }
    } catch {
        Write-Error "Error stopping DirectoryWatcher service: $($_.Exception.Message)"
        return $false
    }
}

function Restart-DirectoryWatcherService {
    if (-not (Test-DirectoryWatcherServiceScript)) {
        return $false
    }
    
    try {
        Log-Info "Restarting DirectoryWatcher service..."
        $result = & PowerShell.exe -ExecutionPolicy Bypass -File $DirectoryWatcherServiceScript -Restart
        
        if ($LASTEXITCODE -eq 0) {
            Log-Info "DirectoryWatcher service restarted successfully"
            return $true
        } else {
            Write-Error "DirectoryWatcher service restart failed with exit code: $LASTEXITCODE"
            return $false
        }
    } catch {
        Write-Error "Error restarting DirectoryWatcher service: $($_.Exception.Message)"
        return $false
    }
}

function Remove-DirectoryWatcherService {
    if (-not (Test-DirectoryWatcherServiceScript)) {
        return $false
    }
    
    try {
        Log-Info "Removing DirectoryWatcher service..."
        $result = & PowerShell.exe -ExecutionPolicy Bypass -File $DirectoryWatcherServiceScript -Remove
        
        if ($LASTEXITCODE -eq 0) {
            Log-Info "DirectoryWatcher service removed successfully"
            return $true
        } else {
            Write-Error "DirectoryWatcher service removal failed with exit code: $LASTEXITCODE"
            return $false
        }
    } catch {
        Write-Error "Error removing DirectoryWatcher service: $($_.Exception.Message)"
        return $false
    }
}

function Get-DirectoryWatcherServiceStatus {
    if (-not (Test-DirectoryWatcherServiceScript)) {
        return $false
    }
    
    try {
        $result = & PowerShell.exe -ExecutionPolicy Bypass -File $DirectoryWatcherServiceScript -Status
        return $result
    } catch {
        Write-Error "Error getting DirectoryWatcher service status: $($_.Exception.Message)"
        return $false
    }
}

#-----------------------------------------------------------------------------#
#                    Service Dependency Validation                           #
#-----------------------------------------------------------------------------#

function Test-ServiceDependencies {
    param([string]$ActionName)
    
    # For operational actions (not status), enforce that both services must exist together
    # Exception: Allow DirectoryWatcherOnly for remove action (needed for uninstall.ps1)
    if ($ActionName -in @("install", "start", "stop", "restart", "enable", "disable")) {
        if ($DirectoryWatcherOnly) {
            Write-Warning "DirectoryWatcher service cannot be managed independently."
            Write-Warning "Both OpenTelemetry Collector and DirectoryWatcher services are managed together."
            Write-Host "Use: .\am-service.ps1 -Action $ActionName (without -DirectoryWatcherOnly)" -ForegroundColor Yellow
            return $false
        }
    }
    
    return $true
}

#-----------------------------------------------------------------------------#
#                         Unified Service Management                         #
#-----------------------------------------------------------------------------#

function Invoke-ServiceAction {
    param(
        [string]$ActionName,
        [scriptblock]$OtelAction,
        [scriptblock]$DirectoryWatcherAction
    )
    
    # Check if installation mode is set (only for operational actions, not status)
    if ($ActionName -ne "status") {
        $mode = Get-InstallMode
        if ($mode -eq "not_installed") {
            Write-Error "Installation mode is not selected. Please run the install.ps1 script first."
            return $false
        }
        
        if ($mode -ne "service") {
            Write-Warning "Current installation mode is '$mode'. Service management is only available when installed in 'service' mode."
            Write-Host "To use service management, run install.ps1 and select 'service' mode." -ForegroundColor Yellow
            return $false
        }
    }
    
    # Validate service dependencies
    if (-not (Test-ServiceDependencies -ActionName $ActionName)) {
        return $false
    }
    
    $success = $true
    
    # For operational actions, always manage both services (except when OtelOnly is specified)
    # For status, allow selective display
    if ($ActionName -eq "status") {
        $manageOtel = -not $DirectoryWatcherOnly
        $manageDirectoryWatcher = -not $OtelOnly
    } else {
        # For all operational actions, manage both services unless OtelOnly is specified
        $manageOtel = $true
        $manageDirectoryWatcher = -not $OtelOnly
        
        if ($OtelOnly) {
            Write-Host "WARNING: Managing only OpenTelemetry Collector (DirectoryWatcher will not be affected)" -ForegroundColor Yellow
        }
    }
    
    if ($manageOtel) {
        Write-Host "=== OpenTelemetry Collector $ActionName ===" -ForegroundColor Yellow
        try {
            $result = & $OtelAction
            if (-not $result) {
                $success = $false
                Write-Error "OpenTelemetry Collector $ActionName failed"
            }
        } catch {
            $success = $false
            Write-Error "OpenTelemetry Collector $ActionName error: $($_.Exception.Message)"
        }
    }
    
    if ($manageDirectoryWatcher) {
        Write-Host "`n=== DirectoryWatcher Service $ActionName ===" -ForegroundColor Yellow
        try {
            $result = & $DirectoryWatcherAction
            if (-not $result) {
                $success = $false
                Write-Error "DirectoryWatcher service $ActionName failed"
            }
        } catch {
            $success = $false
            Write-Error "DirectoryWatcher service $ActionName error: $($_.Exception.Message)"
        }
    }
    
    return $success
}

#-----------------------------------------------------------------------------#
#                              Main Dispatcher                               #
#-----------------------------------------------------------------------------#

switch ($Action) {
    "install" { 
        $success = Invoke-ServiceAction -ActionName "install" -OtelAction {
            Install-OtelService
        } -DirectoryWatcherAction {
            Install-DirectoryWatcherService
        }
        
        if ($success) { 
            Write-Host "`n[SUCCESS] All services installed successfully" -ForegroundColor Green
            exit 0 
        } else { 
            Write-Host "`n[ERROR] Some service installations failed" -ForegroundColor Red
            exit 1 
        }
    }
    
    "start" {
        $success = Invoke-ServiceAction -ActionName "start" -OtelAction {
            Set-GoMemLimit      # Configure GOMEMLIMIT using shared function
            Install-OtelService  # Ensure it's installed first
            Start-OtelService
        } -DirectoryWatcherAction {
            Install-DirectoryWatcherService  # Ensure it's installed first
            Start-DirectoryWatcherService
        }
        
        if ($success) { 
            Write-Host "`n[SUCCESS] All services started successfully" -ForegroundColor Green
            exit 0 
        } else { 
            Write-Host "`n[ERROR] Some services failed to start" -ForegroundColor Red
            exit 1 
        }
    }
    
    "stop" { 
        $success = Invoke-ServiceAction -ActionName "stop" -OtelAction {
            Stop-OtelService
        } -DirectoryWatcherAction {
            Stop-DirectoryWatcherService
        }
        
        if ($success) { 
            Write-Host "`n[SUCCESS] All services stopped successfully" -ForegroundColor Green
            exit 0 
        } else { 
            Write-Host "`n[ERROR] Some services failed to stop" -ForegroundColor Red
            exit 1 
        }
    }
    
    "restart" {
        $success = Invoke-ServiceAction -ActionName "restart" -OtelAction {
            Install-OtelService  # Ensure it's installed first
            Restart-OtelService
        } -DirectoryWatcherAction {
            Install-DirectoryWatcherService  # Ensure it's installed first
            Restart-DirectoryWatcherService
        }
        
        if ($success) { 
            Write-Host "`n[SUCCESS] All services restarted successfully" -ForegroundColor Green
            exit 0 
        } else { 
            Write-Host "`n[ERROR] Some services failed to restart" -ForegroundColor Red
            exit 1 
        }
    }
    
    "enable" { 
        $success = Invoke-ServiceAction -ActionName "enable" -OtelAction {
            Enable-OtelService
        } -DirectoryWatcherAction {
            # DirectoryWatcher service is always set to automatic during installation
            $dwService = Get-Service -Name $DirectoryWatcherServiceName -ErrorAction SilentlyContinue
            if ($dwService) {
                Set-Service -Name $DirectoryWatcherServiceName -StartupType Automatic
                Log-Info "DirectoryWatcher service enabled to start automatically."
                return $true
            } else {
                Write-Warning "DirectoryWatcher service not found."
                return $false
            }
        }
        
        if ($success) { 
            Write-Host "`n[SUCCESS] All services enabled successfully" -ForegroundColor Green
            exit 0 
        } else { 
            Write-Host "`n[ERROR] Some services failed to enable" -ForegroundColor Red
            exit 1 
        }
    }
    
    "disable" { 
        $success = Invoke-ServiceAction -ActionName "disable" -OtelAction {
            Disable-OtelService
        } -DirectoryWatcherAction {
            $dwService = Get-Service -Name $DirectoryWatcherServiceName -ErrorAction SilentlyContinue
            if ($dwService) {
                Set-Service -Name $DirectoryWatcherServiceName -StartupType Disabled
                Log-Info "DirectoryWatcher service disabled from starting automatically."
                return $true
            } else {
                Write-Warning "DirectoryWatcher service not found."
                return $false
            }
        }
        
        if ($success) { 
            Write-Host "`n[SUCCESS] All services disabled successfully" -ForegroundColor Green
            exit 0 
        } else { 
            Write-Host "`n[ERROR] Some services failed to disable" -ForegroundColor Red
            exit 1 
        }
    }
    
    "remove" { 
        $success = Invoke-ServiceAction -ActionName "remove" -OtelAction {
            Remove-OtelService
        } -DirectoryWatcherAction {
            Remove-DirectoryWatcherService
        }
        
        if ($success) { 
            Write-Host "`n[SUCCESS] All services removed successfully" -ForegroundColor Green
            exit 0 
        } else { 
            Write-Host "`n[ERROR] Some services failed to remove" -ForegroundColor Red
            exit 1 
        }
    }
    
    "status" {
        try {
            $manageOtel = -not $DirectoryWatcherOnly
            $manageDirectoryWatcher = -not $OtelOnly
            
            if ($manageOtel) {
                Write-Host "=== OpenTelemetry Collector Status ===" -ForegroundColor Yellow
                if(Is-Otel-Collector-Running) {
                    Write-Host "[RUNNING] OpenTelemetry Collector is running." -ForegroundColor Green
                } else {
                    Write-Host "[STOPPED] OpenTelemetry Collector is not running." -ForegroundColor Red
                }
                Get-OtelServiceStatus | Out-Null
            }
            
            if ($manageDirectoryWatcher) {
                # Check install mode for appropriate header
                $mode = Get-InstallMode
                if ($mode -eq "service") {
                    Write-Host "`n=== DirectoryWatcher Service Status ===" -ForegroundColor Yellow
                } else {
                    Write-Host "`n=== DirectoryWatcher Process Status ===" -ForegroundColor Yellow
                }
                $dwStatus = Get-DirectoryWatcherServiceStatus
                
                # Check if DirectoryWatcher Java process is running
                $javaProcess = Get-Process java -ErrorAction SilentlyContinue | Where-Object {
                    (Get-WmiObject Win32_Process -Filter "ProcessId=$($_.Id)" -ErrorAction SilentlyContinue).CommandLine -match "DirectoryWatcher"
                }
                
                if ($javaProcess) {
                    Write-Host "[RUNNING] DirectoryWatcher Java process is running (PID: $($javaProcess.Id))" -ForegroundColor Green
                } else {
                    Write-Host "[STOPPED] DirectoryWatcher Java process is not running" -ForegroundColor Red
                }
            }
            
            exit 0
        } catch {
            Write-Error "Error checking status: $_"
            exit 2
        }
    }
    
    default { 
        Write-Warning "Unsupported action: $Action"
        Write-Host "`nUsage: .\am-service.ps1 -Action <action> [options]"
        Write-Host "Actions: install, start, stop, restart, enable, disable, remove, status"
        Write-Host "Options:"
        Write-Host "  -OtelOnly              Manage only OpenTelemetry Collector (operational actions)"
        Write-Host "                        Display only OpenTelemetry Collector status (status action)"
        Write-Host "  -DirectoryWatcherOnly  Display only DirectoryWatcher status (status action only)"
        Write-Host "                        WARNING: DirectoryWatcher cannot be managed independently for operational actions"
        Write-Host "`nExamples:"
        Write-Host "  .\am-service.ps1 -Action install                    # Install both services"
        Write-Host "  .\am-service.ps1 -Action start -OtelOnly            # Start only OTEL Collector"
        Write-Host "  .\am-service.ps1 -Action status -DirectoryWatcherOnly  # Show only DirectoryWatcher status"
        Write-Host "  .\am-service.ps1 -Action stop                       # Stop both services (required)"
        exit 1 
    }
}

exit 0