DBAdmin/toolbox/DBAToolBox.ps1

Add-Type -AssemblyName PresentationFramework, PresentationCore, WindowsBase

# Import Custom Modules
Import-Module (Join-Path $PSScriptRoot "Modules\SQLManager.psm1") -Force

# Configuration Paths
$configPath = Join-Path $PSScriptRoot "Config\settings.json"

# Load Settings
if (Test-Path $configPath) {
    $settings = Get-Content $configPath | ConvertFrom-Json
} else {
    $settings = [pscustomobject]@{ LastServer = ""; Theme = "Light"; Servers = @("localhost") }
}

# Ensure Servers is a collection and initialized
if ($null -eq $settings.Servers) {
    if ($settings -is [System.Management.Automation.PSCustomObject]) {
        $props = [ordered]@{
            LastServer = $settings.LastServer
            Theme = $settings.Theme
            Servers = @("localhost")
        }
        $settings = [pscustomobject]$props
    } else {
        $settings = [pscustomobject]@{ LastServer = ""; Theme = "Light"; Servers = @("localhost") }
    }
}

if ($settings.Servers -is [string]) {
    $settings.Servers = @($settings.Servers)
}

# 1. Path to XAML
$xamlFile = Join-Path $PSScriptRoot "UI\MainWindow.xaml"

# 2. Load XAML
[xml]$xml = Get-Content $xamlFile
$reader = New-Object System.Xml.XmlNodeReader($xml)
$Window = [Windows.Markup.XamlReader]::Load($reader)

# 3. Extract controls from XAML
$btnConnect = $Window.FindName("btnConnect")
$txtServerName = $Window.FindName("txtServerName")
$txtLogs = $Window.FindName("txtLogs")
$dgServerInfo = $Window.FindName("dgServerInfo")
$dgServerErrors = $Window.FindName("dgServerErrors")
$tvServers = $Window.FindName("tvServers")
$btnAddServer = $Window.FindName("btnAddServer")
$btnRemoveServer = $Window.FindName("btnRemoveServer")
$gridDetails = $Window.FindName("gridDetails")

# Initialize connection states cache per server
$script:ServerConnectionStates = @{}

# Set initial values from settings
$txtServerName.Text = $settings.LastServer

# 4. Helper function for logging
function Write-Log {
    param([string]$Message)
    $timestamp = Get-Date -Format "HH:mm:ss"
    $txtLogs.Dispatcher.Invoke({
        $txtLogs.AppendText("[$timestamp] $Message`r`n")
        $txtLogs.ScrollToEnd()
    })
}

# 4.1. TreeView population and refresh
function Refresh-ServerTree {
    $tvServers.Dispatcher.Invoke({
        $tvServers.Items.Clear()
        
        $rootNode = New-Object System.Windows.Controls.TreeViewItem
        $rootNode.Header = "SQL Servers"
        $rootNode.IsExpanded = $true
        $rootNode.Tag = "__ROOT__"
        
        foreach ($server in $settings.Servers) {
            $serverNode = New-Object System.Windows.Controls.TreeViewItem
            $serverNode.Header = $server
            $serverNode.Tag = $server
            
            # Add a dummy loading node so the expansion arrow is displayed
            $dummyNode = New-Object System.Windows.Controls.TreeViewItem
            $dummyNode.Header = "Loading..."
            $serverNode.Items.Add($dummyNode)
            
            # Event handler when a server node is expanded
            $serverNode.Add_Expanded({
                param($sender, $e)
                # Prevent this handler from running for nested child expansions (bubbling)
                if ($e.OriginalSource -ne $sender) { return }
                
                $node = $sender
                if ($node.Items.Count -eq 1 -and $node.Items[0].Header -eq "Loading...") {
                    $node.Items.Clear()
                    $srvName = $node.Tag
                    
                    Write-Log "Fetching databases for $srvName..."
                    # Force UI to process the pending write-log before connection blocks the thread
                    [System.Windows.Threading.Dispatcher]::CurrentDispatcher.Invoke(
                        [System.Windows.Threading.DispatcherPriority]::Background, 
                        [System.Action]{}
                    )
                    
                    try {
                        $dbs = Get-SQLDatabases -ServerInstance $srvName
                        if ($dbs -and $dbs.Count -gt 0) {
                            foreach ($db in $dbs) {
                                $dbNode = New-Object System.Windows.Controls.TreeViewItem
                                $dbNode.Header = "$($db.Name) ($($db.Status))"
                                $dbNode.Tag = $db
                                $node.Items.Add($dbNode)
                            }
                            Write-Log "Successfully loaded $($dbs.Count) databases for $srvName."
                        } else {
                            $noDbNode = New-Object System.Windows.Controls.TreeViewItem
                            $noDbNode.Header = "(No databases found)"
                            $node.Items.Add($noDbNode)
                            Write-Log "No databases found on $srvName."
                        }
                    } catch {
                        Write-Log "Error fetching databases: $($_.Exception.Message)"
                        $errNode = New-Object System.Windows.Controls.TreeViewItem
                        $errNode.Header = "(Error loading databases)"
                        $node.Items.Add($errNode)
                    }
                }
            })
            
            $rootNode.Items.Add($serverNode)
        }
        
        $tvServers.Items.Add($rootNode)
    })
}

