# ====================== SSH Host Management Tools (English Version) ======================
# Place this in your $PROFILE together with ssh-copy-id

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

# ==================== Internal Functions (no need to call manually) ====================
function Get-SshConfigEntries {
    $entries = @()
    $current = $null

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

    Get-Content $sshConfigPath -Encoding UTF8 | ForEach-Object {
        $line = ($_ -replace '#.*$', '').Trim()
        if ([string]::IsNullOrWhiteSpace($line)) { return }

        if ($line -match '^Host\s+(.+)$') {
            if ($current) { $entries += $current }
            $current = [PSCustomObject]@{
                Host         = $matches[1].Trim()
                HostName     = $null
                Port         = $null
                User         = $null
                IdentityFile = $null
            }
        }
        elseif ($current) {
            switch -Regex ($line) {
                '^HostName\s+(.+)$'     { $current.HostName     = $matches[1].Trim() }
                '^Port\s+(.+)$'         { $current.Port         = $matches[1].Trim() }
                '^User\s+(.+)$'         { $current.User         = $matches[1].Trim() }
                '^IdentityFile\s+(.+)$' { $current.IdentityFile = $matches[1].Trim() }
            }
        }
    }
    if ($current) { $entries += $current }
    return $entries
}

function Save-SshConfigEntries {
    param([Parameter(Mandatory = $true)] $Entries)

    $content = @()
    foreach ($entry in $Entries) {
        $content += "Host $($entry.Host)"
        if ($entry.HostName)     { $content += "    HostName $($entry.HostName)" }
        if ($entry.Port)         { $content += "    Port $($entry.Port)" }
        if ($entry.User)         { $content += "    User $($entry.User)" }
        if ($entry.IdentityFile) { $content += "    IdentityFile $($entry.IdentityFile)" }
        $content += ""  # blank line separator
    }

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

    # Timestamped backup for safety
    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
    }

    $content | Set-Content $sshConfigPath -Encoding UTF8
    Write-Host "✅ SSH config file has been successfully saved!" -ForegroundColor Green
}

# ==================== 1. List SSH Hosts ====================
<#
.SYNOPSIS
    List all hosts defined in ~/.ssh/config (shortcut: sg)
.DESCRIPTION
    Supports keyword search. Run without parameters to list everything.
.EXAMPLE
    sg
    sg web
    sg --help
#>
function Get-SshHost {
    [CmdletBinding()]
    param(
        [Parameter(Position = 0)]
        [string]$Name,           # keyword for fuzzy search

        [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 = $entries
    if (-not [string]::IsNullOrWhiteSpace($Name)) {
        $filtered = $entries | Where-Object { $_.Host -like "*$Name*" -or $_.HostName -like "*$Name*" }
    }

    if ($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}}
}

# ==================== 2. Add / Update Host ====================
<#
.SYNOPSIS
    Add or update a Host entry in ~/.ssh/config (shortcut: sa)
.EXAMPLE
    sa myserver 192.168.1.100 -User ubuntu -Port 2222
    sa web -Destination ubuntu@web.example.com
    sa --help
#>
function Add-SshHost {
    [CmdletBinding()]
    param(
        [Parameter(Position = 0)]
        [string]$Name,                    # Alias name (Host)

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

        [string]$Destination,             # supports user@host format

        [string]$Port,
        [string]$User,
        [string]$IdentityFile,

        [Alias('h', '--help')]
        [switch]$Help
    )

    # Show help if --help is used or no Name is provided (this prevents mandatory prompt)
    if ($Help -or [string]::IsNullOrWhiteSpace($Name)) {
        Get-Help $MyInvocation.MyCommand.Name -Detailed
        return
    }

    # Support -Destination user@host
    if ($Destination -and $Destination -match '^(.+)@(.+)$') {
        if (-not $User)      { $User = $matches[1] }
        if (-not $HostName)  { $HostName = $matches[2] }
    }

    if (-not $HostName) {
        Write-Error "You must provide a hostname/IP! Use -HostName or -Destination user@host"
        return
    }

    $entries = Get-SshConfigEntries
    $existing = $entries | Where-Object { $_.Host -eq $Name }

    if ($existing) {
        Write-Host "⚠️  Host '$Name' already exists, updating..." -ForegroundColor Yellow
        $existing.HostName     = $HostName
        if ($PSBoundParameters.ContainsKey('Port'))         { $existing.Port = $Port }
        if ($PSBoundParameters.ContainsKey('User'))         { $existing.User = $User }
        if ($PSBoundParameters.ContainsKey('IdentityFile')) { $existing.IdentityFile = $IdentityFile }
    } else {
        $newEntry = [PSCustomObject]@{
            Host         = $Name.Trim()
            HostName     = $HostName.Trim()
            Port         = $Port
            User         = $User
            IdentityFile = $IdentityFile
        }
        $entries += $newEntry
        Write-Host "✅ Added new Host '$Name'" -ForegroundColor Green
    }

    Save-SshConfigEntries $entries
}

# ==================== 3. Remove Host ====================
<#
.SYNOPSIS
    Remove a Host entry from ~/.ssh/config (shortcut: sr)
.EXAMPLE
    sr myserver
    sr myserver -Confirm
#>
function Remove-SshHost {
    [CmdletBinding(SupportsShouldProcess = $true, ConfirmImpact = 'Medium')]
    param(
        [Parameter(Position = 0)]
        [string]$Name,

        [Alias('h', '--help')]
        [switch]$Help
    )

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

    $entries = Get-SshConfigEntries
    $toRemove = $entries | Where-Object { $_.Host -eq $Name }

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

    if ($PSCmdlet.ShouldProcess($Name, "Remove from SSH config")) {
        $entries = $entries | Where-Object { $_.Host -ne $Name }
        Save-SshConfigEntries $entries
        Write-Host "✅ Host '$Name' has been successfully removed" -ForegroundColor Green
    }
}

# ==================== Shortcuts ====================
New-Alias -Name sg -Value Get-SshHost -Force -Scope Global
New-Alias -Name sa -Value Add-SshHost -Force -Scope Global
New-Alias -Name sr -Value Remove-SshHost -Force -Scope Global

Write-Host "🎉 SSH Host Management Tools loaded successfully!" -ForegroundColor Cyan
Write-Host "Available shortcuts:" -NoNewline
Write-Host "  sg → List hosts" -ForegroundColor White
Write-Host "  sa → Add/Update host" -ForegroundColor White
Write-Host "  sr → Remove host" -ForegroundColor White
Write-Host "Examples:" -ForegroundColor Gray
Write-Host "   sa web 192.168.1.100 -User ubuntu" -ForegroundColor Gray
Write-Host "   sg" -ForegroundColor Gray