$sshConfigPath = Join-Path $env:USERPROFILE ".ssh\config"

# ==================== Core Internal Functions ====================

function Get-SshConfigEntries {
    [OutputType([System.Collections.Generic.List[PSCustomObject]])]
    $entries = [System.Collections.Generic.List[PSCustomObject]]::new()
    $current = $null

    if (-not (Test-Path $sshConfigPath)) { return ,$entries }

    Get-Content $sshConfigPath -Encoding UTF8 | ForEach-Object {
        # 剔除不可见字符、BOM、零宽字符，清除注释和首尾空白
        $line = ($_ -replace '[\x00-\x1F\x7F-\x9F\u200B-\u200D\uFEFF]', '').Trim()
        $line = ($line -replace '#.*$', '').Trim()

        if ([string]::IsNullOrWhiteSpace($line) -or $line -notmatch '^[a-zA-Z]') { return }

        if ($line -match '^Host\s+(.+)$') {
            $hostAlias = $matches[1].Trim()
            if ([string]::IsNullOrWhiteSpace($hostAlias)) { return }

            if ($null -ne $current) { $entries.Add($current) }
            $current = [PSCustomObject]@{
                Host                = $hostAlias
                HostName            = $null
                Port                = $null
                User                = $null
                IdentityFile        = $null
                ServerAliveInterval = $null
                ServerAliveCountMax = $null
                TCPKeepAlive        = $null
            }
        }
        elseif ($null -ne $current) {
            if     ($line -match '^HostName\s+(\S+)')           { $current.HostName            = $matches[1].Trim() }
            elseif ($line -match '^Port\s+(\S+)')               { $current.Port                = $matches[1].Trim() }
            elseif ($line -match '^User\s+(\S+)')               { $current.User                = $matches[1].Trim() }
            elseif ($line -match '^IdentityFile\s+(\S+)')       { $current.IdentityFile        = $matches[1].Trim() }
            elseif ($line -match '^ServerAliveInterval\s+(\S+)') { $current.ServerAliveInterval = $matches[1].Trim() }
            elseif ($line -match '^ServerAliveCountMax\s+(\S+)') { $current.ServerAliveCountMax = $matches[1].Trim() }
            elseif ($line -match '^TCPKeepAlive\s+(\S+)')       { $current.TCPKeepAlive        = $matches[1].Trim() }
        }
    }
    if ($null -ne $current) { $entries.Add($current) }

    return ,$entries
}

function Save-SshConfigEntries {
    param(
        [Parameter(Mandatory = $true)]
        [System.Collections.Generic.List[PSCustomObject]]$Entries
    )

    $lines = [System.Collections.Generic.List[string]]::new()

    foreach ($entry in $Entries) {
        if ([string]::IsNullOrWhiteSpace($entry.Host)) { continue }

        $lines.Add("Host $($entry.Host.Trim())")
        if (-not [string]::IsNullOrWhiteSpace($entry.HostName))            { $lines.Add("    HostName $($entry.HostName.Trim())") }
        if (-not [string]::IsNullOrWhiteSpace($entry.Port))                { $lines.Add("    Port $($entry.Port.Trim())") }
        if (-not [string]::IsNullOrWhiteSpace($entry.User))                { $lines.Add("    User $($entry.User.Trim())") }
        if (-not [string]::IsNullOrWhiteSpace($entry.IdentityFile))        { $lines.Add("    IdentityFile $($entry.IdentityFile.Trim())") }
        if (-not [string]::IsNullOrWhiteSpace($entry.ServerAliveInterval)) { $lines.Add("    ServerAliveInterval $($entry.ServerAliveInterval.Trim())") }
        if (-not [string]::IsNullOrWhiteSpace($entry.ServerAliveCountMax)) { $lines.Add("    ServerAliveCountMax $($entry.ServerAliveCountMax.Trim())") }
        if (-not [string]::IsNullOrWhiteSpace($entry.TCPKeepAlive))        { $lines.Add("    TCPKeepAlive $($entry.TCPKeepAlive.Trim())") }
        $lines.Add("")
    }

    $dir = Split-Path $sshConfigPath -Parent
    if (-not (Test-Path $dir)) {
        New-Item -Path $dir -ItemType Directory -Force | Out-Null
    }

    if (Test-Path $sshConfigPath) {
        $backupPath = "$sshConfigPath.$(Get-Date -Format 'yyyyMMdd_HHmmss').bak"
        Copy-Item $sshConfigPath $backupPath -Force
        Write-Host "Backed up original config → $backupPath" -ForegroundColor Yellow

        $oldBaks = Get-ChildItem -Path $dir -Filter "config.*.bak" |
                   Sort-Object LastWriteTime -Descending |
                   Select-Object -Skip 5
        if ($oldBaks) {
            $oldBaks | Remove-Item -Force
            Write-Host "Cleaned up $($oldBaks.Count) old backup(s)." -ForegroundColor DarkGray
        }
    }

    if ($PSVersionTable.PSVersion.Major -ge 6) {
        $lines | Set-Content $sshConfigPath -Encoding utf8NoBOM
    } else {
        [System.IO.File]::WriteAllLines(
            $sshConfigPath,
            [string[]]$lines,
            (New-Object System.Text.UTF8Encoding($false))
        )
    }

    Write-Host "✅ SSH config saved successfully!" -ForegroundColor Green
}