# Handle TreeView selection
$tvServers.Add_SelectedItemChanged({
    param($sender, $e)
    $selectedItem = $tvServers.SelectedItem
    if ($null -ne $selectedItem -and $null -ne $selectedItem.Tag -and $selectedItem.Tag -ne "__ROOT__") {
        # Show detail panel
        $gridDetails.Visibility = "Visible"
        
        if ($selectedItem.Tag -is [string]) {
            # It's a server node, copy its name to TextBox and update its cached connection state
            $srvName = $selectedItem.Tag
            $txtServerName.Text = $srvName
            Write-Log "Selected server: $srvName"
            
            # Check cached connection state for this server
            $connState = $script:ServerConnectionStates[$srvName]
            if ($null -ne $connState -and $connState.Success) {
                # Server is connected: show its version/details in the DataGrid
                $dgServerInfo.ItemsSource = @([PSCustomObject]$connState | Select-Object Instance, Version, Edition, Status)
            } else {
                # Server is not connected: clear the DataGrid so no old server details are visible
                $dgServerInfo.ItemsSource = $null
            }
        } elseif ($selectedItem.Tag -is [System.Management.Automation.PSCustomObject]) {
            # It's a database node
            $db = $selectedItem.Tag
            Write-Log "Selected database: $($db.Name) (Size: $($db.Size), Status: $($db.Status))"
            
            # Show database details in the main information grid
            $dgServerInfo.ItemsSource = @($db | Select-Object Name, Status, Size, RecoveryModel)
            
            # Also, update txtServerName with the parent server's tag so they see which server it belongs to
            $parentItem = $selectedItem.Parent
            if ($null -ne $parentItem -and $parentItem.Tag -is [string]) {
                $txtServerName.Text = $parentItem.Tag
            }
        }
    } else {
        # Hide detail panel if selection is null or ROOT
        $gridDetails.Visibility = "Collapsed"
        $dgServerInfo.ItemsSource = $null
    }
})

# Handle Add Server button
$btnAddServer.Add_Click({
    $server = ""
    try {
        Add-Type -AssemblyName Microsoft.VisualBasic
        $server = [Microsoft.VisualBasic.Interaction]::InputBox("Enter SQL Server Instance Name (e.g., localhost, srv\SQLEXPRESS):", "Add Server", "")
    } catch {
        # Fallback using current textbox value
        $server = $txtServerName.Text
        if ([string]::IsNullOrWhiteSpace($server)) {
            [System.Windows.MessageBox]::Show("Visual Basic assembly not found. Please type the server name in 'Server Name' textbox first, and then click 'Add Server'.", "Add Server Fallback", "OK", "Information")
            return
        }
    }
    
    if ([string]::IsNullOrWhiteSpace($server)) { return }
    $server = $server.Trim()
    
    if ($settings.Servers -contains $server) {
        [System.Windows.MessageBox]::Show("Server '$server' is already in the list.", "Server Already Exists", "OK", "Information")
        return
    }
    
    # Save to settings
    $settings.Servers = $settings.Servers + $server
    $settings | ConvertTo-Json | Set-Content $configPath
    
    Refresh-ServerTree
    Write-Log "Successfully added server to list: $server"
})

# Handle Remove Server button
$btnRemoveServer.Add_Click({
    $selectedItem = $tvServers.SelectedItem
    if ($null -eq $selectedItem -or $selectedItem.Tag -eq "__ROOT__") {
        [System.Windows.MessageBox]::Show("Please select a server node in the tree to remove.", "Select Server", "OK", "Warning")
        return
    }
    
    $serverName = $selectedItem.Tag
    # Verify that we selected a server node (it has a string Tag and is in settings.Servers)
    if ($serverName -is [string] -and ($settings.Servers -contains $serverName)) {
        $confirm = [System.Windows.MessageBox]::Show("Are you sure you want to remove '$serverName' from the list?", "Confirm Server Removal", "YesNo", "Question")
        if ($confirm -eq "Yes") {
            $settings.Servers = $settings.Servers | Where-Object { $_ -ne $serverName }
            $settings | ConvertTo-Json | Set-Content $configPath
            
            Refresh-ServerTree
            Write-Log "Removed server from list: $serverName"
        }
    } else {
        [System.Windows.MessageBox]::Show("Only server nodes can be removed.", "Invalid Selection", "OK", "Warning")
    }
})

# 5. Connection logic
$btnConnect.Add_Click({
    $server = $txtServerName.Text
    if ([string]::IsNullOrWhiteSpace($server)) {
        Write-Log "Error: Please enter a server name."
        return
    }

    $btnConnect.IsEnabled = $false
    Write-Log "Attempting to connect to $server..."

    # Save last server
    $settings.LastServer = $server
    $settings | ConvertTo-Json | Set-Content $configPath
    
    try {
        $result = Test-SQLConnection -ServerInstance $server
        
        if ($result.Success) {
            Write-Log "Connected successfully to $($result.Instance)!"
            $dgServerInfo.ItemsSource = @([PSCustomObject]$result | Select-Object Instance, Version, Edition, Status)
            # Cache successful connection state
            $script:ServerConnectionStates[$server] = $result
        } else {
            Write-Log "Connection failed: $($result.Error)"
            $script:ServerConnectionStates[$server] = $null
            $dgServerInfo.ItemsSource = $null
        }
    } catch {
        Write-Log "Critical Error: $($_.Exception.Message)"
        $script:ServerConnectionStates[$server] = $null
        $dgServerInfo.ItemsSource = $null
    } finally {
        $btnConnect.IsEnabled = $true
    }
})

# Startup Dependency Check and Tree Initialization
try {
    Import-DBAToolBoxDependencies
    Write-Log "DBAToolBox started. dbatools module loaded."
    Refresh-ServerTree
} catch {
    [System.Windows.MessageBox]::Show($_.Exception.Message, "Dependency Error", "OK", "Error")
    $Window.Close()
}

# 6. Show Window
$Window.ShowDialog() | Out-Null