param(
    [string]$Distro = "Ubuntu"
)

$ErrorActionPreference = "Stop"

function Write-Info {
    param([string]$Message)
    Write-Host $Message
}

function Ensure-WslInstalled {
    if (-not (Get-Command wsl.exe -ErrorAction SilentlyContinue)) {
        Write-Warning "WSL is not installed. Install it with: wsl --install -d $Distro"
        throw "WSL is required."
    }
}

function Ensure-WslDistro {
    $distros = @()
    try {
        $distros = wsl -l -q 2>$null
    } catch {
        $distros = @()
    }

    if (-not $distros -or ($distros -notcontains $Distro)) {
        Write-Warning "WSL distro '$Distro' is not installed. Install it with: wsl --install -d $Distro"
        throw "WSL distro not found."
    }
}

function Ensure-WslConfig {
    $configPath = Join-Path $env:USERPROFILE ".wslconfig"
    $lines = [System.Collections.Generic.List[string]]::new()

    if (Test-Path $configPath) {
        $content = Get-Content -Path $configPath -ErrorAction SilentlyContinue
        foreach ($line in $content) {
            $null = $lines.Add([string]$line)
        }
    }

    $wsl2Index = $lines.FindIndex({ param($l) $l -match '^\s*\[wsl2\]\s*$' })

    if ($wsl2Index -lt 0) {
        if ($lines.Count -gt 0 -and $lines[$lines.Count - 1] -ne "") {
            $lines.Add("")
        }
        $lines.Add("[wsl2]")
        $lines.Add("networkingMode=mirrored")
        $lines.Add("guiApplications=true")
    } else {
        $blockEnd = $lines.Count
        for ($i = $wsl2Index + 1; $i -lt $lines.Count; $i++) {
            if ($lines[$i] -match '^\s*\[.+\]\s*$') {
                $blockEnd = $i
                break
            }
        }

        $netIndex = -1
        $guiIndex = -1
        for ($i = $wsl2Index + 1; $i -lt $blockEnd; $i++) {
            if ($lines[$i] -match '^\s*networkingMode=') { $netIndex = $i }
            if ($lines[$i] -match '^\s*guiApplications=') { $guiIndex = $i }
        }

        if ($netIndex -ge 0) {
            $lines[$netIndex] = "networkingMode=mirrored"
        } else {
            $lines.Insert($blockEnd, "networkingMode=mirrored")
            $blockEnd++
        }

        if ($guiIndex -ge 0) {
            $lines[$guiIndex] = "guiApplications=true"
        } else {
            $lines.Insert($blockEnd, "guiApplications=true")
        }
    }

    Set-Content -Path $configPath -Value $lines -Encoding ASCII
}

