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 & SplunkHosts are collection and initialized
if ($null -eq $settings.Servers -or $null -eq $settings.SplunkHosts) {
    $servers = if ($null -ne $settings.Servers) { $settings.Servers } else { @("localhost") }
    $splunkHosts = if ($null -ne $settings.SplunkHosts) { $settings.SplunkHosts } else { @("sf-host-prod-01", "sf-host-prod-02", "sf-host-staging-01") }
    
    if ($settings -is [System.Management.Automation.PSCustomObject]) {
        $props = [ordered]@{
            LastServer = $settings.LastServer
            Theme = $settings.Theme
            Servers = $servers
            SplunkHosts = $splunkHosts
            SplunkRealm = if ($null -ne $settings.SplunkRealm) { $settings.SplunkRealm } else { "eu0" }
            SplunkToken = if ($null -ne $settings.SplunkToken) { $settings.SplunkToken } else { "YOUR_TOKEN" }
        }
        $settings = [pscustomobject]$props
    } else {
        $settings = [pscustomobject]@{ 
            LastServer = ""; Theme = "Light"; Servers = $servers; SplunkHosts = $splunkHosts; SplunkRealm = "eu0"; SplunkToken = "YOUR_TOKEN" 
        }
    }
}

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

# 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")
$dgDatabases = $Window.FindName("dgDatabases")
$dgServerErrors = $Window.FindName("dgServerErrors")
$tvServers = $Window.FindName("tvServers")
$btnAddServer = $Window.FindName("btnAddServer")
$btnRemoveServer = $Window.FindName("btnRemoveServer")
$gridDetails = $Window.FindName("gridDetails")
$menuSqlServers = $Window.FindName("menuSqlServers")
$menuHostSplunk = $Window.FindName("menuHostSplunk")
$grpDirectory = $Window.FindName("grpDirectory")

# Current view state: "SQL" or "Splunk"
$script:CurrentView = "SQL"
$script:SplunkHosts = @()

