@(set "0=%~f0"^)#) & powershell -nop -c iex([io.file]::ReadAllText($env:0)) & exit /b
$CMDLINE = "root=PARTLABEL=linux loglevel=3"

function Log {
  Write-Host -NoNewline "==> " -ForegroundColor Green
  Write-Host "$args"
}

function LogErr {
  Write-Host -NoNewline "==> " -ForegroundColor Red
  Write-Host "$args"
}


$PYTHON = if (Get-Command python -ErrorAction SilentlyContinue) {
  "python"
} else {
  LogErr "Python 3 not found"
  exit 1
}


# Begin script

if ($args.Count -eq 2) {
  $VMLINUZ = $args[0]
  $DTB = $args[1]
  $INTERACTIVE = $false
} else {
  $VMLINUZ = Read-Host "Drag and drop vmlinuz here and press enter"
  $DTB = Read-Host "Drag and drop dtb here and press enter"
  $INTERACTIVE = $true
}

# Check if files exist
if (-not (Test-Path $VMLINUZ)) {
  LogErr "vmlinuz not found at $VMLINUZ"
  exit 1
}

if (-not (Test-Path $DTB)) {
  LogErr "dtb not found at $DTB"
  exit 1
}

# Download mkbootimg and related files if not present
if (-not (Test-Path ./mkbootimg.py) -or -not (Test-Path ./generate_gki_certificate.py)) {
  Log "Downloading mkbootimg from Google"

  $err = $false

  try {
    $mkbootimgUrl = "https://android.googlesource.com/platform/system/tools/mkbootimg/+/refs/heads/main/mkbootimg.py?format=text"
    $mkbootimgContent = (Invoke-WebRequest -Uri $mkbootimgUrl).Content
    $mkbootimgDecoded = [System.Text.Encoding]::UTF8.GetString([System.Convert]::FromBase64String($mkbootimgContent))
    $mkbootimgDecoded = $mkbootimgDecoded -replace 'from gki.generate_gki_certificate import generate_gki_certificate', 'from generate_gki_certificate import generate_gki_certificate'
    $mkbootimgDecoded | Out-File -FilePath ./mkbootimg.py -Encoding UTF8
  } catch {
    $err = $true
  }

  try {
    $gkiCertUrl = "https://android.googlesource.com/platform/system/tools/mkbootimg/+/refs/heads/main/gki/generate_gki_certificate.py?format=text"
    $gkiCertContent = (Invoke-WebRequest -Uri $gkiCertUrl).Content
    $gkiCertDecoded = [System.Text.Encoding]::UTF8.GetString([System.Convert]::FromBase64String($gkiCertContent))
    $gkiCertDecoded | Out-File -FilePath ./generate_gki_certificate.py -Encoding UTF8
  } catch {
    $err = $true
  }

  if ($err) {
    LogErr "Failed to download mkbootimg"
    exit 1
  }
}

# Combine vmlinuz and dtb into a single zImage file
try {
  Get-Content -Path $VMLINUZ, $DTB -AsByteStream -Read 1024 | Set-Content -Path .\zImage -AsByteStream
} catch {
  try {
    Get-Content -Path $VMLINUZ, $DTB -Encoding Byte -Read 1024 | Set-Content -Path .\zImage -Encoding Byte
  } catch {
    LogErr "Failed to create zImage"
    exit 1 
  }
}

# Create the boot image using Python script
$pythonCmd = "$PYTHON mkbootimg.py --kernel zImage --cmdline `"$CMDLINE`" --base 0x00000000 --kernel_offset 0x00008000 --tags_offset 0x00000100 --pagesize 4096 --id -o linux.boot.img"
Invoke-Expression $pythonCmd | Out-Null

if ($LASTEXITCODE -ne 0) {
  LogErr "Failed to create boot image"
  Remove-Item .\zImage
  exit 1
} else {
  Log "Created: linux.boot.img"
}

# Clean up
Remove-Item .\zImage
Remove-Item -Recurse -Force .\__pycache__

if ($INTERACTIVE) {
  $null = $Host.UI.RawUI.ReadKey('NoEcho,IncludeKeyDown');
}