# 内部连接字符串解析器
function Invoke-ParseConnectionString {
    param([string]$ConnectionString)

    $parsed = @{ User = $null; HostName = $null; Port = $null }
    if ([string]::IsNullOrWhiteSpace($ConnectionString)) { return $parsed }

    if ($ConnectionString -match '^(?:([^@]+)@)?([^:]+)(?::(\d+))?$') {
        if ($matches[1]) { $parsed.User     = $matches[1].Trim() }
        if ($matches[2]) { $parsed.HostName = $matches[2].Trim() }
        if ($matches[3]) { $parsed.Port     = $matches[3].Trim() }
    }
    return $parsed
}

# 查找 entry 的内部辅助函数
function Find-SshEntry {
    param(
        [System.Collections.Generic.List[PSCustomObject]]$Entries,
        [string]$Name
    )
    return $Entries | Where-Object {
        $_.Host -eq $Name -or ($_.Host -split '\s+') -contains $Name
    } | Select-Object -First 1
}

# ==================== Part 1. List SSH Hosts (Shortcut: sg) ====================
function Get-SshHost {
    [CmdletBinding()]
    param(
        [Parameter(Position = 0)]
        [string]$Name,
        [Alias('h', '--help')][switch]$Help
    )

    if ($Help) { Get-Help $MyInvocation.MyCommand.Name -Detailed; return }

    $entries = Get-SshConfigEntries
    if ($entries.Count -eq 0) {
        Write-Host "No Host entries found in ~/.ssh/config." -ForegroundColor Yellow; return
    }

    $filtered = if ([string]::IsNullOrWhiteSpace($Name)) {
        $entries
    } else {
        $entries | Where-Object { $_.Host -like "*$Name*" -or $_.HostName -like "*$Name*" }
    }

    if (-not $filtered -or @($filtered).Count -eq 0) {
        Write-Host "No hosts matched '$Name'." -ForegroundColor Yellow; return
    }

    Write-Host "✅ Current SSH Config entries:" -ForegroundColor Green
    @($filtered) | Sort-Object Host |
        Format-Table -AutoSize `
            @{Label='Host (Alias)'; Expression={$_.Host}}, `
            @{Label='Hostname / IP'; Expression={$_.HostName}}, `
            @{Label='Port';          Expression={$_.Port}}, `
            @{Label='User';          Expression={$_.User}}, `
            @{Label='IdentityFile';  Expression={$_.IdentityFile}}
}

# ==================== Part 2. Add / Update Host (Shortcut: sa) ====================
function Add-SshHost {
    [CmdletBinding()]
    param(
        [Parameter(Position = 0)]
        [string]$Name,

        [Parameter(Position = 1)]
        [string]$ConnectionString,

        [Alias('p')][string]$Port,
        [string]$User,
        [string]$IdentityFile,
        [Alias('h', '--help')][switch]$Help
    )

    if ($Help) { Get-Help $MyInvocation.MyCommand.Name -Detailed; return }

    if ([string]::IsNullOrWhiteSpace($Name)) {
        Write-Error "Missing required parameter: -Name (alias for the host)." -Category InvalidArgument
        return
    }

    $conn          = Invoke-ParseConnectionString $ConnectionString
    $finalHostName = $conn.HostName
    $finalUser     = if ($PSBoundParameters.ContainsKey('User')) { $User } else { $conn.User }
    $finalPort     = if ($PSBoundParameters.ContainsKey('Port')) { $Port } else { $conn.Port }

    if ([string]::IsNullOrWhiteSpace($finalHostName)) {
        Write-Error "Missing target host! Provide an IP/Domain or 'user@host:port' as the second parameter." -Category InvalidArgument
        Write-Host "💡 Usage: sa <Alias> <User@IP:Port>" -ForegroundColor Cyan
        return
    }

    $entries = Get-SshConfigEntries
    $existing = Find-SshEntry $entries $Name

    if ($existing -and ($existing.Host -split '\s+').Count -gt 1) {
        Write-Error "❌ '$Name' belongs to a multi-alias block '$($existing.Host)'. Use 'sc' to edit it manually." -Category InvalidOperation
        return
    }

    if ($existing) {
        Write-Host "⚠️  Host '$Name' found, updating..." -ForegroundColor Yellow
        $existing.HostName = $finalHostName
        if ($PSBoundParameters.ContainsKey('Port')         -or $finalPort) { $existing.Port         = $finalPort }
        if ($PSBoundParameters.ContainsKey('User')         -or $finalUser) { $existing.User         = $finalUser }
        if ($PSBoundParameters.ContainsKey('IdentityFile'))                 { $existing.IdentityFile = $IdentityFile }

        # 为已有主机自动补全 KeepAlive 参数
        if (-not $existing.ServerAliveInterval) { $existing.ServerAliveInterval = "30" }
        if (-not $existing.ServerAliveCountMax) { $existing.ServerAliveCountMax = "6" }
        if (-not $existing.TCPKeepAlive)        { $existing.TCPKeepAlive        = "yes" }
    } else {
        $newEntry = [PSCustomObject]@{
            Host                = $Name.Trim()
            HostName            = $finalHostName
            Port                = $finalPort
            User                = $finalUser
            IdentityFile        = if ($PSBoundParameters.ContainsKey('IdentityFile')) { $IdentityFile } else { $null }
            # 自动添加防断连心跳配置
            ServerAliveInterval = "30"
            ServerAliveCountMax = "6"
            TCPKeepAlive        = "yes"
        }
        $entries.Add($newEntry)
        Write-Host "✅ Adding new Host '$Name' (with KeepAlive options)..." -ForegroundColor Green
    }

    Save-SshConfigEntries $entries
}

# ==================== Part 3. Edit Host Properties (Shortcut: se) ====================
function Edit-SshHost {
    [CmdletBinding()]
    param(
        [Parameter(Position = 0, Mandatory = $true)]
        [string]$Name,

        [Parameter(Position = 1)]
        [string]$ConnectionString,

        [Alias('p')][string]$Port,
        [string]$User,
        [string]$IdentityFile,
        [Alias('h', '--help')][switch]$Help
    )

    if ($Help) { Get-Help $MyInvocation.MyCommand.Name -Detailed; return }

    $conn = Invoke-ParseConnectionString $ConnectionString

    $hasChanges = $conn.HostName -or $conn.User -or $conn.Port -or
                  $PSBoundParameters.ContainsKey('Port') -or
                  $PSBoundParameters.ContainsKey('User') -or
                  $PSBoundParameters.ContainsKey('IdentityFile')

    if (-not $hasChanges) {
        Write-Error "You must specify at least one property or a connection string to modify." -Category InvalidArgument
        return
    }

    $entries  = Get-SshConfigEntries
    $existing = Find-SshEntry $entries $Name

    if (-not $existing) {
        Write-Error "❌ Host '$Name' not found in config." -Category ObjectNotFound
        return
    }

    if (($existing.Host -split '\s+').Count -gt 1) {
        Write-Error "❌ '$Name' belongs to a multi-alias block '$($existing.Host)'. Use 'sc' to edit it manually." -Category InvalidOperation
        return
    }

    if ($conn.HostName)                                 { $existing.HostName     = $conn.HostName }
    if ($PSBoundParameters.ContainsKey('User'))        { $existing.User         = $User }
    elseif ($conn.User)                                 { $existing.User         = $conn.User }
    if ($PSBoundParameters.ContainsKey('Port'))        { $existing.Port         = $Port }
    elseif ($conn.Port)                                 { $existing.Port         = $conn.Port }
    if ($PSBoundParameters.ContainsKey('IdentityFile')) { $existing.IdentityFile = $IdentityFile }

    Write-Host "⚠️  Updating Host '$Name' ($($existing.Host))..." -ForegroundColor Yellow
    Save-SshConfigEntries $entries
}

# ==================== Part 4. Remove Host (Shortcut: sr) ====================
function Remove-SshHost {
    [CmdletBinding(SupportsShouldProcess = $true, ConfirmImpact = 'Medium')]
    param(
        [Parameter(Position = 0)]
        [string]$Name,
        [Alias('h', '--help')][switch]$Help
    )

    if ($Help) { Get-Help $MyInvocation.MyCommand.Name -Detailed; return }

    if ([string]::IsNullOrWhiteSpace($Name)) {
        Write-Error "Missing required parameter: -Name" -Category InvalidArgument; return
    }

    $entries  = Get-SshConfigEntries
    $toRemove = Find-SshEntry $entries $Name

    if (-not $toRemove) {
        Write-Host "❌ No Host named '$Name' found." -ForegroundColor Red; return
    }

    if ($PSCmdlet.ShouldProcess($Name, "Remove block '$($toRemove.Host)' from SSH config")) {
        $targetHost = $toRemove.Host
        $newList = [System.Collections.Generic.List[PSCustomObject]]::new()
        foreach ($e in $entries) {
            if ($e.Host -ne $targetHost) { $newList.Add($e) }
        }
        Save-SshConfigEntries $newList
        Write-Host "✅ Host block '$($toRemove.Host)' removed." -ForegroundColor Green
    }
}

# ==================== Part 5. Open Config File (Shortcut: sc) ====================
function Open-SshConfig {
    if (-not (Test-Path $sshConfigPath)) {
        $dir = Split-Path $sshConfigPath -Parent
        if (-not (Test-Path $dir)) { New-Item -Path $dir -ItemType Directory -Force | Out-Null }
        New-Item -Path $sshConfigPath -ItemType File -Force | Out-Null
    }
    Write-Host "📝 Opening SSH config..." -ForegroundColor Yellow

    $editor = $null
    foreach ($candidate in @('code', 'notepad++', 'notepad')) {
        if (Get-Command $candidate -ErrorAction SilentlyContinue) {
            $editor = $candidate; break
        }
    }

    if ($editor -eq 'code') {
        & code --wait $sshConfigPath
    } else {
        Start-Process $editor -ArgumentList "`"$sshConfigPath`""
    }
}

