#!/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,
    
    [Parameter()]
    [switch]$OtelOnly
)

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"
$DW_PIDFILE = Join-Path $env:AM_HOME "directorywatcher.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"

#-----------------------------------------------------------------------------#
#                    DirectoryWatcher Process Management                      #
#-----------------------------------------------------------------------------#

# Removed Test-DirectoryWatcherRunning wrapper - using Is-DirectoryWatcher-Running directly

function Start-DirectoryWatcherProcess {
    # Check multiple times to avoid race conditions
    for ($i = 1; $i -le 3; $i++) {
        if (Is-DirectoryWatcher-Running) {
            Log-Info "DirectoryWatcher is already running (check $i/3)"
            return $true
        }
        if ($i -lt 3) {
            Start-Sleep -Milliseconds 500  # Brief pause between checks
        }
    }

    # Check if Java is available using serviceHelper function
    if (-not (Test-JavaAvailable)) {
        Log-Info "ERROR: Java is not installed or not in PATH"
        return $false
    }

    # Construct the classpath
    $classpath = Join-Path -Path $env:AM_HOME -ChildPath "lib\*"
    
    try {
        Log-Info "Starting DirectoryWatcher process..."
        $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

        # Write PID to file
        $dwProcess.Id | Out-File -FilePath $DW_PIDFILE -Force
        Start-Sleep -Seconds 2

        if (Get-Process -Id $dwProcess.Id -ErrorAction SilentlyContinue) {
            Log-Info "DirectoryWatcher started with PID $($dwProcess.Id)"
            return $true
        } else {
            Log-Info "Failed to start DirectoryWatcher"
            return $false
        }
    } catch {
        Log-Info "ERROR: Exception starting DirectoryWatcher: $($_.Exception.Message)"
        return $false
    }
}

function Stop-DirectoryWatcherProcess {
    $success = $true
    
    # Stop all DirectoryWatcher Java processes by name
    $javaProcesses = Get-Process java -ErrorAction SilentlyContinue | Where-Object {
        (Get-CimInstance Win32_Process -Filter "ProcessId=$($_.Id)" -ErrorAction SilentlyContinue).CommandLine -match "com\.mulesoft\.dias\.mule\.agent\.DirectoryWatcher"
    }
    
    if ($javaProcesses) {
        $javaProcesses | ForEach-Object {
            Log-Info "Attempting to stop DirectoryWatcher process with PID $($_.Id)"
            try {
                $_ | Stop-Process -Force
                Log-Info "Stopped DirectoryWatcher (PID $($_.Id))"
            } catch {
                Log-Info "Failed to stop DirectoryWatcher (PID $($_.Id)): $_"
                $success = $false
            }
        }
    }

    # Check PID file and clean up
    if (Test-Path $DW_PIDFILE) {
        $dwPid = Get-Content $DW_PIDFILE | Select-Object -First 1
        
        if ([int]::TryParse($dwPid, [ref]$null)) {
            $process = Get-Process -Id $dwPid -ErrorAction SilentlyContinue
            if ($process) {
                Log-Info "Stopping DirectoryWatcher process from PID file: $dwPid"
                try {
                    Stop-Process -Id $dwPid -Force
                } catch {
                    Log-Info "Failed to stop DirectoryWatcher from PID file: $_"
                    $success = $false
                }
            }
        }
        Remove-Item $DW_PIDFILE -Force
    }
    
    return $success
}

# DirectoryWatcher service functions removed - service mode is handled by delegation to am-service.ps1

# Individual DirectoryWatcher management functions removed - 
# Both components are now managed together via Start-BothComponents/Stop-BothComponents/Restart-BothComponents

#-----------------------------------------------------------------------------#
#                    OTEL Collector Process Management                       #
#-----------------------------------------------------------------------------#

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
    }

    # Process configuration files with FileProcessingUtils before starting OTEL collector
    Log-Info "Processing configuration files before starting OTEL collector..."
    if (-not (Invoke-FileProcessingUtils -MuleHome $HomeDir -AmHome $env:AM_HOME)) {
        Log-Info "ERROR: FileProcessingUtils failed - cannot start OTEL collector"
        return $false
    }

    # Configure GOMEMLIMIT using shared function
    Set-GoMemLimit

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

# OTEL service functions removed - service mode is handled by delegation to am-service.ps1

function Start-BothComponents {
    # Check if am is installed as process or service
    $mode = Get-InstallMode
    if ($mode -eq "not_installed") {
        Write-Error "Installation mode is not selected. Please run the install.ps1 script first."
        exit 1
    }
    
    if ($mode -eq "process") {
        if ($OtelOnly) {
            # Start only OTEL Collector in process mode
            Log-Info "Starting only OTEL Collector in process mode (DirectoryWatcher will NOT be started)..."
            $otelResult = Start-OtelCollector-Process
            return $otelResult
        } else {
            # Start both OTEL Collector and DirectoryWatcher in process mode
            Log-Info "Starting both OTEL Collector and DirectoryWatcher in process mode..."
            $otelResult = Start-OtelCollector-Process
            $dwResult = Start-DirectoryWatcherProcess
            return ($otelResult -and $dwResult)
        }
    } else {
        # Delegate to am-service.ps1 for service mode
        if ($OtelOnly) {
            Log-Info "Starting only OTEL service in service mode (DirectoryWatcher will NOT be started)......"
            Start-Process -FilePath "powershell.exe" -ArgumentList "-File `"$serviceScript`" start -OtelOnly" -NoNewWindow -Wait
        } else {
            Log-Info "Starting both services via am-service.ps1..."
            Start-Process -FilePath "powershell.exe" -ArgumentList "-File `"$serviceScript`" start" -NoNewWindow -Wait
        }
        return $true
    }
}



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
}

