param(
    [string]$Router = "root@192.168.1.1",
    [string]$OutputDir = $PSScriptRoot,
    [int64]$ChunkSize = 16MB
)

$ErrorActionPreference = "Stop"

function Invoke-Router {
    param([string]$Command)
    & ssh $Router $Command
    if ($LASTEXITCODE -ne 0) {
        throw "SSH command failed: $Command"
    }
}

function Copy-FromRouter {
    param(
        [string]$RemotePath,
        [string]$LocalPath
    )
    & scp -O "${Router}:$RemotePath" $LocalPath
    if ($LASTEXITCODE -ne 0) {
        throw "SCP failed: $RemotePath"
    }
}

function Sanitize-Name {
    param([string]$Name)
    if ([string]::IsNullOrWhiteSpace($Name)) { return "unnamed" }
    return ($Name -replace '[\\/:*?"<>| ]', '_')
}

if ([string]::IsNullOrWhiteSpace($OutputDir)) {
    $OutputDir = (Get-Location).Path
}

New-Item -ItemType Directory -Force -Path $OutputDir | Out-Null
$OutputDir = (Resolve-Path $OutputDir).Path

$restoreFile = Join-Path $OutputDir "restore_commands.txt"
$manifestFile = Join-Path $OutputDir "backup_manifest.txt"
$restoreLines = New-Object System.Collections.Generic.List[string]
$manifestLines = New-Object System.Collections.Generic.List[string]

Write-Host "Fetching MTD partition list..."
$mtdLines = & ssh $Router "cat /proc/mtd"
if ($LASTEXITCODE -ne 0) {
    throw "Failed to read /proc/mtd"
}

Write-Host "Cleaning router temp files..."
Invoke-Router "rm -f /tmp/*.bin /tmp/*.part*"

$partitions = @()

foreach ($line in $mtdLines) {
    if ($line -match '^(mtd\d+):\s+([0-9a-fA-F]+)\s+[0-9a-fA-F]+\s+"([^"]+)"$') {
        $dev = $matches[1]
        $nameFromProc = $matches[3]

        $sysName = (& ssh $Router "cat /sys/class/mtd/$dev/name").Trim()
        if ($LASTEXITCODE -ne 0) { throw "Failed to read /sys/class/mtd/$dev/name" }

        $sysSize = (& ssh $Router "cat /sys/class/mtd/$dev/size").Trim()
        if ($LASTEXITCODE -ne 0) { throw "Failed to read /sys/class/mtd/$dev/size" }

        $sysType = (& ssh $Router "cat /sys/class/mtd/$dev/type").Trim()
        if ($LASTEXITCODE -ne 0) { throw "Failed to read /sys/class/mtd/$dev/type" }

        $size = [int64]$sysSize
        $name = if ($sysName) { $sysName } else { $nameFromProc }

        $partitions += [pscustomobject]@{
            Dev = $dev
            Name = $name
            SafeName = (Sanitize-Name $name)
            Size = $size
            SizeHex = ('0x{0:X}' -f $size)
            Type = $sysType
        }
    }
}

if (-not $partitions) {
    throw "No MTD partitions found"
}

Write-Host ""
Write-Host "Detected partitions:"
$partitions | ForEach-Object {
    Write-Host ("  {0}  {1}  {2} bytes  type={3}" -f $_.Dev, $_.Name, $_.Size, $_.Type)
}

$manifestLines.Add("Backup manifest")
$manifestLines.Add("----------------------------------------")
$manifestLines.Add(("Router: {0}" -f $Router))
$manifestLines.Add(("Output directory: {0}" -f $OutputDir))
$manifestLines.Add(("Generated: {0}" -f (Get-Date -Format "yyyy-MM-dd HH:mm:ss zzz")))
$manifestLines.Add("")