# ==================== Part 6. SSH Copy ID Tool (Shortcut: ssh-copy-id) ====================
function ssh-copy-id {
    [CmdletBinding()]
    param(
        [Parameter(Position = 0)]
        [string]$Destination,

        [Alias('i', '--identity-file')][string]$IdentityFile,
        [Alias('p', '--port')][string]$Port,
        [Alias('o', '--ssh-option')][string[]]$SshOption = @(),
        [Alias('f', '--force')][switch]$Force,
        [Alias('h', '--help', '?')][switch]$Help
    )

    if ($Help -or [string]::IsNullOrWhiteSpace($Destination)) {
        Get-Help $MyInvocation.MyCommand.Name -Detailed; return
    }

    if ($Destination -notmatch '^([^@]+)@([^:]+)(?::(\d+))?$') {
        Write-Error "Invalid format! Use user@host or user@host:port." -Category InvalidArgument
        return
    }

    $finalUser = $matches[1].Trim()
    $finalHost = $matches[2].Trim()
    $finalPort = if ($PSBoundParameters.ContainsKey('Port') -and $Port) {
        $Port
    } elseif ($matches[3]) {
        $matches[3].Trim()
    } else {
        "22"
    }

    if (-not $IdentityFile) {
        foreach ($key in @(
            "$env:USERPROFILE\.ssh\id_ed25519.pub",
            "$env:USERPROFILE\.ssh\id_rsa.pub",
            "$env:USERPROFILE\.ssh\id_ecdsa.pub"
        )) {
            if (Test-Path $key) { $IdentityFile = $key; break }
        }
    }

    if (-not $IdentityFile -or -not (Test-Path $IdentityFile)) {
        Write-Error "Public key not found! Run 'ssh-keygen' first or specify with '-i'." -Category ObjectNotFound
        return
    }

    $pubKey     = (Get-Content -Path $IdentityFile -Raw -ErrorAction Stop).Trim()
    $forceFlag  = if ($Force) { "1" } else { "0" }

    $remoteBashScript = @'
set -e
FORCE_FLAG=__FORCE__
PUB_KEY='__PUBKEY__'
mkdir -p ~/.ssh && chmod 700 ~/.ssh
if [ "$FORCE_FLAG" = "0" ] && grep -qF "$PUB_KEY" ~/.ssh/authorized_keys 2>/dev/null; then
    echo "Key already exists, skipping."
    exit 0
fi
echo "$PUB_KEY" >> ~/.ssh/authorized_keys
chmod 600 ~/.ssh/authorized_keys
echo "Public key added to ~/.ssh/authorized_keys"
'@

    $escapedPubKey = $pubKey -replace "'", "'\'''"
    $remoteBashScript = $remoteBashScript `
        -replace '__FORCE__', $forceFlag `
        -replace '__PUBKEY__', $escapedPubKey

    $remoteBashScript = $remoteBashScript -replace "`r", ""
    $base64Script = [Convert]::ToBase64String(
        [System.Text.Encoding]::UTF8.GetBytes($remoteBashScript)
    )

    $sshCmd = "echo '$base64Script' | { base64 -d 2>/dev/null || base64 -D 2>/dev/null || openssl enc -d -base64 2>/dev/null; } | bash"

    Write-Host "🚀 Connecting to ${finalUser}@${finalHost}:${finalPort}..." -ForegroundColor Green

    $sshArgs = @()
    if ($finalPort -ne "22") { $sshArgs += "-p", $finalPort }
    foreach ($opt in $SshOption) { $sshArgs += $opt }

    & ssh @sshArgs "${finalUser}@${finalHost}" $sshCmd

    if ($LASTEXITCODE -eq 0) {
        $portSuffix = if ($finalPort -ne '22') { " -p $finalPort" } else { "" }
        Write-Host "🎉 Done! Login with: ssh ${finalUser}@${finalHost}${portSuffix}" -ForegroundColor Green
    } else {
        Write-Error "Failed. Check network, credentials, or SSH server config." -Category ConnectionError
    }
}

# ==================== Shortcuts & Global Registration ====================
New-Alias -Name sg -Value Get-SshHost    -Force -Scope Global
New-Alias -Name sa -Value Add-SshHost    -Force -Scope Global
New-Alias -Name se -Value Edit-SshHost   -Force -Scope Global
New-Alias -Name sr -Value Remove-SshHost -Force -Scope Global
New-Alias -Name sc -Value Open-SshConfig -Force -Scope Global

Write-Host "🎉 SSH Suite loaded!" -ForegroundColor Cyan
Write-Host "  sg=List  sa=Add/Update  se=Edit  sr=Remove  sc=Open  ssh-copy-id=CopyKey" -ForegroundColor White