# Stop-OtelCollector-Service removed - service mode is handled by delegation to am-service.ps1

function Stop-BothComponents {
    # Check if am is installed as process or service
    $mode = Get-InstallMode
    if ($mode -eq "not_installed") {
        Write-Error "Installation mode is not selected. Please run the install.ps1 script first."
        exit 1
    }
    
    if ($mode -eq "process") {
        if ($OtelOnly) {
            # Stop only OTEL Collector in process mode
            Log-Info "Stopping only OTEL Collector in process mode..."
            $otelResult = Stop-OtelCollector-Process
            return $otelResult
        } else {
            # Stop both OTEL Collector and DirectoryWatcher in process mode
            Log-Info "Stopping both OTEL Collector and DirectoryWatcher in process mode..."
            $otelResult = Stop-OtelCollector-Process
            $dwResult = Stop-DirectoryWatcherProcess
            return ($otelResult -and $dwResult)
        }
    } else {
        # Delegate to am-service.ps1 for service mode
        if ($OtelOnly) {
            Log-Info "Stopping only OTEL service via am-service.ps1..."
            Start-Process -FilePath "powershell.exe" -ArgumentList "-File `"$serviceScript`" stop -OtelOnly" -NoNewWindow -Wait
        } else {
            Log-Info "Stopping both services via am-service.ps1..."
            Start-Process -FilePath "powershell.exe" -ArgumentList "-File `"$serviceScript`" stop" -NoNewWindow -Wait
        }
        return $true
    }
}

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

# Restart-OtelCollector-Service removed - service mode is handled by delegation to am-service.ps1

function Restart-BothComponents {
    # Check if am is installed as process or service
    $mode = Get-InstallMode
    if ($mode -eq "not_installed") {
        Write-Error "Installation mode is not selected. Please run the install.ps1 script first."
        exit 1
    }
    
    if ($mode -eq "process") {
        if ($OtelOnly) {
            # Restart only OTEL Collector in process mode
            Log-Info "Restarting only OTEL Collector in process mode..."
            $stopResult = Stop-OtelCollector-Process
            if (-not $stopResult) {
                Log-Info "Failed to stop OTEL Collector during restart."
                return $false
            }
            Start-Sleep -Seconds 2
            $startResult = Start-OtelCollector-Process
            if ($startResult) {
                Log-Info "OTEL Collector restarted successfully."
                return $true
            } else {
                Log-Info "Failed to start OTEL Collector during restart."
                return $false
            }
        } else {
            # Restart both OTEL Collector and DirectoryWatcher in process mode
            Log-Info "Restarting both OTEL Collector and DirectoryWatcher in process mode..."
            $stopResult = Stop-BothComponents
            if (-not $stopResult) {
                Log-Info "Failed to stop components during restart."
                return $false
            }
            Start-Sleep -Seconds 2
            $startResult = Start-BothComponents
            if ($startResult) {
                Log-Info "Both components restarted successfully."
                return $true
            } else {
                Log-Info "Failed to start components during restart."
                return $false
            }
        }
    } else {
        # Delegate to am-service.ps1 for service mode
        if ($OtelOnly) {
            Log-Info "Restarting only OTEL service via am-service.ps1..."
            Start-Process -FilePath "powershell.exe" -ArgumentList "-File `"$serviceScript`" restart -OtelOnly" -NoNewWindow -Wait
        } else {
            Log-Info "Restarting both services via am-service.ps1..."
            Start-Process -FilePath "powershell.exe" -ArgumentList "-File `"$serviceScript`" restart" -NoNewWindow -Wait
        }
        return $true
    }
}

switch ($Action) {
    "start"   { 
        if ($OtelOnly) {
            $script:OtelOnly = $true
        }
        if (Start-BothComponents) { exit 0 } else { exit 1 } 
    }
    "stop"    { 
        if ($OtelOnly) {
            $script:OtelOnly = $true
        }
        if (Stop-BothComponents) { exit 0 } else { exit 1 } 
    }
    "restart" { 
        if ($OtelOnly) {
            $script:OtelOnly = $true
        }
        if (Restart-BothComponents) { exit 0 } else { exit 1 } 
    }
    "status"   {
        try {
            $otelRunning = Is-Otel-Collector-Running
            $dwRunning = Is-DirectoryWatcher-Running
            
            if ($otelRunning) {
                Write-Host "OpenTelemetry Collector is running."
            } else {
                Write-Host "OpenTelemetry Collector is not running."
            }
            
            if ($dwRunning) {
                # Get DirectoryWatcher Java process ID
                $javaProcess = Get-Process java -ErrorAction SilentlyContinue | Where-Object {
                    (Get-WmiObject Win32_Process -Filter "ProcessId=$($_.Id)" -ErrorAction SilentlyContinue).CommandLine -match "DirectoryWatcher"
                }
                
                if ($javaProcess) {
                    Write-Host "DirectoryWatcher is running (PID: $($javaProcess.Id))."
                } else {
                    Write-Host "DirectoryWatcher is running."
                }
            } else {
                Write-Host "DirectoryWatcher is not running."
            }
            
            if ($otelRunning -and $dwRunning) {
                exit 0
            } else {
                exit 1
            }
        } catch {
            Write-Error "Error checking status: $_"
            exit 2
        }
    }
    default   { Write-Host "Usage: .\am.ps1 -Action [start|stop|restart|status] -m <HomeDir>" }
}