# AeThex OS - Fix Fastboot Driver & Flash Boot Image
# This script MUST be run as Administrator
# It will: 1) Remove broken Zadig drivers, 2) Reboot to bootloader, 3) Flash boot image

param([switch]$Elevate)

# Self-elevate if not admin
$isAdmin = ([Security.Principal.WindowsPrincipal] [Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
if (-not $isAdmin) {
    Write-Host "Requesting Administrator privileges..." -ForegroundColor Yellow
    Start-Process PowerShell -Verb RunAs -ArgumentList "-ExecutionPolicy Bypass -File `"$PSCommandPath`" -Elevate"
    exit
}

$FASTBOOT = "C:\Users\PCOEM\platform-tools\fastboot.exe"
$BOOT_IMG = "C:\Users\PCOEM\AeThexOS\magisk_patched_boot.img"

Write-Host ""
Write-Host "==========================================" -ForegroundColor Cyan
Write-Host "  AeThex OS - Fix Driver & Flash Boot" -ForegroundColor Cyan
Write-Host "==========================================" -ForegroundColor Cyan
Write-Host ""

# Verify boot image exists
if (-not (Test-Path $BOOT_IMG)) {
    Write-Host "ERROR: Boot image not found at $BOOT_IMG" -ForegroundColor Red
    Read-Host "Press Enter to exit"
    exit 1
}
$size = (Get-Item $BOOT_IMG).Length
Write-Host "[OK] Boot image: $size bytes" -ForegroundColor Green

# ================================================
# PHASE 1: Remove broken Zadig USB drivers
# ================================================
Write-Host ""
Write-Host "--- PHASE 1: Fixing USB Drivers ---" -ForegroundColor Yellow
Write-Host ""

# Find and remove all Zadig (libwdi) drivers for MediaTek and Android Bootloader
$output = pnputil /enum-drivers 2>&1 | Out-String
$blocks = $output -split "(?=Published Name:)"
$removed = 0

foreach ($block in $blocks) {
    if ($block -match "libwdi") {
        if ($block -match "Published Name:\s+(oem\d+\.inf)") {
            $drvName = $Matches[1]
            Write-Host "  Removing Zadig driver: $drvName" -ForegroundColor Red
            $result = pnputil /delete-driver $drvName /force 2>&1
            if ($result -match "deleted") {
                Write-Host "    -> Removed!" -ForegroundColor Green
                $removed++
            } else {
                Write-Host "    -> $($result -join ' ')" -ForegroundColor DarkGray
            }
        }
    }
}

Write-Host "  Removed $removed Zadig driver(s)" -ForegroundColor Cyan

# Remove phantom device entries
Write-Host ""
Write-Host "  Cleaning phantom device entries..." -ForegroundColor Yellow
$phantoms = Get-PnpDevice | Where-Object {
    ($_.InstanceId -like '*VID_18D1&PID_4EE0*' -or
     $_.InstanceId -like '*VID_0E8D&PID_201C*') -and
    $_.Status -ne 'OK'
}
foreach ($p in $phantoms) {
    Write-Host "    Removing phantom: $($p.FriendlyName)" -ForegroundColor Gray
    pnputil /remove-device "$($p.InstanceId)" 2>&1 | Out-Null
}

# Ensure Google USB driver is in driver store
Write-Host ""
Write-Host "  Ensuring Google USB driver is available..." -ForegroundColor Yellow
$googleInfPath = "$env:LOCALAPPDATA\Android\Sdk\extras\google\usb_driver\android_winusb.inf"
if (Test-Path $googleInfPath) {
    pnputil /add-driver $googleInfPath /install 2>&1 | Out-Null
    Write-Host "    Google USB driver ready" -ForegroundColor Green
}

# Rescan
pnputil /scan-devices 2>&1 | Out-Null
Write-Host "  Hardware scan complete" -ForegroundColor Green

# ================================================
# PHASE 2: Reboot to bootloader and flash
# ================================================
Write-Host ""
Write-Host "--- PHASE 2: Flash Boot Image ---" -ForegroundColor Yellow
Write-Host ""

# Check if device is connected via ADB
$adbResult = adb devices 2>&1 | Out-String
if ($adbResult -match "device$" -or $adbResult -match "device\s") {
    Write-Host "[OK] Device connected via ADB" -ForegroundColor Green

    # Reboot to bootloader
    Write-Host "  Rebooting to bootloader..." -ForegroundColor Yellow
    adb reboot bootloader 2>&1 | Out-Null
} else {
    Write-Host "[!] Device not connected via ADB" -ForegroundColor Yellow
    Write-Host "    Please reboot tablet to bootloader manually:" -ForegroundColor Yellow
    Write-Host "    Power off -> Hold Vol Up + Power -> Select FASTBOOT" -ForegroundColor White
}

# Poll for fastboot device
Write-Host ""
Write-Host "  Waiting for fastboot device..." -ForegroundColor Yellow
Write-Host "  (If 'SELECT BOOT MODE' appears, choose FASTBOOT)" -ForegroundColor DarkGray

$found = $false
$attempts = 0

while (-not $found -and $attempts -lt 120) {
    $attempts++
    $result = & $FASTBOOT devices 2>&1 | Out-String
    if ($result -match "\S+\s+fastboot") {
        $found = $true
        Write-Host ""
        Write-Host "  DEVICE FOUND!" -ForegroundColor Green
        Write-Host "  $($result.Trim())" -ForegroundColor White
    } else {
        if ($attempts % 10 -eq 0) {
            Write-Host "    Still waiting... ($([math]::Round($attempts/2))s)" -ForegroundColor DarkGray
        }
        Start-Sleep -Milliseconds 500
    }
}

if (-not $found) {
    Write-Host ""
    Write-Host "  TIMEOUT waiting for fastboot device." -ForegroundColor Red
    Write-Host ""
    Write-Host "  Trying to check what USB devices exist..." -ForegroundColor Yellow
    Get-PnpDevice -Status 'OK' | Where-Object { $_.InstanceId -like '*0E8D*' -or $_.InstanceId -like '*18D1*' } |
        Select-Object Status, Class, InstanceId, FriendlyName | Format-Table -AutoSize
    Write-Host ""
    Write-Host "  If a device IS connected but fastboot can't see it:" -ForegroundColor Yellow
    Write-Host "  The device may use MediaTek's proprietary fastboot." -ForegroundColor Yellow
    Write-Host "  Try: python C:\Users\PCOEM\mtkclient\mtk.py w boot_b $BOOT_IMG" -ForegroundColor White
    Read-Host "Press Enter to exit"
    exit 1
}

# Flash boot_b
Write-Host ""
Write-Host "  Flashing boot_b partition..." -ForegroundColor Cyan
& $FASTBOOT flash boot_b $BOOT_IMG 2>&1

if ($LASTEXITCODE -eq 0) {
    Write-Host ""
    Write-Host "  =============================" -ForegroundColor Green
    Write-Host "  SUCCESS! Magisk boot flashed!" -ForegroundColor Green
    Write-Host "  =============================" -ForegroundColor Green
    Write-Host ""
    Write-Host "  Rebooting device..." -ForegroundColor Yellow
    & $FASTBOOT reboot 2>&1
} else {
    Write-Host "  boot_b failed, trying 'boot'..." -ForegroundColor Yellow
    & $FASTBOOT flash boot $BOOT_IMG 2>&1
    if ($LASTEXITCODE -eq 0) {
        Write-Host "  SUCCESS with 'boot' partition!" -ForegroundColor Green
        & $FASTBOOT reboot 2>&1
    } else {
        Write-Host ""
        Write-Host "  Both partition names failed." -ForegroundColor Red
        Write-Host "  Checking available partitions..." -ForegroundColor Yellow
        & $FASTBOOT getvar all 2>&1
    }
}

Write-Host ""
Read-Host "Press Enter to exit"