param(
    [string]$usbLetter,
    [string]$srcPath = ".\src",
    [string]$distPath = ".\dist\main.exe"
)

function Write-Header {
    param([string]$title)
    $width = 80
    $padding = [math]::Floor(($width - $title.Length - 2) / 2)
    Write-Host ""
    Write-Host ("=" * $width) -ForegroundColor Cyan
    Write-Host ("=" + (" " * $padding) + $title + (" " * $padding) + "=") -ForegroundColor Cyan
    Write-Host ("=" * $width) -ForegroundColor Cyan
    Write-Host ""
}

function Write-Step {
    param([string]$message)
    Write-Host "[*] " -NoNewline -ForegroundColor Green
    Write-Host $message -ForegroundColor White
}

function Write-ErrorMsg {
    param([string]$message)
    Write-Host "[X] " -NoNewline -ForegroundColor Red
    Write-Host $message -ForegroundColor Red
}

function Write-Success {
    param([string]$message)
    Write-Host "[+] " -NoNewline -ForegroundColor Green
    Write-Host $message -ForegroundColor Green
}

function Write-Warning {
    param([string]$message)
    Write-Host "[!] " -NoNewline -ForegroundColor Yellow
    Write-Host $message -ForegroundColor Yellow
}

function Write-ProgressBar {
    param(
        [int]$current,
        [int]$total,
        [string]$activity
    )
    $percent = [math]::Min(100, [math]::Floor(($current / $total) * 100))
    $barWidth = 50
    $completed = [math]::Floor(($percent / 100) * $barWidth)
    $remaining = $barWidth - $completed
    
    $completedBar = "#" * $completed
    $remainingBar = "-" * $remaining
    $bar = "[" + $completedBar + $remainingBar + "]"
    
    Write-Host ("`r$activity $bar $percent% ($current/$total)") -NoNewline -ForegroundColor Cyan
    
    if ($current -eq $total) {
        Write-Host ""
    }
}

