#!/usr/bin/env pwsh

Write-Host "=== Azure Bastion Connection Helper ===" -ForegroundColor Cyan
Write-Host ""

# Step 1: Get all VMs with their bastion-related info
Write-Host "Loading VMs and Bastions..." -ForegroundColor Yellow
$vms = az vm list -d --query "[].{Name:name, ResourceGroup:resourceGroup, PowerState:powerState, VnetId:networkProfile.networkInterfaces[0].id}" | ConvertFrom-Json

if ($vms.Count -eq 0) {
    Write-Host "No VMs found in subscription." -ForegroundColor Red
    exit 1
}

$bastions = az network bastion list --query "[].{Name:name, ResourceGroup:resourceGroup, VnetId:virtualNetwork.id}" | ConvertFrom-Json

Write-Host ""
Write-Host "=== Available VMs ===" -ForegroundColor Cyan
for ($i = 0; $i -lt $vms.Count; $i++) {
    $vm = $vms[$i]
    Write-Host "[$($i + 1)] $($vm.Name) (RG: $($vm.ResourceGroup), State: $($vm.PowerState))"
}

Write-Host ""
$vmSelection = Read-Host "Select VM number (1-$($vms.Count))"
$selectedVM = $vms[$vmSelection - 1]

if (-not $selectedVM) {
    Write-Host "Invalid selection." -ForegroundColor Red
    exit 1
}

Write-Host ""
Write-Host "Selected VM: $($selectedVM.Name)" -ForegroundColor Green
Write-Host "Resource Group: $($selectedVM.ResourceGroup)" -ForegroundColor Green
Write-Host "Power State: $($selectedVM.PowerState)" -ForegroundColor Green

# Check if VM is running
if ($selectedVM.PowerState -ne "VM running") {
    Write-Host ""
    Write-Host "WARNING: VM is not running!" -ForegroundColor Yellow
    $startVM = Read-Host "Would you like to start it? (y/n)"
    if ($startVM -eq "y") {
        Write-Host "Starting VM..." -ForegroundColor Yellow
        az vm start --resource-group $selectedVM.ResourceGroup --name $selectedVM.Name
        Write-Host "VM started successfully." -ForegroundColor Green
    }
}

# Step 2: Get VM's VNet
Write-Host ""
Write-Host "Getting VM network information..." -ForegroundColor Yellow
$vmDetails = az vm show --resource-group $selectedVM.ResourceGroup --name $selectedVM.Name | ConvertFrom-Json
$vmId = $vmDetails.id

# Get NIC details to find VNet
$nicId = $vmDetails.networkProfile.networkInterfaces[0].id
$nic = az network nic show --ids $nicId | ConvertFrom-Json
$vmVnetId = $nic.ipConfigurations[0].subnet.id -replace '/subnets/.*', ''

Write-Host "VM VNet ID: $vmVnetId" -ForegroundColor Gray

# Step 3: Find bastions
Write-Host ""
Write-Host "=== Available Bastions ===" -ForegroundColor Cyan

if ($bastions.Count -eq 0) {
    Write-Host "No Bastions found in subscription." -ForegroundColor Red
    exit 1
}

# Try to match bastion to VM's VNet
$matchedBastion = $null
for ($i = 0; $i -lt $bastions.Count; $i++) {
    $bastion = $bastions[$i]
    $isMatch = $bastion.VnetId -eq $vmVnetId
    $matchIndicator = if ($isMatch) { " (MATCHED TO VM)" } else { "" }
    
    if ($isMatch) {
        $matchedBastion = $bastion
    }
    
    Write-Host "[$($i + 1)] $($bastion.Name) (RG: $($bastion.ResourceGroup))$matchIndicator"
}

Write-Host ""
if ($matchedBastion) {
    Write-Host "Auto-selecting matched bastion: $($matchedBastion.Name)" -ForegroundColor Green
    $selectedBastion = $matchedBastion
} else {
    $bastionSelection = Read-Host "Select Bastion number (1-$($bastions.Count))"
    $selectedBastion = $bastions[$bastionSelection - 1]
    
    if (-not $selectedBastion) {
        Write-Host "Invalid selection." -ForegroundColor Red
        exit 1
    }
}

Write-Host ""
Write-Host "Selected Bastion: $($selectedBastion.Name)" -ForegroundColor Green

# Step 4: Select connection type
Write-Host ""
Write-Host "=== Connection Type ===" -ForegroundColor Cyan
Write-Host "[1] SSH (port 22)"
Write-Host "[2] RDP (port 3389)"
Write-Host "[3] Custom port"
Write-Host ""
$connType = Read-Host "Select connection type (1-3)"

switch ($connType) {
    "1" { 
        $remotePort = 22
        $localPort = 50022
        $protocol = "SSH"
    }
    "2" { 
        $remotePort = 3389
        $localPort = 53389
        $protocol = "RDP"
    }
    "3" {
        $remotePort = Read-Host "Enter remote port"
        $localPort = Read-Host "Enter local port"
        $protocol = "Custom"
    }
    default {
        Write-Host "Invalid selection." -ForegroundColor Red
        exit 1
    }
}

# Step 5: Display connection command
Write-Host ""
Write-Host "=== Connection Details ===" -ForegroundColor Cyan
Write-Host "VM: $($selectedVM.Name)" -ForegroundColor White
Write-Host "Bastion: $($selectedBastion.Name)" -ForegroundColor White
Write-Host "Protocol: $protocol" -ForegroundColor White
Write-Host "Remote Port: $remotePort" -ForegroundColor White
Write-Host "Local Port: $localPort" -ForegroundColor White
Write-Host ""

$tunnelCommand = "az network bastion tunnel ``
  --name $($selectedBastion.Name) ``
  --resource-group $($selectedBastion.ResourceGroup) ``
  --target-resource-id $vmId ``
  --resource-port $remotePort ``
  --port $localPort"

Write-Host "=== Tunnel Command ===" -ForegroundColor Cyan
Write-Host $tunnelCommand -ForegroundColor Yellow
Write-Host ""

if ($protocol -eq "SSH") {
    Write-Host "After tunnel is established, connect with:" -ForegroundColor Cyan
    Write-Host "  ssh -p $localPort <username>@localhost" -ForegroundColor White
} elseif ($protocol -eq "RDP") {
    Write-Host "After tunnel is established, connect with:" -ForegroundColor Cyan
    Write-Host "  mstsc /v:localhost:$localPort" -ForegroundColor White
}

Write-Host ""
$proceed = Read-Host "Start tunnel now? (y/n)"

if ($proceed -eq "y") {
    Write-Host ""
    Write-Host "Starting tunnel... (Press Ctrl+C to stop)" -ForegroundColor Green
    Write-Host ""
    
    Invoke-Expression $tunnelCommand
} else {
    Write-Host ""
    Write-Host "Tunnel command saved. Run it manually when ready." -ForegroundColor Yellow
}