#!/usr/bin/pwsh

param(
    [Alias('h')][switch]$ShowHelp,
    [Alias('m')][string]$HomeDir=""
)

if ($ShowHelp) {
    Write-Host "DESCRIPTION: Uninstall Anypoint Monitoring OpenTelemetry Collector and DirectoryWatcher Service"
    Write-Host "Usage: .\$($MyInvocation.MyCommand) [options]"
    Write-Host " "
    Write-Host "Options:"
    Write-Host "-h,                    Show help"
    Write-Host "-m [DIR]               Set home directory - defaults to MULE_HOME"
    Write-Host " "
    Write-Host "Examples:"
    Write-Host ".\uninstall.ps1                              # Uninstall both services"
    Write-Host ".\uninstall.ps1 -m C:\mule-enterprise         # Uninstall with custom MULE_HOME"
    Write-Host " "
    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 {
        $scriptDir = Split-Path -Parent $MyInvocation.MyCommand.Definition
        $HomeDir = (Resolve-Path -Path (Join-Path -Path $scriptDir -ChildPath "..\..")).Path
        Write-Host "Determined the MULE_HOME"
    }
}

# 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

# Java program running on Windows treats paths differently when the trailing character is \
# If MULE_HOME env is not set, HomeDir is autodetected and things are merry.
# But if HomeDir is supplied as an option, with a valid path but ending in \, then AgentConfiguratorApp complains:
# ERROR: The provided Mule Home doesn't exist, please check if the provided directory is correct.
# Hence the TrimEnd
$HomeDir = $HomeDir.TrimEnd('\','/')
$env:MULE_HOME = $HomeDir
$env:AM_HOME = Join-Path -Path $HomeDir -ChildPath "am"

Log-Info "MULE_HOME is set to $HomeDir"
Log-Info "AM_HOME is set to $env:AM_HOME"

# Service configuration
$DirectoryWatcherServiceName = "DirectoryWatcherService"
$DirectoryWatcherServiceScript = Join-Path -Path $env:AM_HOME -ChildPath "bin\dw-service.ps1"
$amScript = Join-Path -Path $env:AM_HOME -ChildPath "bin\am.ps1"

#-----------------------------------------------------------------------------#
#                              Utility Functions                             #
#-----------------------------------------------------------------------------#

# Function to check if running as administrator
function Test-Administrator {
    $currentUser = New-Object Security.Principal.WindowsPrincipal([Security.Principal.WindowsIdentity]::GetCurrent())
    return $currentUser.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

# Function to unset GOMEMLIMIT environment variable
# Uninstall should clean up this universal GOMEMLIMIT env variable
function Unset-GoMemLimit {
    try {
        Log-Info "Checking for GOMEMLIMIT environment variable set at Machine level..."
        $currentValue = [Environment]::GetEnvironmentVariable("GOMEMLIMIT", "Machine")

        if ($null -ne $currentValue) {
            Log-Info "GOMEMLIMIT is currently set to: $currentValue"
            Log-Info "Removing GOMEMLIMIT environment variable..."
            [Environment]::SetEnvironmentVariable("GOMEMLIMIT", $null, "Machine")

            # Verify removal
            $newValue = [Environment]::GetEnvironmentVariable("GOMEMLIMIT", "Machine")
            if ($null -eq $newValue) {
                Log-Info "GOMEMLIMIT environment variable has been successfully removed."
                return $true
            } else {
                Log-Info "Warning: Failed to remove GOMEMLIMIT. Current value: $newValue"
                return $false
            }
        } else {
            Log-Info "GOMEMLIMIT environment variable is not set. No action needed."
            return $true
        }
    } catch {
        Log-Info "Error removing GOMEMLIMIT environment variable: $($_.Exception.Message)"
        return $false
    }
}

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

function Stop-DirectoryWatcherJavaProcess {
    try {
        Log-Info "Attempting to stop DirectoryWatcher Java processes..."
        $targetClass = "com.mulesoft.dias.mule.agent.DirectoryWatcher"
        
        # Get all java.exe processes with full command lines
        $javaProcesses = Get-CimInstance Win32_Process | Where-Object {
            $_.Name -eq "java.exe" -and $_.CommandLine -match [regex]::Escape($targetClass)
        }

        $className = $targetClass.Split('.')[-1]

        if ($javaProcesses.Count -eq 0) {
            Log-Info "No running Java process found for class $className."
            return $true
        }

        foreach ($proc in $javaProcesses) {
            Log-Info "Stopping Java process with PID $($proc.ProcessId) running class $className..."
            try {
                Stop-Process -Id $proc.ProcessId -Force -ErrorAction Stop
                Log-Info "Successfully stopped Java process PID $($proc.ProcessId)"
            } catch {
                Log-Info "Warning: Failed to stop Java process PID $($proc.ProcessId): $($_.Exception.Message)"
            }
        }

        # Wait a moment for processes to terminate
        Start-Sleep 2

        # Verify all processes are stopped
        $remainingProcesses = Get-CimInstance Win32_Process | Where-Object {
            $_.Name -eq "java.exe" -and $_.CommandLine -match [regex]::Escape($targetClass)
        }

        if ($remainingProcesses.Count -eq 0) {
            Log-Info "All DirectoryWatcher Java processes stopped successfully."
            return $true
        } else {
            Log-Info "Warning: $($remainingProcesses.Count) DirectoryWatcher Java processes are still running."
            return $false
        }
    } catch {
        Log-Info "Error stopping DirectoryWatcher Java processes: $($_.Exception.Message)"
        return $false
    }
}

function Uninstall-DirectoryWatcherService {
    Log-Info "=== DirectoryWatcher Service Uninstallation ===" 
    $success = $true

    try {
        # Check install.mode to determine how to uninstall
        $installModeFile = Join-Path -Path $env:AM_HOME -ChildPath "install.mode"
        $installMode = "process"  # Default to process if file not found
        
        if (Test-Path $installModeFile) {
            $installMode = (Get-Content $installModeFile -Raw).Trim()
            Log-Info "Found install.mode: $installMode"
        } else {
            Log-Info "install.mode file not found. Assuming process mode for uninstall."
        }
        
        if ($installMode -eq "service") {
            Log-Info "Service mode detected - service management will be handled centrally."
        } else {
            Log-Info "DirectoryWatcher was installed as process mode - no service to remove."
            Log-Info "Will only stop DirectoryWatcher Java processes and clean up files."
        }

        # Stop any remaining DirectoryWatcher Java processes
        if (-not (Stop-DirectoryWatcherJavaProcess)) {
            $success = $false
        }

        # Clean up DirectoryWatcher-related files
        Log-Info "Cleaning up DirectoryWatcher-related files..."
        try {
            # Remove PID file
            $pidFile = Join-Path -Path $env:AM_HOME -ChildPath "directorywatcher.pid"
            if (Test-Path $pidFile) {
                Remove-Item -Path $pidFile -Force -ErrorAction SilentlyContinue
                Log-Info "Removed DirectoryWatcher PID file."
            }

        } catch {
            Log-Info "Warning: Error during DirectoryWatcher file cleanup: $($_.Exception.Message)"
            $success = $false
        }

        if ($success) {
            Log-Info "DirectoryWatcher service uninstallation completed successfully."
        } else {
            Log-Info "DirectoryWatcher service uninstallation completed with warnings."
        }

        return $success

    } catch {
        Log-Info "Error during DirectoryWatcher service uninstallation: $($_.Exception.Message)"
        return $false
    }
}

#-----------------------------------------------------------------------------#
#                      Centralized Service Management                        #
#-----------------------------------------------------------------------------#

function Invoke-ServiceUninstallation {
    Log-Info "=== Centralized Service Management Uninstallation ===" 
    $success = $true
    
    # Use am-service.ps1 as single point of truth for all service operations
    $amServiceScript = Join-Path -Path $env:AM_HOME -ChildPath "bin\am-service.ps1"
    
    if (Test-Path $amServiceScript) {
        # Stop both services via am-service.ps1
        try {
            Log-Info "Stopping both services via am-service.ps1..."
            Start-Process -FilePath "powershell.exe" -ArgumentList "-ExecutionPolicy Bypass -File `"$amServiceScript`" -Action stop" -NoNewWindow -Wait
            
            if ($LASTEXITCODE -eq 0) {
                Log-Info "Both services stopped successfully via am-service.ps1."
            } else {
                Log-Info "Warning: Service stop via am-service.ps1 returned exit code: $LASTEXITCODE"
            }
        } catch {
            Log-Info "Warning: Error stopping services via am-service.ps1: $($_.Exception.Message)"
            $success = $false
        }

        # Remove both services via am-service.ps1
        try {
            Log-Info "Removing both services via am-service.ps1..."
            Start-Process -FilePath "powershell.exe" -ArgumentList "-ExecutionPolicy Bypass -File `"$amServiceScript`" -Action remove" -NoNewWindow -Wait
            
            if ($LASTEXITCODE -eq 0) {
                Log-Info "Both services removed successfully via am-service.ps1."
            } else {
                Log-Info "Warning: Service removal via am-service.ps1 returned exit code: $LASTEXITCODE"
                $success = $false
            }
        } catch {
            Log-Info "Warning: Error removing services via am-service.ps1: $($_.Exception.Message)"
            $success = $false
        }
    } else {
        Log-Info "Error: am-service.ps1 not found at: $amServiceScript"
        Log-Info "Cannot proceed with service uninstallation."
        return $false
    }
    
    return $success
}

#-----------------------------------------------------------------------------#
#                      OTEL Collector Uninstall Functions                   #
#-----------------------------------------------------------------------------#

function Uninstall-OtelCollector {
    Log-Info "=== OpenTelemetry Collector Uninstallation ==="
    $success = $true

    try {
        # Check install.mode to determine if we need to handle process stopping
        $installModeFile = Join-Path -Path $env:AM_HOME -ChildPath "install.mode"
        $installMode = "process"  # Default to process if file not found
        
        if (Test-Path $installModeFile) {
            $installMode = (Get-Content $installModeFile -Raw).Trim()
        }
        
        if ($installMode -eq "process" -and (Test-Path $amScript)) {
            # Only call am.ps1 stop for process mode (services are handled centrally)
            Log-Info "Stopping OpenTelemetry Collector process via am.ps1..."
            Start-Process -FilePath "powershell.exe" -ArgumentList "-File `"$amScript`" stop" -NoNewWindow -Wait
            Log-Info "OpenTelemetry Collector process stopped."
        } else {
            Log-Info "OpenTelemetry Collector service stopping handled by centralized service management."
        }
        
        # Remove directories and files
        Log-Info "Removing configuration files and directories..."
        Remove-Item -Path (Join-Path -Path $env:AM_HOME -ChildPath "config\prospectors") -Recurse -Force -ErrorAction SilentlyContinue
        Remove-Item -Path (Join-Path -Path $env:AM_HOME -ChildPath "config\certs") -Recurse -Force -ErrorAction SilentlyContinue
        Remove-Item -Path (Join-Path -Path $env:AM_HOME -ChildPath "config\proxy.url") -Force -ErrorAction SilentlyContinue
        Remove-Item -Path (Join-Path -Path $env:AM_HOME -ChildPath "install.mode") -Force -ErrorAction SilentlyContinue

        # Remove specific JAR files
        $jarPaths = @(
            "$HomeDir\server-plugins\mule-agent-plugin\lib\mule-agent-dias-*.jar",
            "$HomeDir\server-plugins\mule-agent-plugin\lib\analytics-metrics-collector-*.jar",
            "$HomeDir\server-plugins\mule-agent-plugin\lib\modules\mule-agent-dias-*.jar",
            "$HomeDir\server-plugins\mule-agent-plugin\lib\modules\analytics-metrics-collector-*.jar",
            "$HomeDir\plugins\mule-agent-plugin\lib\modules\mule-agent-dias-*.jar",
            "$HomeDir\plugins\mule-agent-plugin\lib\modules\analytics-metrics-collector-*.jar"
        )

        foreach ($jarPath in $jarPaths) {
            Remove-Item -Path $jarPath -Force -ErrorAction SilentlyContinue
        }

        Log-Info "OpenTelemetry Collector uninstallation completed successfully."
        return $true

    } catch {
        Log-Info "Error during OpenTelemetry Collector uninstallation: $($_.Exception.Message)"
        return $false
    }
}

#-----------------------------------------------------------------------------#
#                              Main Uninstall Logic                          #
#-----------------------------------------------------------------------------#

function Invoke-UninstallProcess {
    Log-Info "=== Starting Anypoint Monitoring Uninstallation Process ==="
    
    # Check install.mode to determine correct terminology
    $installModeFile = Join-Path -Path $env:AM_HOME -ChildPath "install.mode"
    $installMode = "process"  # Default to process if file not found
    
    if (Test-Path $installModeFile) {
        $installMode = (Get-Content $installModeFile -Raw).Trim()
    }
    
    if ($installMode -eq "service") {
        Log-Info "Uninstalling both OpenTelemetry Collector and DirectoryWatcher services"
    } else {
        Log-Info "Stopping both OpenTelemetry Collector and DirectoryWatcher processes"
    }
    
    $overallSuccess = $true
    
    # Check install.mode to determine uninstall approach
    $installModeFile = Join-Path -Path $env:AM_HOME -ChildPath "install.mode"
    $installMode = "process"  # Default to process if file not found
    
    if (Test-Path $installModeFile) {
        $installMode = (Get-Content $installModeFile -Raw).Trim()
    }
    
    if ($installMode -eq "service") {
        # Use centralized service management for service mode
        if (-not (Invoke-ServiceUninstallation)) {
            $overallSuccess = $false
        }
        
        # Clean up DirectoryWatcher processes and files
        if (-not (Stop-DirectoryWatcherJavaProcess)) {
            $overallSuccess = $false
        }
        
        # Clean up DirectoryWatcher-related files
        Log-Info "Cleaning up DirectoryWatcher-related files..."
        try {
            $pidFile = Join-Path -Path $env:AM_HOME -ChildPath "directorywatcher.pid"
            if (Test-Path $pidFile) {
                Remove-Item -Path $pidFile -Force -ErrorAction SilentlyContinue
                Log-Info "Removed DirectoryWatcher PID file."
            }
        } catch {
            Log-Info "Warning: Error during DirectoryWatcher file cleanup: $($_.Exception.Message)"
            $overallSuccess = $false
        }
    } else {
        # Process mode: handle each component separately
        if (-not (Uninstall-DirectoryWatcherService)) {
            $overallSuccess = $false
        }
    }
    
    # Handle OTEL Collector cleanup (file removal, etc.)
    if (-not (Uninstall-OtelCollector)) {
        $overallSuccess = $false
    }
    
    # Handle GOMEMLIMIT cleanup
    if (Test-Administrator) {
        if (-not (Unset-GoMemLimit)) {
            $overallSuccess = $false
        }
    } else {
        $currentValue = [Environment]::GetEnvironmentVariable("GOMEMLIMIT", "Machine")
        if ($null -ne $currentValue) {
            Log-Info "Failed to unset GOMEMLIMIT ($currentValue). You must run as administrator to remove system environment variables."
            Log-Info "Please restart with administrator privileges to complete cleanup."
            $overallSuccess = $false
        }
    }
    
    # Final status
    if ($overallSuccess) {
        Log-Info "=== Anypoint Monitoring uninstallation completed successfully ==="
        if ($installMode -eq "service") {
            Log-Info "Both OpenTelemetry Collector and DirectoryWatcher services have been removed."
        } else {
            Log-Info "Both OpenTelemetry Collector and DirectoryWatcher processes have been stopped and cleaned up."
        }
        return $true
    } else {
        Log-Info "=== Anypoint Monitoring uninstallation completed with errors ==="
        Log-Info "Please check the logs above for specific issues that occurred during uninstallation."
        return $false
    }
}

#-----------------------------------------------------------------------------#
#                              Execute Uninstallation                        #
#-----------------------------------------------------------------------------#

try {
    $result = Invoke-UninstallProcess
    
    if ($result) {
        Write-Host "`n[SUCCESS] Uninstallation completed successfully!" -ForegroundColor Green
        exit 0
    } else {
        Write-Host "`n[WARNING] Uninstallation completed with warnings. Check logs for details." -ForegroundColor Yellow
        exit 1
    }
} catch {
    Log-Info "Fatal error during uninstallation: $($_.Exception.Message)"
    Write-Host "`n[ERROR] Uninstallation failed with errors!" -ForegroundColor Red
    exit 1
}