function New-StartMenuShortcut {
    param(
        [string]$ShortcutName,
        [string]$ScriptPath,
        [string]$WorkingDirectory
    )

    $startMenuDir = Join-Path $env:APPDATA "Microsoft\Windows\Start Menu\Programs"
    $shortcutPath = Join-Path $startMenuDir "$ShortcutName.lnk"
    $iconPath = Join-Path $env:SystemRoot "System32\imageres.dll"

    $shell = New-Object -ComObject WScript.Shell
    $shortcut = $shell.CreateShortcut($shortcutPath)
    $shortcut.TargetPath = (Get-Command powershell.exe).Source
    $shortcut.Arguments = "-NoProfile -ExecutionPolicy Bypass -File `"$ScriptPath`""
    $shortcut.WorkingDirectory = $WorkingDirectory
    $shortcut.IconLocation = "$iconPath,15"
    $shortcut.WindowStyle = 1
    $shortcut.Save()
}

function Test-IsAdmin {
    $current = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = New-Object Security.Principal.WindowsPrincipal($current)
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Add-FirewallRules {
    $rulePrefix = "ScreenShare AirPlay"
    New-NetFirewallRule -DisplayName "$rulePrefix TCP" -Direction Inbound -Action Allow -Protocol TCP -LocalPort 35000-35002 | Out-Null
    New-NetFirewallRule -DisplayName "$rulePrefix UDP" -Direction Inbound -Action Allow -Protocol UDP -LocalPort 35000-35002 | Out-Null
    New-NetFirewallRule -DisplayName "$rulePrefix mDNS" -Direction Inbound -Action Allow -Protocol UDP -LocalPort 5353 | Out-Null
}

function Write-Config {
    param(
        [string]$ConfigPath,
        [string]$DefaultName
    )

    $lines = @()
    if (-not [string]::IsNullOrWhiteSpace($DefaultName)) {
        $safeName = $DefaultName.Replace('"', '`"')
        $lines += "`$DefaultName = `"$safeName`""
    }

    Set-Content -Path $ConfigPath -Value $lines -Encoding ASCII
}

Ensure-WslInstalled
Ensure-WslDistro

Write-Info "Updating .wslconfig to enable mirrored networking and GUI apps..."
Ensure-WslConfig

$repoRoot = Resolve-Path (Join-Path $PSScriptRoot "..")
$wslScriptWindows = Resolve-Path (Join-Path $repoRoot "wsl\setup-uxplay.sh")
$wslScriptWindowsPath = $wslScriptWindows.Path
$wslScriptWindowsForWsl = $wslScriptWindowsPath -replace "\\", "/"

$wslScriptLinux = (wsl -d $Distro -- wslpath -a "$wslScriptWindowsForWsl" 2>$null)
if (-not $wslScriptLinux) {
    throw "Failed to resolve WSL path for setup script: $wslScriptWindowsPath"
}
$wslScriptLinux = $wslScriptLinux.Trim()

Write-Info "Running WSL setup script..."
$bashCmd = "chmod +x '$wslScriptLinux'; '$wslScriptLinux'"
& wsl -d $Distro -- bash -lc $bashCmd

Write-Info "Restarting WSL to apply systemd and networking changes..."
& wsl --shutdown

Write-Info "Attempting to enable Avahi after restart (may prompt for sudo password)..."
$enableCmd = "if command -v systemctl >/dev/null 2>&1; then sudo systemctl enable --now avahi-daemon; else echo 'systemctl not available'; fi"
& wsl -d $Distro -- bash -lc $enableCmd

$startScriptPath = Resolve-Path (Join-Path $repoRoot "scripts\start-airplay.ps1")
$startShortcutName = "ScreenShare"
$prompt = "Create Start Menu shortcut for $startShortcutName? (Y/N) [Y]"
$response = Read-Host $prompt
if ([string]::IsNullOrWhiteSpace($response) -or $response -match '^[Yy]') {
    New-StartMenuShortcut -ShortcutName $startShortcutName -ScriptPath $startScriptPath.Path -WorkingDirectory $repoRoot.Path
    Write-Info "Start Menu shortcut created: $startShortcutName"
} else {
    Write-Info "Skipped Start Menu shortcut."
}

$firewallPrompt = "Add Windows Firewall rules for ports 35000-35002 and 5353? (Y/N) [Y]"
$firewallResponse = Read-Host $firewallPrompt
if ([string]::IsNullOrWhiteSpace($firewallResponse) -or $firewallResponse -match '^[Yy]') {
    if (Test-IsAdmin) {
        Add-FirewallRules
        Write-Info "Firewall rules added."
    } else {
        Write-Info "Skipping firewall rules (requires an elevated PowerShell)."
    }
} else {
    Write-Info "Skipped firewall rules."
}

$configPath = Join-Path $repoRoot "scripts\config.ps1"
$namePrompt = "Set default AirPlay name? (Y/N) [Y]"
$nameResponse = Read-Host $namePrompt
if ([string]::IsNullOrWhiteSpace($nameResponse) -or $nameResponse -match '^[Yy]') {
    $defaultName = Read-Host "Default name (blank keeps ScreenShare)"
    if ([string]::IsNullOrWhiteSpace($defaultName)) {
        $defaultName = "ScreenShare"
    }
    Write-Config -ConfigPath $configPath -DefaultName $defaultName
    Write-Info "Default name set to: $defaultName"
} else {
    Write-Info "Skipped default name configuration."
}

Write-Info "Setup complete. You can now run scripts\\start-airplay.ps1"