# ============================================================================
# Script Parameters
# ============================================================================
param(
    [switch]$Debug
)

# ============================================================================
# Debug Logging Function
# ============================================================================
function Write-DebugLog {
    param([string]$Message)
    if ($Debug) {
        Write-Host "[DEBUG] $Message" -ForegroundColor DarkGray
    }
}

# ============================================================================
# Administrator Privilege Check
# ============================================================================
Write-Host "`n========================================" -ForegroundColor Cyan
Write-Host "Firmware Update Tool" -ForegroundColor Cyan
Write-Host "========================================`n" -ForegroundColor Cyan

Write-Host "[INFO] Checking administrator privileges..." -ForegroundColor Yellow

# Debug: Show debug mode status
if ($Debug) {
    Write-Host "[DEBUG] Debug mode ENABLED" -ForegroundColor DarkGray
}

$isAdmin = ([Security.Principal.WindowsPrincipal] [Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)

Write-DebugLog "Admin check result: $isAdmin"

if (-not $isAdmin) {
    Write-Host "`n[ERROR] Administrator privileges required!" -ForegroundColor Red
    Write-Host "[ERROR] Please run this script as Administrator." -ForegroundColor Red
    Write-Host "`nPress any key to exit..." -ForegroundColor Yellow
    $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
    exit 1
}

# Change to the script's directory
Set-Location -Path $PSScriptRoot

# Store script directory path for reliable executable access
$scriptDir = $PSScriptRoot

Write-DebugLog "Script directory: $scriptDir"
Write-DebugLog "PSScriptRoot: $PSScriptRoot"
Write-DebugLog "Current Location: $(Get-Location)"

# ============================================================================
# Check for Marvell Driver and Version
# ============================================================================
Write-Host "[INFO] Checking for Marvell network driver..." -ForegroundColor Yellow

$requiredVersion = [Version]"3.1.11"

# Get all Marvell devices using Get-PnpDevice (more reliable for updated drivers without reboot)
# WMI Win32_PnPSignedDriver caches driver info and requires reboot to refresh
$marvellDevices = Get-PnpDevice | Where-Object { $_.InstanceId -like "*PCI\VEN_1D6A*" }

Write-DebugLog "PnpDevice Query completed for Marvell devices"
Write-DebugLog "Found devices count: $($marvellDevices.Count)"

if (-not $marvellDevices) {
    Write-Host "`n[ERROR] No Marvell network devices found in the system!" -ForegroundColor Red
    Write-Host "`nPress any key to exit..." -ForegroundColor Yellow
    $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
    exit 1
}

Write-Host "[SUCCESS] Found $($marvellDevices.Count) Marvell device(s)`n" -ForegroundColor Green

# Check driver versions using Get-PnpDeviceProperty for real-time driver version
$allDriversValid = $true
$deviceCount = 0

foreach ($device in $marvellDevices) {
    $deviceCount++
    $deviceID = $device.InstanceId
    
    # Get driver version from device properties (reflects updates without reboot)
    $driverVersionProp = Get-PnpDeviceProperty -InstanceId $deviceID -KeyName "DEVPKEY_Device_DriverVersion" -ErrorAction SilentlyContinue
    $driverVersion = $driverVersionProp.Data
    
    # Get provider name
    $providerProp = Get-PnpDeviceProperty -InstanceId $deviceID -KeyName "DEVPKEY_Device_DriverProvider" -ErrorAction SilentlyContinue
    $driverProvider = $providerProp.Data
    
    # Get INF name
    $infProp = Get-PnpDeviceProperty -InstanceId $deviceID -KeyName "DEVPKEY_Device_DriverInfPath" -ErrorAction SilentlyContinue
    $infName = $infProp.Data
    
    Write-Host "Device #$deviceCount" -ForegroundColor Magenta
    Write-Host "  Device ID: $deviceID" -ForegroundColor Cyan
    Write-Host "  Driver Version: $driverVersion" -ForegroundColor Cyan
    Write-Host "  Provider: $driverProvider" -ForegroundColor Cyan
    Write-Host "  INF File: $infName" -ForegroundColor Cyan
    
    # Try to parse the driver version
    if ($driverVersion) {
        try {
            $currentVersion = [Version]$driverVersion
            
            if ($currentVersion -ge $requiredVersion) {
                Write-Host "  Status: Version OK ($driverVersion >= $requiredVersion)" -ForegroundColor Green
            } else {
                Write-Host "  Status: Driver update required (Current: $driverVersion, Required: $requiredVersion or higher)" -ForegroundColor Yellow
                $allDriversValid = $false
            }
        } catch {
            Write-Host "  Status: Could not parse version '$driverVersion'" -ForegroundColor Red
            $allDriversValid = $false
        }
    } else {
        Write-Host "  Status: No version information available" -ForegroundColor Red
        $allDriversValid = $false
    }
    Write-Host ""
}

if (-not $allDriversValid) {
    Write-Host "`n========================================" -ForegroundColor Red
    Write-Host "DRIVER UPDATE REQUIRED" -ForegroundColor Red
    Write-Host "========================================" -ForegroundColor Red
    Write-Host "`nYour current Marvell driver version is outdated." -ForegroundColor Yellow
    Write-Host "To proceed with the firmware update, please:" -ForegroundColor Yellow
    Write-Host "  1. Download Marvell driver version $requiredVersion or newer" -ForegroundColor Cyan
    Write-Host "  2. Install the updated driver" -ForegroundColor Cyan
    Write-Host "  3. Run this firmware update tool again" -ForegroundColor Cyan
    Write-Host "`nPress any key to exit..." -ForegroundColor Yellow
    $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
    exit 1
}

Write-Host "[SUCCESS] Driver version check passed`n" -ForegroundColor Green

# ============================================================================
# Verify Devices with pnputil
# ============================================================================
$targetHardwareId = "PCI\\VEN_1D6A"
$pnpDevices = pnputil /enum-devices | Select-String -Pattern $targetHardwareId

if (-not $pnpDevices) {
    Write-Host "`n[ERROR] No devices with hardware ID '$targetHardwareId' found!" -ForegroundColor Red
    Write-Host "`nPress any key to exit..." -ForegroundColor Yellow
    $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
    exit 1
}

# ============================================================================
# Get NDISIO Device IDs
# ============================================================================
Write-Host "[INFO] Enumerating NDISIO devices..." -ForegroundColor Yellow
Write-DebugLog "Executing: $scriptDir\listDevices.exe"
$output1 = & "$scriptDir\listDevices.exe"
Write-DebugLog "listDevices.exe output: $output1"
$devIds = $output1 | Select-String "NDISIO:" | ForEach-Object { ($_ -split '\s+') | Where-Object { $_ -match '^NDISIO:' } }
Write-DebugLog "Parsed NDISIO device IDs: $devIds"

if (-not $devIds -or $devIds.Count -eq 0) {
    Write-Host "`n[ERROR] No NDISIO devices found!" -ForegroundColor Red
    Write-Host "`nPress any key to exit..." -ForegroundColor Yellow
    $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
    exit 1
}

Write-Host "[SUCCESS] Found $($devIds.Count) NDISIO device(s)`n" -ForegroundColor Green

# ============================================================================
# Filter out Dual 10G devices
# ============================================================================
Write-Host "[INFO] Checking for Dual 10G devices..." -ForegroundColor Yellow

$filteredDevIds = @()
$skippedDevIds = @()

# Parse bus numbers from all device IDs
$deviceBusMap = @{}
foreach ($devId in $devIds) {
    # Extract bus number from format: NDISIO:04C0-0D2:00.0
    if ($devId -match 'NDISIO:\w+-([0-9A-Fa-f]+):') {
        $busNumber = [Convert]::ToInt32($Matches[1], 16)
        $deviceBusMap[$devId] = $busNumber
        Write-DebugLog "Device: $devId, Bus Number: $busNumber (0x$($Matches[1]))"
    }
}

# Check for Dual 10G devices and filter - UPDATE standalone, SKIP Dual 10G
foreach ($devId in $devIds) {
    $currentBus = $deviceBusMap[$devId]
    $hasConsecutive = $false
    
    # Check if any other device is part of a Dual 10G adapter
    foreach ($otherDevId in $devIds) {
        if ($devId -ne $otherDevId) {
            $otherBus = $deviceBusMap[$otherDevId]
            
            # Check if devices are Dual 10G (differ by 1)
            if ([Math]::Abs($currentBus - $otherBus) -eq 1) {
                Write-DebugLog "Found Dual 10G adapter: $devId (0x$($currentBus.ToString('X'))) and $otherDevId (0x$($otherBus.ToString('X')))"
                
                # Skip BOTH ports of Dual 10G adapter
                $hasConsecutive = $true
                #Write-Host "  [SKIP] Device $devId - part of Dual 10G adapter with $otherDevId" -ForegroundColor Yellow
                $skippedDevIds += $devId
                break
            }
        }
    }
    
    if (-not $hasConsecutive) {
        # Process standalone devices
        Write-Host "  [UPDATE] Device $devId - standalone device" -ForegroundColor Green
        $filteredDevIds += $devId
    }
}

if ($skippedDevIds.Count -gt 0) {
    #Write-Host "[INFO] Skipped $($skippedDevIds.Count) Dual 10G device(s)" -ForegroundColor Yellow
    Write-Host "[INFO] Will process $($filteredDevIds.Count) standalone device(s)`n" -ForegroundColor Green
} else {
    Write-Host "[INFO] Processing all devices`n" -ForegroundColor Green
}

# Use filtered device list for processing
$devIds = $filteredDevIds

if ($devIds.Count -eq 0) {
    Write-Host "`n[ERROR] No devices to process after filtering!" -ForegroundColor Red
    Write-Host "`nPress any key to exit..." -ForegroundColor Yellow
    $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
    exit 1
}

# ============================================================================
# Process Each Device
# ============================================================================
$totalDevices = $devIds.Count
$currentDevice = 0
$successfulUpdates = 0

foreach($devId in $devIds) {
    $currentDevice++
    Write-Host "========================================" -ForegroundColor Cyan
    Write-Host "Processing Device $currentDevice of $totalDevices" -ForegroundColor Cyan
    Write-Host "Device ID: $devId" -ForegroundColor Cyan
    Write-Host "========================================`n" -ForegroundColor Cyan

    # Getting four-part IDs for the device
    Write-Host "[INFO] Reading device IDs..." -ForegroundColor Yellow
    Write-DebugLog "Executing: $scriptDir\flashUpdate2.exe -d $devId -i"
    $output2 = & "$scriptDir\flashUpdate2.exe" -d $devId -i
    Write-DebugLog "flashUpdate2.exe output: $output2"
    $ids = $output2 | Select-String "FourPartIDs:" | ForEach-Object { $_.Line.Split(":") }
    Write-DebugLog "Parsed IDs array: $($ids -join ', ')"

    If ($ids -eq $null -or $ids.count -eq 0) {
        Write-Host "[ERROR] Failed to get four part IDs for device $devId" -ForegroundColor Red
        Write-Host "[ERROR] Skipping this device...`n" -ForegroundColor Red
        continue
    }

    Write-Host "[SUCCESS] Retrieved IDs: VID=$($ids[1]), DID=$($ids[2]), SVID=$($ids[3]), SDID=$($ids[4])" -ForegroundColor Green

    # Validate the four-part IDs
    If (($ids.count -ne 5) -or ($ids[1] -eq "0x0000") -or ($ids[2] -eq "0x0000") -or ($ids[3] -eq "0x0000") -or ($ids[4] -eq "0x0000")) {
        Write-Host "[ERROR] Invalid four part IDs for device ${devId}" -ForegroundColor Red
        Write-Host "[ERROR] VID: $($ids[1]), DID: $($ids[2]), SVID: $($ids[3]), SDID: $($ids[4])" -ForegroundColor Red
        Write-Host "[ERROR] Skipping this device...`n" -ForegroundColor Red
        continue
    }
    Write-Host "[SUCCESS] IDs validation passed" -ForegroundColor Green
    
    # Check if device matches required IDs
    # Supported configuration:
    #   VID=0x1D6A, DID=0x04C0, SVID=1043, SDID=8890
    $deviceSupported = $false

    If (($ids[1] -eq "0x1D6A") -and ($ids[2] -eq "0x04C0") -and ($ids[3] -eq "0x1043") -and ($ids[4] -eq "0x8890")) {
        $deviceSupported = $true
    }

    If (-not $deviceSupported) {
        Write-Host "[SKIP] Device does not match required IDs" -ForegroundColor Yellow
        Write-Host "[INFO] Found: VID=$($ids[1]), DID=$($ids[2]), SVID=$($ids[3]), SDID=$($ids[4])" -ForegroundColor Yellow
        Write-Host "[INFO] Required: VID=0x1D6A, DID=0x04C0, SVID=0x1043, SDID=0x8890" -ForegroundColor Yellow
        Write-Host "[INFO] Skipping this device...`n" -ForegroundColor Yellow
        continue
    }
    Write-Host "[SUCCESS] Device IDs match - proceeding with firmware update" -ForegroundColor Green

    # Flash the device with firmware
    Write-Host "[INFO] Flashing device (this may take a moment)..." -ForegroundColor Yellow
    Write-DebugLog "Executing: $scriptDir\flashUpdate2.exe -d $devId --reflash $scriptDir\ATL2-8.0.1_agent_SIGN.clx $scriptDir\AQC113-Lenovo-Vulcan-1.5.48_aqsign.clx"
    $flashOutput = & "$scriptDir\flashUpdate2.exe" -d $devId --reflash "$scriptDir\ATL2-8.0.1_agent_SIGN.clx" "$scriptDir\AQC113-Lenovo-Vulcan-1.5.48_aqsign.clx" 2>&1
    Write-DebugLog "flashUpdate2.exe output: $flashOutput"
    Write-DebugLog "flashUpdate2.exe exit code: $LASTEXITCODE"
    
    # Check if flash was successful
    if ($LASTEXITCODE -eq 0 -or $flashOutput -match "Status: Success") {
        Write-Host "[SUCCESS] Device flashed successfully" -ForegroundColor Green
        Write-Host "[SUCCESS] Firmware upgraded to version 1.5.48`n" -ForegroundColor Green
        $successfulUpdates++
    } else {
        Write-Host "[ERROR] Flash operation failed!" -ForegroundColor Red
        Write-Host $flashOutput
        Write-Host "[ERROR] Skipping this device...`n" -ForegroundColor Red
        continue
    }

    # -----------------------------------------------------------------------
    # Disable / Enable network adapter so new firmware rings are allocated
    # -----------------------------------------------------------------------
    Write-Host "[INFO] Cycling network adapter to allocate new firmware rings..." -ForegroundColor Yellow

    # Extract bus number from NDISIO device ID (format: NDISIO:DDDD-BBB:SS.F)
    if ($devId -match 'NDISIO:\w+-([0-9A-Fa-f]+):') {
        $busHex    = $Matches[1]
        $busNumber = [Convert]::ToInt32($busHex, 16)
        Write-DebugLog "Looking for PnP device on bus 0x$busHex (decimal: $busNumber)"

        # Find the matching Marvell PnP device by bus number
        $matchedDevice = $null
        foreach ($pnpDev in $marvellDevices) {
            $busProp = Get-PnpDeviceProperty -InstanceId $pnpDev.InstanceId `
                           -KeyName "DEVPKEY_Device_BusNumber" -ErrorAction SilentlyContinue
            Write-DebugLog "  Checking $($pnpDev.InstanceId) - bus $($busProp.Data)"
            if ($busProp.Data -eq $busNumber) {
                $matchedDevice = $pnpDev
                Write-DebugLog "  Matched: $($pnpDev.InstanceId)"
                break
            }
        }

        if ($matchedDevice) {
            Write-Host "[INFO] Disabling adapter: $($matchedDevice.FriendlyName)..." -ForegroundColor Yellow
            Disable-PnpDevice -InstanceId $matchedDevice.InstanceId -Confirm:$false -ErrorAction SilentlyContinue
            Start-Sleep -Seconds 3

            Write-Host "[INFO] Enabling adapter: $($matchedDevice.FriendlyName)..." -ForegroundColor Yellow
            Enable-PnpDevice -InstanceId $matchedDevice.InstanceId -Confirm:$false -ErrorAction SilentlyContinue
            Start-Sleep -Seconds 3

            Write-Host "[SUCCESS] Network adapter cycled successfully`n" -ForegroundColor Green
        } else {
            Write-Host "[WARNING] Could not find matching PnP device for bus 0x$busHex - skipping adapter cycle`n" -ForegroundColor Yellow
        }
    } else {
        Write-Host "[WARNING] Could not parse bus number from device ID: $devId - skipping adapter cycle`n" -ForegroundColor Yellow
    }
}

# ============================================================================
# Script Completion
# ============================================================================
Write-Host "`n========================================" -ForegroundColor Green
Write-Host "Firmware Update Completed Successfully!" -ForegroundColor Green
Write-Host "========================================`n" -ForegroundColor Green

# Only ask for reboot if at least one device was successfully updated
if ($successfulUpdates -gt 0) {
    Write-Host "[INFO] Successfully updated $successfulUpdates device(s)" -ForegroundColor Green

    # Show Windows message box informing about reboot requirement
    Add-Type -AssemblyName System.Windows.Forms
    [System.Windows.Forms.MessageBox]::Show(
        "Firmware update completed successfully!`n`nSuccessfully updated $successfulUpdates device(s).`n`nPlease reboot your system to complete the firmware update.",
        "Firmware Update Complete",
        [System.Windows.Forms.MessageBoxButtons]::OK,
        [System.Windows.Forms.MessageBoxIcon]::Information
    )

    Write-Host "`n[IMPORTANT] Please reboot your system to complete the firmware update." -ForegroundColor Yellow
    Write-Host "`nPress any key to exit..." -ForegroundColor Yellow
    $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
} else {
    Write-Host "[WARNING] No devices were successfully updated." -ForegroundColor Yellow
    Write-Host "`nPress any key to exit..." -ForegroundColor Yellow
    $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
}