# 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.0. Get Splunk API Hosts function
function Get-SplunkAPIHosts {
    $realm = $settings.SplunkRealm
    if ([string]::IsNullOrWhiteSpace($realm)) { $realm = "eu0" }
    
    $token = $env:SF_TOKEN
    if ([string]::IsNullOrWhiteSpace($token)) {
        $token = $settings.SplunkToken
    }
    
    # Bezpieczne wymuszenie protokołu TLS 1.2 dla starszych wersji PowerShell
    [System.Net.ServicePointManager]::SecurityProtocol = [System.Net.SecurityProtocolType]::Tls12

    # Jeśli brak tokenu lub jest domyślny, spróbujemy poprosić użytkownika o podanie, lub użyjemy demo
    if ([string]::IsNullOrWhiteSpace($token) -or $token -eq "YOUR_TOKEN") {
        Write-Log "DEBUG: Brak tokenu w SF_TOKEN lub settings.json. Pytam użytkownika..."
        try {
            Add-Type -AssemblyName Microsoft.VisualBasic
            $inputToken = [Microsoft.VisualBasic.Interaction]::InputBox("Wprowadź swój token API Splunk Observability (pozostaw puste dla trybu DEMO):", "Wymagana Autoryzacja Splunk", "")
            if (-not [string]::IsNullOrWhiteSpace($inputToken)) {
                $token = $inputToken.Trim()
                $settings.SplunkToken = $token
                $settings | ConvertTo-Json | Set-Content $configPath
                Write-Log "DEBUG: Token został zapisany w settings.json."
            }
        } catch {
            # Ignoruj błędy ładowania VisualBasic
        }
    }

    if ([string]::IsNullOrWhiteSpace($token) -or $token -eq "YOUR_TOKEN") {
        Write-Log "DEBUG: Używam symulacji - brak prawidłowego tokenu API."
        Write-Log "DEBUG: Pokazuję pierwszych 20 wyników demonstracyjnych (zgodnie z ograniczeniem)."
        return @(
            "demo-splunk-host-eu-prod-01", "demo-splunk-host-eu-prod-02", "demo-splunk-host-eu-prod-03", "demo-splunk-host-eu-prod-04", "demo-splunk-host-eu-prod-05",
            "demo-splunk-host-eu-prod-06", "demo-splunk-host-eu-prod-07", "demo-splunk-host-eu-prod-08", "demo-splunk-host-eu-prod-09", "demo-splunk-host-eu-prod-10",
            "demo-splunk-host-eu-staging-01", "demo-splunk-host-eu-staging-02", "demo-splunk-host-eu-staging-03", "demo-splunk-host-eu-staging-04", "demo-splunk-host-eu-staging-05",
            "demo-splunk-host-eu-dev-01", "demo-splunk-host-eu-dev-02", "demo-splunk-host-eu-dev-03", "demo-splunk-host-eu-dev-04", "demo-splunk-host-eu-dev-05"
        )
    }

    $headers = @{
        "X-SF-TOKEN"   = $token
        "Content-Type" = "application/json"
    }

    # Wewnętrzny, kompatybilny helper do zapytań z logowaniem postępu na żywo
    $InvokeSplunkRequest = {
        param([string]$targetUri)
        Write-Log "DEBUG: Zapytanie HTTP GET do: $targetUri"
        
        $params = @{
            Uri         = $targetUri
            Method      = "GET"
            Headers     = $headers
            ErrorAction = "Stop"
        }
        
        # Kompatybilność parametrów timeout w zależności od wersji PowerShell (PS 5.1 vs Core)
        if ($PSVersionTable.PSVersion.Major -ge 6) {
            $params.TimeoutSec = 10
        } else {
            $params.Timeout = 10000  # 10 sekund w ms dla PS 5.1
        }
        
        $startTime = Get-Date
        $responseObj = Invoke-RestMethod @params
        $duration = [Math]::Round(((Get-Date) - $startTime).TotalMilliseconds)
        Write-Log "DEBUG: Odpowiedź otrzymana pomyślnie w $duration ms."
        return $responseObj
    }

    $hostsSet = New-Object System.Collections.Generic.HashSet[string] ([System.StringComparer]::OrdinalIgnoreCase)

    # 1. Metoda główna: GET /v2/dimension?query=host
    try {
        Write-Log "DEBUG: Rozpoczynam pobieranie wymiarów (Metoda 1: GET /v2/dimension?query=host)..."
        $limit = 20
        $offset = 0
        $hasMore = $true

        while ($hasMore) {
            $uri = "https://api.$realm.observability.splunkcloud.com/v2/dimension?query=host&limit=$limit&offset=$offset"
            $response = &$InvokeSplunkRequest -targetUri $uri
            
            $results = $null
            if ($null -ne $response) {
                if ($null -ne $response.results) { $results = $response.results }
                elseif ($response -is [System.Array]) { $results = $response }
            }

            if (($null -eq $results -or $results.Count -eq 0) -and $offset -eq 0) {
                Write-Log "DEBUG: Metoda query=host pusta. Próba z precyzyjnym query=key:host..."
                $uri = "https://api.$realm.observability.splunkcloud.com/v2/dimension?query=key:host&limit=$limit&offset=$offset"
                $response = &$InvokeSplunkRequest -targetUri $uri
                if ($null -ne $response -and $null -ne $response.results) {
                    $results = $response.results
                }
            }

            if ($null -ne $results -and $results.Count -gt 0) {
                foreach ($item in $results) {
                    if ($item.key -eq "host" -and -not [string]::IsNullOrWhiteSpace($item.value)) {
                        [void]$hostsSet.Add($item.value.Trim())
                    }
                    
                    # Przerywamy natychmiast, gdy zebraliśmy już 20 unikalnych hostów
                    if ($hostsSet.Count -ge 20) {
                        $hasMore = $false
                        break
                    }
                }
                Write-Log "DEBUG: Pobrano $($results.Count) obiektów wymiarów (offset: $offset). Dotychczas unikalnych hostów: $($hostsSet.Count)"
                
                if ($hostsSet.Count -ge 20) {
                    $hasMore = $false
                } elseif ($results.Count -lt $limit) {
                    $hasMore = $false
                } else {
                    $offset += $limit
                }
            } else {
                $hasMore = $false
            }
        }
    }
    catch {
        Write-Log "DEBUG: Metoda 1 zakończyła się błędem: $($_.Exception.Message)"
    }

    # 2. Metoda Fallback: GET /v2/metrictimeseries?query=host:*
    if ($hostsSet.Count -eq 0) {
        try {
            Write-Log "DEBUG: Metoda 1 nie zwróciła hostów. Uruchamiam Metodę zapasową 2 (GET /v2/metrictimeseries)..."
            $limit = 20
            $offset = 0
            $hasMore = $true

            while ($hasMore) {
                $uri = "https://api.$realm.observability.splunkcloud.com/v2/metrictimeseries?query=host:*&limit=$limit&offset=$offset"
                $response = &$InvokeSplunkRequest -targetUri $uri
                
                $results = $null
                if ($null -ne $response) {
                    if ($null -ne $response.results) { $results = $response.results }
                    elseif ($response -is [System.Array]) { $results = $response }
                }

                if ($null -ne $results -and $results.Count -gt 0) {
                    foreach ($mts in $results) {
                        if ($null -ne $mts.dimensions -and $null -ne $mts.dimensions.host) {
                            $val = $mts.dimensions.host
                            if (-not [string]::IsNullOrWhiteSpace($val)) {
                                [void]$hostsSet.Add($val.Trim())
                            }
                        }
                        
                        # Przerywamy natychmiast, gdy zebraliśmy już 20 unikalnych hostów
                        if ($hostsSet.Count -ge 20) {
                            $hasMore = $false
                            break
                        }
                    }
                    Write-Log "DEBUG: Pobrano $($results.Count) rekordów MTS (offset: $offset). Dotychczas unikalnych hostów: $($hostsSet.Count)"
                    
                    if ($hostsSet.Count -ge 20) {
                        $hasMore = $false
                    } elseif ($results.Count -lt $limit) {
                        $hasMore = $false
                    } else {
                        $offset += $limit
                    }
                } else {
                    $hasMore = $false
                }
            }
        }
        catch {
            Write-Log "DEBUG: Metoda 2 zakończyła się błędem: $($_.Exception.Message)"
        }
    }

    if ($hostsSet.Count -gt 0) {
        $sorted = $hostsSet | Sort-Object
        # Ograniczenie do pierwszych 20 wyników zgodnie z żądaniem
        $limited = $sorted | Select-Object -First 20
        Write-Log "DEBUG: Pobrano łącznie $($sorted.Count) unikalnych hostów. Ograniczam listę i pokazuję pierwszych $($limited.Count) wyników."
        return $limited
    } else {
        Write-Log "DEBUG: API nie zwróciło żadnych danych. Ładuję listę demonstracyjną (20 hostów)."
        return @(
            "demo-splunk-host-eu-prod-01", "demo-splunk-host-eu-prod-02", "demo-splunk-host-eu-prod-03", "demo-splunk-host-eu-prod-04", "demo-splunk-host-eu-prod-05",
            "demo-splunk-host-eu-prod-06", "demo-splunk-host-eu-prod-07", "demo-splunk-host-eu-prod-08", "demo-splunk-host-eu-prod-09", "demo-splunk-host-eu-prod-10",
            "demo-splunk-host-eu-staging-01", "demo-splunk-host-eu-staging-02", "demo-splunk-host-eu-staging-03", "demo-splunk-host-eu-staging-04", "demo-splunk-host-eu-staging-05",
            "demo-splunk-host-eu-dev-01", "demo-splunk-host-eu-dev-02", "demo-splunk-host-eu-dev-03", "demo-splunk-host-eu-dev-04", "demo-splunk-host-eu-dev-05"
        )
    }
}