foreach ($p in $partitions) {
    $finalFile = Join-Path $OutputDir ($p.SafeName + ".bin")
    Write-Host ""
    Write-Host ("Backing up {0} ({1}, {2} bytes, type={3})..." -f $p.Dev, $p.Name, $p.Size, $p.Type)

    if ($p.Type -eq "nor") {
        $remote = "/tmp/$($p.SafeName).bin"
        Invoke-Router "rm -f '$remote' && dd if=/dev/$($p.Dev) of='$remote' bs=64k"
        Copy-FromRouter -RemotePath $remote -LocalPath $finalFile
        Invoke-Router "rm -f '$remote'"
    }
    elseif ($p.Type -in @("nand", "mlc-nand")) {
        if ($p.Size -le $ChunkSize) {
            $remote = "/tmp/$($p.SafeName).bin"
            Invoke-Router "rm -f '$remote' && nanddump --omitoob -f '$remote' /dev/$($p.Dev)"
            Copy-FromRouter -RemotePath $remote -LocalPath $finalFile
            Invoke-Router "rm -f '$remote'"
        }
        else {
            $partFiles = @()
            $offset = 0
            $index = 1

            while ($offset -lt $p.Size) {
                $remaining = $p.Size - $offset
                $thisChunk = [Math]::Min($ChunkSize, $remaining)

                $remotePart = "/tmp/$($p.SafeName).part$index.bin"
                $localPart  = Join-Path $OutputDir "$($p.SafeName).part$index.bin"

                $offsetHex = ('0x{0:X}' -f $offset)
                $lengthHex = ('0x{0:X}' -f $thisChunk)

                Write-Host ("  Chunk {0}: offset {1}, length {2}" -f $index, $offsetHex, $lengthHex)

                Invoke-Router "rm -f '$remotePart' && nanddump --omitoob -s $offsetHex -l $lengthHex -f '$remotePart' /dev/$($p.Dev)"
                Copy-FromRouter -RemotePath $remotePart -LocalPath $localPart
                Invoke-Router "rm -f '$remotePart'"

                $partFiles += $localPart
                $offset += $thisChunk
                $index += 1
            }

            $out = [System.IO.File]::Open($finalFile, [System.IO.FileMode]::Create, [System.IO.FileAccess]::Write)
            try {
                foreach ($pf in $partFiles) {
                    $in = [System.IO.File]::OpenRead($pf)
                    try {
                        $in.CopyTo($out)
                    }
                    finally {
                        $in.Dispose()
                    }
                }
            }
            finally {
                $out.Dispose()
            }

            if ((Get-Item $finalFile).Length -gt 0) {
                Remove-Item $partFiles -Force
            }
            else {
                throw "Combined file is empty: $finalFile"
            }
        }
    }
    else {
        throw "Unsupported MTD type '$($p.Type)' for $($p.Dev)"
    }

    $saved = Get-Item $finalFile
    $hash = (Get-FileHash -Algorithm SHA256 -Path $finalFile).Hash

    Write-Host ("  Saved: {0}" -f $finalFile)
    Write-Host ("  SHA256: {0}" -f $hash)

    $manifestLines.Add(("[{0}] {1}" -f $p.Dev, $p.Name))
    $manifestLines.Add(("  Type: {0}" -f $p.Type))
    $manifestLines.Add(("  Reported size: {0} bytes ({1})" -f $p.Size, $p.SizeHex))
    $manifestLines.Add(("  Output file: {0}" -f $saved.Name))
    $manifestLines.Add(("  Output size: {0} bytes" -f $saved.Length))
    $manifestLines.Add(("  SHA256: {0}" -f $hash))
    $manifestLines.Add("")
}

$restoreLines.Add("USB restore prep")
$restoreLines.Add("----------------------------------------")
$restoreLines.Add("Install USB storage and filesystem support if needed:")
$restoreLines.Add("  opkg update")
$restoreLines.Add("  opkg install kmod-usb-storage kmod-usb3 kmod-fs-vfat block-mount ntfs-3g")
$restoreLines.Add("")
$restoreLines.Add("Plug in the USB flash drive, then identify it:")
$restoreLines.Add("  dmesg | tail -n 50")
$restoreLines.Add("  block info")
$restoreLines.Add("  ls /dev/sd*")
$restoreLines.Add("")
$restoreLines.Add("Create a mount point and mount the USB drive (example uses /dev/sda1):")
$restoreLines.Add("  mkdir -p /mnt/usb")
$restoreLines.Add("  mount /dev/sda1 /mnt/usb")
$restoreLines.Add("")
$restoreLines.Add("Verify your backup files are present:")
$restoreLines.Add("  ls -l /mnt/usb")
$restoreLines.Add("")
$restoreLines.Add("Restore commands")
$restoreLines.Add("----------------------------------------")
$restoreLines.Add("")

foreach ($p in $partitions) {
    $fileName = "$($p.SafeName).bin"
    $usbFile = "/mnt/usb/$fileName"
    $mtdNum = ($p.Dev -replace '^mtd','')

    $restoreLines.Add("[$($p.Dev)] $($p.Name)")

    if ($p.Name -eq "ubi" -and $p.Type -in @("nand", "mlc-nand")) {
        $restoreLines.Add("Restore from USB (UBI image):")
        $restoreLines.Add("  ubidetach -m $mtdNum 2>/dev/null")
        $restoreLines.Add("  ubiformat /dev/$($p.Dev) -f $usbFile")
    }
    elseif ($p.Type -in @("nand", "mlc-nand")) {
        $restoreLines.Add("Restore from USB (raw NAND MTD):")
        $restoreLines.Add("  flash_erase /dev/$($p.Dev) 0 0")
        $restoreLines.Add("  nandwrite -p /dev/$($p.Dev) $usbFile")
    }
    elseif ($p.Type -eq "nor") {
        $restoreLines.Add("Restore from USB (raw NOR MTD):")
        $restoreLines.Add("  mtd_debug erase /dev/$($p.Dev) 0 $($p.Size)")
        $restoreLines.Add("  dd if=$usbFile of=/dev/$($p.Dev) bs=64k")
    }

    if ($p.Name -in @("art","mac")) {
        $restoreLines.Add("WARNING: device-specific calibration/identity partition. Restore only your own backup.")
    }
    elseif ($p.Name -match "u-boot") {
        $restoreLines.Add("WARNING: bootloader-related partition. Restoring the wrong image can brick the device.")
    }

    $restoreLines.Add("")
}

Set-Content -Path $restoreFile -Value $restoreLines -Encoding UTF8
Set-Content -Path $manifestFile -Value $manifestLines -Encoding UTF8

Write-Host ""
Write-Host "Backup complete."
Write-Host ("Restore commands written to: {0}" -f $restoreFile)
Write-Host ("Backup manifest written to: {0}" -f $manifestFile)
Write-Host ""
Write-Host "Recommended restore commands:"
Write-Host "----------------------------------------"
Get-Content $restoreFile