function Test-Prerequisites {
    $errors = @()
    
    # Check if running as admin (for autorun.inf)
    $isAdmin = ([Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
    if (-not $isAdmin) {
        Write-Warning "Not running as administrator. Autorun.inf may not work properly."
    }
    
    # Validate USB letter format
    if ($usbLetter -notmatch '^[A-Z]:?$') {
        $errors += "Invalid USB letter format. Use format like 'E:' or 'E'"
    }
    
    # Normalize USB letter
    $script:usbLetter = $usbLetter.TrimEnd(':') + ':'
    
    # Check if paths exist
    if (-not (Test-Path $srcPath)) {
        $errors += "Source path not found: $srcPath"
    }
    
    if (-not (Test-Path $distPath)) {
        $errors += "Distribution executable not found: $distPath"
    }
    
    # Check USB drive
    $usbRoot = "$($script:usbLetter)\"
    if (-not (Test-Path $usbRoot)) {
        $errors += "USB drive not found: $usbRoot"
    } else {
        # Check if USB has enough space
        try {
            $drive = Get-PSDrive -Name $script:usbLetter.TrimEnd(':') -ErrorAction Stop
            $freeSpaceMB = [math]::Round($drive.Free / 1MB, 2)
            
            # Estimate required space
            $srcSize = (Get-ChildItem -Path $srcPath -Recurse -ErrorAction SilentlyContinue | 
                       Where-Object { $_.FullName -notmatch "__pycache__" } | 
                       Measure-Object -Property Length -Sum).Sum
            $distSize = (Get-Item $distPath).Length
            $requiredMB = [math]::Round(($srcSize + $distSize) / 1MB * 1.1, 2)
            
            if ($freeSpaceMB -lt $requiredMB) {
                $errors += "Insufficient space on USB. Required: ~$requiredMB MB, Available: $freeSpaceMB MB"
            }
        } catch {
            Write-Warning "Could not check USB free space: $_"
        }
    }
    
    return $errors
}

Clear-Host
Write-Header "wxIRC USB Writer"

# Show usage if no USB letter provided
if (-not $usbLetter) {
    Write-Host "Usage: .\write.ps1 -usbLetter E:" -ForegroundColor Yellow
    Write-Host ""
    Write-Host "Parameters:" -ForegroundColor Cyan
    Write-Host "  -usbLetter  : Target USB drive letter (required)" -ForegroundColor White
    Write-Host "  -srcPath    : Source directory (default: .\src)" -ForegroundColor White
    Write-Host "  -distPath   : Executable path (default: .\dist\main.exe)" -ForegroundColor White
    Write-Host ""
    exit 1
}

Write-Step "Validating prerequisites..."
$validationErrors = Test-Prerequisites

if ($validationErrors.Count -gt 0) {
    Write-Host ""
    foreach ($err in $validationErrors) {
        Write-ErrorMsg $err
    }
    Write-Host ""
    exit 1
}

Write-Success "All prerequisites validated"

# Convert to absolute paths
try {
    $srcPath = (Resolve-Path $srcPath -ErrorAction Stop).Path
    $distPath = (Resolve-Path $distPath -ErrorAction Stop).Path
} catch {
    Write-ErrorMsg "Failed to resolve paths: $_"
    exit 1
}

$usbRoot = "$usbLetter\"

Write-Host ""
Write-Host "Configuration:" -ForegroundColor Cyan
Write-Host "  Source      : $srcPath" -ForegroundColor White
Write-Host "  Executable  : $distPath" -ForegroundColor White
Write-Host "  USB Target  : $usbRoot" -ForegroundColor White
Write-Host ""

# Confirm before proceeding
Write-Host "This will overwrite existing files on the USB drive." -ForegroundColor Yellow
$response = Read-Host "Continue? (Y/N)"
if ($response -ne 'Y' -and $response -ne 'y') {
    Write-Host "Operation cancelled." -ForegroundColor Yellow
    exit 0
}

Write-Host ""
$startTime = Get-Date

Write-Step "Preparing source directory..."
$srcDest = Join-Path $usbRoot "src"

if (Test-Path $srcDest) {
    try {
        Remove-Item -Recurse -Force $srcDest -ErrorAction Stop
    } catch {
        Write-ErrorMsg "Failed to remove existing src directory: $_"
        exit 1
    }
}

try {
    New-Item -ItemType Directory -Path $srcDest -ErrorAction Stop | Out-Null
} catch {
    Write-ErrorMsg "Failed to create src directory: $_"
    exit 1
}

Write-Step "Scanning source files..."
$allFiles = @(Get-ChildItem -Path $srcPath -Recurse -ErrorAction SilentlyContinue | Where-Object {
    $_.FullName -notmatch "__pycache__"
})

$totalFiles = $allFiles.Count
$currentFile = 0

Write-Step "Copying $totalFiles files..."

foreach ($item in $allFiles) {
    $currentFile++
    
    try {
        $relative = $item.FullName.Substring($srcPath.Length)
        $target = Join-Path $srcDest $relative
        
        if ($item.PSIsContainer) {
            if (-not (Test-Path $target)) {
                New-Item -ItemType Directory -Path $target -ErrorAction Stop | Out-Null
            }
        } else {
            $targetDir = Split-Path -Parent $target
            if (-not (Test-Path $targetDir)) {
                New-Item -ItemType Directory -Path $targetDir -ErrorAction Stop | Out-Null
            }
            Copy-Item $item.FullName -Destination $target -Force -ErrorAction Stop
        }
        
        Write-ProgressBar -current $currentFile -total $totalFiles -activity "Copying files"
    } catch {
        Write-Host ""
        Write-ErrorMsg "Failed to copy $($item.FullName): $_"
    }
}

Write-Success "Source files copied successfully"

Write-Step "Copying executable..."
try {
    $exeDest = Join-Path $usbRoot "wxIRC.exe"
    $exeSize = (Get-Item $distPath).Length
    $exeSizeMB = [math]::Round($exeSize / 1MB, 2)
    
    Write-Host "  Size: $exeSizeMB MB" -ForegroundColor Gray
    
    # Copy with progress simulation for large files
    if ($exeSize -gt 5MB) {
        $buffer = 1MB
        $read = 0
        $totalChunks = [math]::Ceiling($exeSize / $buffer)
        $currentChunk = 0
        
        $sourceStream = [System.IO.File]::OpenRead($distPath)
        $destStream = [System.IO.File]::Create($exeDest)
        $bufferArray = New-Object byte[] $buffer
        
        try {
            while (($read = $sourceStream.Read($bufferArray, 0, $buffer)) -gt 0) {
                $destStream.Write($bufferArray, 0, $read)
                $currentChunk++
                Write-ProgressBar -current $currentChunk -total $totalChunks -activity "Copying executable"
            }
        } finally {
            $sourceStream.Close()
            $destStream.Close()
        }
    } else {
        # For smaller files, just copy directly
        Copy-Item $distPath -Destination $exeDest -Force -ErrorAction Stop
        Write-ProgressBar -current 1 -total 1 -activity "Copying executable"
    }
    
    Write-Success "Executable copied as wxIRC.exe"
} catch {
    Write-ErrorMsg "Failed to copy executable: $_"
    exit 1
}

Write-Step "Creating autorun.inf..."
try {
    $autorun = Join-Path $usbRoot "autorun.inf"
    @(
        "[AutoRun]"
        "open=wxIRC.exe"
        "label=wxIRC"
        "icon=src\icon.ico"
    ) | Set-Content -Path $autorun -Encoding ASCII -ErrorAction Stop
    
    # Try to set hidden attribute (may fail without admin)
    try {
        $autorunItem = Get-Item $autorun -ErrorAction Stop
        $autorunItem.Attributes = $autorunItem.Attributes -bor [System.IO.FileAttributes]::Hidden
    } catch {
        # Silently continue if can't set hidden attribute
    }
    
    Write-Success "Autorun.inf created"
} catch {
    Write-ErrorMsg "Failed to create autorun.inf: $_"
}

Write-Step "Generating write info..."
$endTime = Get-Date
$duration = ($endTime - $startTime).TotalSeconds

try {
    $sizeMB = (Get-ChildItem -Recurse -Path $usbRoot -ErrorAction SilentlyContinue | 
               Measure-Object Length -Sum).Sum / 1MB
    $speed = if ($duration -gt 0) { [math]::Round(($sizeMB / $duration), 2) } else { 0 }
    
    $infoFile = Join-Path $usbRoot "writeinfo.txt"
    @(
        "wxIRC USB Write Report"
        "=" * 50
        "Completed    : $endTime"
        "Total Size   : $([math]::Round($sizeMB, 2)) MB"
        "Duration     : $([math]::Round($duration, 2)) seconds"
        "Write Speed  : $speed MB/s"
        "Files Copied : $totalFiles"
        ""
        "USB Drive    : $usbRoot"
    ) | Set-Content -Path $infoFile -ErrorAction Stop
    
    Write-Success "Write info saved to writeinfo.txt"
} catch {
    Write-ErrorMsg "Failed to create writeinfo.txt: $_"
}

Write-Host ""
Write-Host ("=" * 80) -ForegroundColor Green
Write-Host "  USB WRITE COMPLETED SUCCESSFULLY!" -ForegroundColor Green
Write-Host ("=" * 80) -ForegroundColor Green
Write-Host ""
Write-Host "Summary:" -ForegroundColor Cyan
Write-Host "  Files Copied : $totalFiles" -ForegroundColor White
Write-Host "  Total Size   : $([math]::Round($sizeMB, 2)) MB" -ForegroundColor White
Write-Host "  Duration     : $([math]::Round($duration, 2)) seconds" -ForegroundColor White
Write-Host "  Write Speed  : $speed MB/s" -ForegroundColor White
Write-Host ""
Write-Host "The USB drive is ready to use!" -ForegroundColor Green
Write-Host ""