# 4.1. TreeView population and refresh
function Refresh-ServerTree {
    $tvServers.Dispatcher.Invoke({
        $tvServers.Items.Clear()
        
        if ($script:CurrentView -eq "SQL") {
            $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) {
                                # Cache the loaded databases
                                if ($null -eq $script:ServerConnectionStates[$srvName]) {
                                    $script:ServerConnectionStates[$srvName] = @{
                                        Success = $true
                                        Instance = $srvName
                                        Version = "Unknown (Expanded)"
                                        Edition = "Unknown"
                                        Status = "Online"
                                        Databases = $dbs
                                        Errors = $null
                                    }
                                } else {
                                    $script:ServerConnectionStates[$srvName].Databases = $dbs
                                }

                                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)
        } else {
            # Splunk Hosts View
            $rootNode = New-Object System.Windows.Controls.TreeViewItem
            $rootNode.Header = "Splunk Observability Hosts"
            $rootNode.IsExpanded = $true
            $rootNode.Tag = "__ROOT__"
            
            foreach ($hostName in $script:SplunkHosts) {
                $hostNode = New-Object System.Windows.Controls.TreeViewItem
                $hostNode.Header = $hostName
                $hostNode.Tag = $hostName
                
                # Add a dummy loading node so the expansion arrow is displayed
                $dummyNode = New-Object System.Windows.Controls.TreeViewItem
                $dummyNode.Header = "Loading..."
                $hostNode.Items.Add($dummyNode)
                
                # Event handler when a Splunk host node is expanded
                $hostNode.Add_Expanded({
                    param($sender, $e)
                    if ($e.OriginalSource -ne $sender) { return }
                    
                    $node = $sender
                    if ($node.Items.Count -eq 1 -and $node.Items[0].Header -eq "Loading...") {
                        $node.Items.Clear()
                        $hName = $node.Tag
                        
                        Write-Log "Fetching active dimensions/properties for Splunk host $hName..."
                        
                        # Generate mock indices / indexes for Splunk host representation
                        $indexes = @(
                            [pscustomobject]@{ Name = "main_index"; Status = "Active"; Size = "1.2 GB"; RecoveryModel = "N/A"; LastFullBackup = "N/A"; LastLogBackup = "N/A" }
                            [pscustomobject]@{ Name = "infrastructure_telemetry"; Status = "Active"; Size = "15.4 GB"; RecoveryModel = "N/A"; LastFullBackup = "N/A"; LastLogBackup = "N/A" }
                            [pscustomobject]@{ Name = "audit_logs"; Status = "Active"; Size = "0.8 GB"; RecoveryModel = "N/A"; LastFullBackup = "N/A"; LastLogBackup = "N/A" }
                        )
                        
                        # Cache Splunk host info
                        $script:ServerConnectionStates[$hName] = @{
                            Success = $true
                            Instance = $hName
                            Version = "Splunk Agent v5.4.1"
                            Edition = "Enterprise Observability"
                            Status = "Active / Online"
                            Databases = $indexes
                            Errors = @(
                                [pscustomobject]@{ LogDate = (Get-Date).AddMinutes(-5).ToString(); Message = "High memory usage alert on host $hName"; Source = "SplunkObservability" }
                                [pscustomobject]@{ LogDate = (Get-Date).AddMinutes(-20).ToString(); Message = "Metric forwarder queue processed successfully"; Source = "Forwarder" }
                            )
                        }

                        foreach ($idx in $indexes) {
                            $idxNode = New-Object System.Windows.Controls.TreeViewItem
                            $idxNode.Header = "$($idx.Name) ($($idx.Status))"
                            $idxNode.Tag = $idx
                            $node.Items.Add($idxNode)
                        }
                        
                        Write-Log "Loaded active indexes for $hName."
                    }
                })
                
                $rootNode.Items.Add($hostNode)
            }
            
            $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)
                
                # Show cached databases if available
                if ($null -ne $connState.Databases) {
                    $dgDatabases.ItemsSource = @($connState.Databases)
                } else {
                    $dgDatabases.ItemsSource = $null
                }

                # Show cached error logs if available
                if ($null -ne $connState.Errors) {
                    $dgServerErrors.ItemsSource = @($connState.Errors)
                } else {
                    $dgServerErrors.ItemsSource = $null
                }
            } else {
                # Server is not connected: clear the DataGrids so no old server details/errors are visible
                $dgServerInfo.ItemsSource = $null
                $dgDatabases.ItemsSource = $null
                $dgServerErrors.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, LastFullBackup, LastLogBackup)
            $dgDatabases.ItemsSource = $null
            $dgServerErrors.ItemsSource = $null
            
            # 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
        $dgDatabases.ItemsSource = $null
        $dgServerErrors.ItemsSource = $null
    }
})

