# === Normalize input and defaults ===
$version="v7.2.8"
$uninstall = $uninstall.ToLowerInvariant()
$debug = $debug -ne "false"
$osArch = (Get-CimInstance Win32_OperatingSystem).OSArchitecture
$win_arch = if ($osArch -eq "32-bit") { 32 } else { 64 }
$isArm = $osArch -like "*ARM*"
$installer_filename = "x"
$installer_path = Join-Path $Env:TEMP "FieldEffect_MDR"
$working_folder = $installer_path
if (-not (Test-Path $working_folder)) {
    New-Item -Path $working_folder -ItemType "Directory" | Out-Null
}

$debug_log = Join-Path $working_folder "mdrAgentPoShInstaller.proc_$env:COMPUTERNAME_$(Get-Date -Format 'yyyyMMddTHHmmZ').log"
function Log {
    param ([string]$msg, [string]$level = "INFO")
    $ts = Get-Date -Format "yyyy/MM/dd HH:mm:ss"
    $entry = "$ts [$level] $msg"
    if ($debug) { Write-Host $entry }
    $entry | Out-File -FilePath $debug_log -Append -Encoding UTF8
}

Log "Script started on $($env:COMPUTERNAME)"

if ($uninstall -eq "true") {    
    Log "Uninstall mode"

    # UNINSTALL === Check for agent in registry
    $agentMSI = Get-ItemProperty HKLM:\Software\Microsoft\Windows\CurrentVersion\Uninstall\* |
                Where-Object { $_.DisplayName -eq 'Field Effect MDR' }
    if (-not $agentMSI) {
        Log "Agent not found - nothing to uninstall"
        return
    }

    # UNINSTALL === Read org ID and agent state from status file
    $statusFilePath = "C:\ProgramData\Field Effect\Covalence\data\status.json"
    if (-not (Test-Path $statusFilePath)) {
        throw "Status file not found: $statusFilePath"
    }
    try {
        $_org_id   = (Get-Content $statusFilePath | ConvertFrom-Json).AgentInfo.CurrentOrgId
        $agent_state = (Get-Content $statusFilePath | ConvertFrom-Json).AgentInfo.StateName
        Log "Org ID: $_org_id, Agent state: $agent_state"
    } catch {
        throw "Failed to read status file: $_"
    }

    # UNINSTALL === Disable agent protection via API (if api_token is available)
    if (-not [string]::IsNullOrEmpty($api_token)) {
        $headers = @{ Authorization = "Bearer $api_token"; ContentType = "application/json" }
        $hostname = $env:COMPUTERNAME

        try {
            $host_list = Invoke-RestMethod -Method Get -Uri "https://services.fieldeffect.net/v1/endpoint_devices?organization_id=$_org_id&hostname=$hostname" -Headers $headers
        } catch {
            Log "API call failed to get endpoint device data for $hostname" -level "ERROR"
        }

        if (-not $host_list.items -or $host_list.items.Count -eq 0) {
            if ($agent_state -eq "ACTIVE") {
                Log "Endpoint not found in portal but local status is ACTIVE - uninstall may fail if protection is enabled" -level "ERROR"
            } else {
                Log "Endpoint not found in portal - will attempt uninstall without modifying protection"
            }
        } else {
            $endpoint_id = $host_list.items[0].id
            Log "Found endpoint ID: $endpoint_id"

            $bulk_update_body = @{
                agent_protection = "Disabled"
                endpoint_id      = $endpoint_id
            } | ConvertTo-Json

            try {
                $protect_status = Invoke-RestMethod -Method Get -Uri "https://services.fieldeffect.net/v1/endpoint_feature_overrides?endpoint_id=$endpoint_id" -Headers $headers
            } catch {
                throw "API call failed to get endpoint feature overrides"
            }

            try {
                if ($protect_status.items -and $protect_status.items.Count -gt 0 -and $protect_status.items[0].id) {
                    $override_id = $protect_status.items[0].id
                    Invoke-RestMethod -Method Patch -Uri "https://services.fieldeffect.net/v1/endpoint_feature_overrides/$override_id" -Headers $headers -Body $bulk_update_body -ContentType "application/json"
                } else {
                    Invoke-RestMethod -Method Post -Uri "https://services.fieldeffect.net/v1/endpoint_feature_overrides" -Headers $headers -Body $bulk_update_body -ContentType "application/json"
                }
                Log "Agent protection disabled for $hostname"
            } catch {
                Log "Failed to disable agent protection: $_" -level "ERROR"
                throw "API call failed to update endpoint protection"
            }
        }
    } else {
        Log "No API token provided - skipping protection disable"
    }

    # UNINSTALL === Stop services
    $services = @("Covalence Endpoint Service", "CovAgent")

    foreach ($service in $services) {
        for ($attempt = 0; $attempt -le 4; $attempt++) {
            try {
                $serviceObj = Get-Service -Name $service -ErrorAction SilentlyContinue
            
                if (-not $serviceObj) {
                    Log "Service $service not found"
                    break
                }
            
                if ($serviceObj.Status -ne 'Running') {
                    Log "Service $service is not running"
                    break
                }
            
                Log "Attempting to stop service: $service (Attempt $($attempt + 1))..."
                Stop-Service -Name $service -Force -ErrorAction Stop
                Log "Service $service stopped successfully"
                break
            }
            catch {
                Log "Failed to stop $service on attempt $($attempt + 1): $_"
                if ($attempt -eq 4) {
                    throw "Failed to stop $service after 5 attempts: $_"
                }
                Log "Retrying in 5 minutes..."
                Start-Sleep -Seconds 300
            }
        }
    }

    # UNINSTALL === Run uninstall
    try {
        $installed_msi  = Get-ItemProperty HKLM:\Software\Microsoft\Windows\CurrentVersion\Uninstall\* |
                          Where-Object DisplayName -eq "Field Effect MDR"
        $uninstallGuid  = $installed_msi.PSChildName
        $proc = Start-Process -FilePath 'msiexec.exe' -Wait -PassThru -ArgumentList '/x', $uninstallGuid, '/qn', '/norestart'

        if ($proc.ExitCode -eq 0) {
            Log "Uninstall completed successfully"
        } else {
            Log "Uninstall completed with exit code: $($proc.ExitCode)" -level "ERROR"
        }
    } catch {
        Log "Uninstall failed: $_" -level "ERROR"
        throw $_
    }

    # UNINSTALL === Verify removal
    Start-Sleep -Seconds 3  # Allow time for uninstall to complete
    $remainingInstall = Get-ItemProperty HKLM:\Software\Microsoft\Windows\CurrentVersion\Uninstall\* |
                        Where-Object { $_.DisplayName -eq 'Field Effect MDR' }
    $remainingServices = $services | Where-Object { Get-Service -Name $_ -ErrorAction SilentlyContinue }
    $remainingPaths = @("${env:ProgramFiles}\Field Effect", "${env:ProgramData}\Field Effect") |
                      Where-Object { Test-Path $_ }

    if (-not $remainingInstall -and $remainingServices.Count -eq 0 -and $remainingPaths.Count -eq 0) {
        Log "Uninstall verification: SUCCESS - agent completely removed"
    } else {
        $issues = @()
        if ($remainingInstall)             { $issues += "registry entry" }
        if ($remainingServices.Count -gt 0) { $issues += "$($remainingServices.Count) service(s)" }
        if ($remainingPaths.Count -gt 0)    { $issues += "$($remainingPaths.Count) director(ies)" }
        Log "Uninstall verification: PARTIAL - remaining: $($issues -join ', ')" -level "ERROR"
    }

} else {
    Log "Install mode"

    # === PreScript ===
    $org_id = "x"
    $use_license = "false"
    $client_valid = "false"

    if (($license_key.Length -eq 0) -and ($client.Length -eq 0)) {
        throw "Either license_key or client must be provided."
    }

    if (($license_key.Length -ne 0) -and ($license_key.Length -ne 156)) {
        throw "License key invalid length ($($license_key.Length))"
    }

    if ($license_key.Length -eq 156) {
        $use_license = "true"
        Log "Using license key"
    } else {
        $headers = @{ Authorization = "Bearer $api_token"; ContentType = "application/json" }
        $client_list = Invoke-RestMethod -Method Get -Uri "https://services.fieldeffect.net/v1/organizations" -Headers $headers
        $client_list = $client_list -creplace '"id"', '"new_ID"' | ConvertFrom-JSON

        foreach ($c in $client_list) {
            if ($client -like $c.name -or $client -like $c.client_id) {
                $client_valid = "true"
                $org_id = $c.ID
                Log "Matched client: $($c.name)"
                break
            }
        }

        if ($client_valid -ne "true") {
            throw "Client name/id not valid"
        }
    }

    # === GetOrgId ===
    if ($use_license -eq "true") {
        Log "Using org from license"
        $headers = @{ Authorization = "Bearer $api_token"; ContentType = "application/json" }
        $orgs = Invoke-RestMethod -Method Get -Uri "https://services.fieldeffect.net/v1/organizations" -Headers $headers
        $orgs = $orgs -creplace '"id"', '"new_ID"' | ConvertFrom-JSON
        $org_id = $orgs[0].ID
        Log "Org ID is $org_id"

    }

    # === Download Installer ===
    $headers = @{ Authorization = "Bearer $api_token"; ContentType = "application/json" }
    $body = @{
        bundle_constituents = @(@{ architecture = "$win_arch bit"; platform = "Windows"; is_arm=$isArm })
        organization_id = $org_id
    } | ConvertTo-Json

    # Ensure TLS 1.2 is enabled
    [Net.ServicePointManager]::SecurityProtocol = 3072

    $resp = Invoke-RestMethod -Method Post -Uri "https://services.fieldeffect.net/v1/endpoint_installer_bundles" -Body $body -Headers $headers -ContentType application/json

    # Define paths
    $zipFileName = $resp.file_name
    $sanitizedName = $zipFileName -replace  '([\[\]\*\?])', ''
    $zipFilePath = Join-Path $working_folder $sanitizedName
    $extractFolderName = [System.IO.Path]::GetFileNameWithoutExtension($sanitizedName)
    $extractPath = Join-Path $working_folder $extractFolderName

    # Remove any leftover files or folders from previous runs
    if (Test-Path $zipFilePath) {
        Remove-Item $zipFilePath -Force
    }
    if (Test-Path $extractPath) {
        Remove-Item $extractPath -Recurse -Force
    }

    # Download the ZIP file 
    try {
        $msi = New-Object System.Net.WebClient
        $msi.DownloadFile($resp.download_link, $zipFilePath)
        $msi.Dispose()
        Log "Downloaded ZIP to $zipFilePath"
    } catch {
        Log "Download failed: $_" 
        throw "Download error: $_"
    }

    # Validate file before extraction
    if (-not (Test-Path $zipFilePath)) {
        throw "ZIP file was not found at $zipFilePath"
    }

    try {
        Expand-Archive -Path $zipFilePath -DestinationPath $extractPath -Force
        Log "Extracted ZIP contents to $extractPath"
    } catch {
        Log "Extraction failed: $_" 
        throw "Failed to extract ZIP file: $_"
    }

    # Update installer path and filename
    $script:installer_path = $extractPath
    $script:installer_filename = Get-ChildItem -Path $extractPath -Filter *.msi | Select-Object -First 1

    if (-not $script:installer_filename) {
        throw "MSI installer not found in extracted contents"
    }


    # === Get License ===
    if ($use_license -ne "true") {
        $license_key = Get-Content -Path (Join-Path $installer_path "license.txt")
        Remove-Item -Path (Join-Path $installer_path "license.txt")
        $masked = $license_key.Substring(0, 8) + ('*' * ($license_key.Length - 16)) + $license_key.Substring($license_key.Length - 8)
        Log "License key: $masked"
    }

    # === Install Agent ===
    $msiPath = Join-Path $installer_path $installer_filename
    $logPath = Join-Path $installer_path "msiexec_log.txt"
    $args = "/i `"$msiPath`" COVALENCE_LICENSE=`"$license_key`" /qn /l*v `"$logPath`""

    $proc = Start-Process msiexec.exe -Wait -PassThru -ArgumentList $args
    Log "MSI exit code: $($proc.ExitCode)"

    if ($proc.ExitCode -ne 0 -and $proc.ExitCode -ne 3010) {
        $msiLog = if (Test-Path $logPath) { Get-Content $logPath -Tail 15 | Out-String } else { "log not found" }
        Log "MSI log (last 15 lines): $msiLog" -level "ERROR"
        Log "Full install logs at $logpath"
        throw "MSI installer failed with exit code: $($proc.ExitCode)"
    }

    # Wait for status file with polling loop
    $statusFilePath = "C:\ProgramData\Field Effect\Covalence\data\status.json"
    $maxWaitSeconds = 60
    $pollIntervalSeconds = 5
    $elapsed = 0
    $statusFileFound = $false

    Log "Waiting for agent status file (max $maxWaitSeconds seconds)..."
    while ($elapsed -lt $maxWaitSeconds) {
        if (Test-Path $statusFilePath) {
            $statusFileFound = $true
            Log "Status file found after $elapsed seconds"
            break
        }
        Start-Sleep -Seconds $pollIntervalSeconds
        $elapsed += $pollIntervalSeconds
        Log "Still waiting for status file... ($elapsed/$maxWaitSeconds seconds)"
    }

    if (-not $statusFileFound) {
        throw "Install failed: status.json not found after $maxWaitSeconds seconds. Check msiexec log at $logPath"
    }

    # Clean up downloaded ZIP and extracted installer folder
    try {
        if (Test-Path $zipFilePath) {
            Remove-Item $zipFilePath -Force
            Log "Deleted ZIP file: $zipFilePath"
        }
        if (Test-Path $extractPath) {
            Remove-Item $extractPath -Recurse -Force
            Log "Deleted extracted folder: $extractPath"
        }
    } catch {
        throw "Install cleanup failed: $_"
    }

    # === PostScript ===
    $epid = Test-NetConnection -ComputerName epid.fieldeffect.net -Port 443
    Log "Identity server reachable: $($epid.TcpTestSucceeded)"

    # Status poll
    for ($i = 1; $i -le 40; $i++) {
        $status_json = Get-Content "C:\ProgramData\Field Effect\Covalence\data\status.json" | ConvertFrom-Json
        if ($status_json.AgentInfo.StateName -eq "ACTIVE") {
            Log "Agent status: ACTIVE"
            break
        }
        Start-Sleep -Seconds 10
    }

    try {
        sc.exe queryex CovAgent | Out-Null
        sc.exe queryex "Covalence Endpoint Service" | Out-Null
        Log "MDR services are running"
    } catch {
        Log "Could not confirm MDR services"
    }
}

Log "Script completed"