# Handle Add Server button
$btnAddServer.Add_Click({
    $server = ""
    $promptMsg = "Enter SQL Server Instance Name (e.g., localhost, srv\SQLEXPRESS):"
    $title = "Add Server"
    if ($script:CurrentView -eq "Splunk") {
        $promptMsg = "Enter Splunk Host Name (e.g., splunk-host-prod-01):"
        $title = "Add Splunk Host"
    }

    try {
        Add-Type -AssemblyName Microsoft.VisualBasic
        $server = [Microsoft.VisualBasic.Interaction]::InputBox($promptMsg, $title, "")
    } 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 name in 'Server Name' textbox first, and then click 'Add Server'.", $title, "OK", "Information")
            return
        }
    }
    
    if ([string]::IsNullOrWhiteSpace($server)) { return }
    $server = $server.Trim()
    
    if ($script:CurrentView -eq "SQL") {
        if ($settings.Servers -contains $server) {
            [System.Windows.MessageBox]::Show("Server '$server' is already in the list.", "Server Already Exists", "OK", "Information")
            return
        }
        $settings.Servers = $settings.Servers + $server
        Write-Log "Successfully added server to list: $server"
    } else {
        if ($settings.SplunkHosts -contains $server) {
            [System.Windows.MessageBox]::Show("Splunk Host '$server' is already in the list.", "Host Already Exists", "OK", "Information")
            return
        }
        $settings.SplunkHosts = $settings.SplunkHosts + $server
        $script:SplunkHosts = $settings.SplunkHosts
        Write-Log "Successfully added Splunk host: $server"
    }

    # Save to settings
    $settings | ConvertTo-Json | Set-Content $configPath
    Refresh-ServerTree
})

# Handle Remove Server button
$btnRemoveServer.Add_Click({
    $selectedItem = $tvServers.SelectedItem
    if ($null -eq $selectedItem -or $selectedItem.Tag -eq "__ROOT__") {
        $msg = if ($script:CurrentView -eq "SQL") { "Please select a server node in the tree to remove." } else { "Please select a Splunk Host node in the tree to remove." }
        [System.Windows.MessageBox]::Show($msg, "Select Node", "OK", "Warning")
        return
    }
    
    $serverName = $selectedItem.Tag
    if ($serverName -is [string]) {
        if ($script:CurrentView -eq "SQL" -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 }
                Write-Log "Removed server from list: $serverName"
            } else { return }
        } elseif ($script:CurrentView -eq "Splunk" -and ($settings.SplunkHosts -contains $serverName)) {
            $confirm = [System.Windows.MessageBox]::Show("Are you sure you want to remove '$serverName' from the list?", "Confirm Host Removal", "YesNo", "Question")
            if ($confirm -eq "Yes") {
                $settings.SplunkHosts = $settings.SplunkHosts | Where-Object { $_ -ne $serverName }
                $script:SplunkHosts = $settings.SplunkHosts
                Write-Log "Removed Splunk host from list: $serverName"
            } else { return }
        } else {
            [System.Windows.MessageBox]::Show("Only host/server nodes can be removed.", "Invalid Selection", "OK", "Warning")
            return
        }

        # Save to settings
        $settings | ConvertTo-Json | Set-Content $configPath
        Refresh-ServerTree
    } else {
        [System.Windows.MessageBox]::Show("Only server/host nodes can be removed.", "Invalid Selection", "OK", "Warning")
    }
})

# Handle Main -> SQL Servers menu click
$menuSqlServers.Add_Click({
    $script:CurrentView = "SQL"
    $grpDirectory.Header = "SQL Server Directory"
    Write-Log "Switched view to SQL Server Directory."
    Refresh-ServerTree
})

# Handle Main -> Host Splunk menu click
$menuHostSplunk.Add_Click({
    $script:CurrentView = "Splunk"
    $grpDirectory.Header = "Splunk Hosts Directory"
    Write-Log "Switched view to Splunk Observability Cloud."
    
    # Fetch Splunk Hosts dynamically from API (or fallback/demo)
    $script:SplunkHosts = Get-SplunkAPIHosts
    Refresh-ServerTree
})

# 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)
            
            # Fetch databases list using Get-SQLDatabases
            Write-Log "Fetching databases list for $server..."
            try {
                $dbs = Get-SQLDatabases -ServerInstance $server
                if ($dbs -and $dbs.Count -gt 0) {
                    $dgDatabases.ItemsSource = @($dbs)
                    $result.Databases = $dbs
                    Write-Log "Loaded $($dbs.Count) databases."
                } else {
                    $dgDatabases.ItemsSource = $null
                    $result.Databases = $null
                    Write-Log "No databases found on $server."
                }
            } catch {
                Write-Log "Warning: Could not load databases: $($_.Exception.Message)"
                $dgDatabases.ItemsSource = $null
                $result.Databases = $null
            }

            # Fetch errors from the last hour using Get-SQLErrors
            Write-Log "Fetching error logs (last hour) for $server..."
            try {
                $errors = Get-SQLErrors -ServerInstance $server
                if ($errors -and $errors.Count -gt 0) {
                    $dgServerErrors.ItemsSource = @($errors)
                    $result.Errors = $errors
                    Write-Log "Loaded $($errors.Count) error log entries."
                } else {
                    $dgServerErrors.ItemsSource = $null
                    $result.Errors = $null
                    Write-Log "No errors found in the last hour."
                }
            } catch {
                Write-Log "Warning: Could not load error log: $($_.Exception.Message)"
                $dgServerErrors.ItemsSource = $null
                $result.Errors = $null
            }

            # Cache successful connection state with its fetched errors and databases
            $script:ServerConnectionStates[$server] = $result
        } else {
            Write-Log "Connection failed: $($result.Error)"
            $script:ServerConnectionStates[$server] = $null
            $dgServerInfo.ItemsSource = $null
            $dgDatabases.ItemsSource = $null
            $dgServerErrors.ItemsSource = $null
        }
    } catch {
        Write-Log "Critical Error: $($_.Exception.Message)"
        $script:ServerConnectionStates[$server] = $null
        $dgServerInfo.ItemsSource = $null
        $dgDatabases.ItemsSource = $null
        $dgServerErrors.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