package org.openjproxy.grpc.client;

import com.openjproxy.grpc.ConnectionDetails;
import com.openjproxy.grpc.PropertyEntry;
import com.openjproxy.grpc.SessionInfo;
import com.openjproxy.grpc.StatementServiceGrpc;
import io.grpc.ManagedChannel;
import io.grpc.StatusRuntimeException;
import org.openjproxy.constants.CommonConstants;
import org.openjproxy.grpc.GrpcChannelFactory;
import org.openjproxy.grpc.ProtoConverter;
import org.openjproxy.jdbc.DatasourcePropertiesLoader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

/**
 * Manages multinode connections to OJP servers with load-aware routing,
 * session stickiness, and failover support.
 * 
 * Server Selection Strategies:
 * - Load-aware (default): Selects the healthy server with the fewest active connections
 * - Round-robin (legacy): Cycles through healthy servers in order
 * 
 * Addressing PR #39 review comments:
 * - Comment #3: Throws exception when session server is unavailable (enforces session stickiness)
 * - Configurable retry attempts and delays (via CommonConstants properties)
 */
public class MultinodeConnectionManager {
    
    private static final Logger log = LoggerFactory.getLogger(MultinodeConnectionManager.class);
    private static final String DNS_PREFIX = "dns:///";

    private final List<ServerEndpoint> serverEndpoints;
    private final Map<ServerEndpoint, ChannelAndStub> channelMap;
    private final Map<String, ServerEndpoint> sessionToServerMap; // sessionUUID -> server
    private final Map<String, List<ServerEndpoint>> connHashToServersMap; // connHash -> list of servers that received connect()
    private final AtomicInteger roundRobinCounter;
    private final int retryAttempts;
    private final long retryDelayMs;
    private final List<ServerHealthListener> healthListeners; // Phase 2: Listeners for server health changes
    
    // Health check and redistribution support
    private final HealthCheckConfig healthCheckConfig;
    private final AtomicLong lastHealthCheckTimestamp;
    private final HealthCheckValidator healthCheckValidator;
    private final ConnectionTracker connectionTracker; // Legacy - kept for backward compatibility
    private final SessionTracker sessionTracker; // New unified tracker
    private final ConnectionRedistributor connectionRedistributor;
    private XAConnectionRedistributor xaConnectionRedistributor;
    private final ScheduledExecutorService healthCheckScheduler;
    
    public MultinodeConnectionManager(List<ServerEndpoint> serverEndpoints) {
        this(serverEndpoints, CommonConstants.DEFAULT_MULTINODE_RETRY_ATTEMPTS, 
             CommonConstants.DEFAULT_MULTINODE_RETRY_DELAY_MS, null);
    }
    
    public MultinodeConnectionManager(List<ServerEndpoint> serverEndpoints, 
                                    int retryAttempts, long retryDelayMs) {
        this(serverEndpoints, retryAttempts, retryDelayMs, null);
    }
    
    public MultinodeConnectionManager(List<ServerEndpoint> serverEndpoints, 
                                    int retryAttempts, long retryDelayMs,
                                    HealthCheckConfig healthCheckConfig) {
        this(serverEndpoints, retryAttempts, retryDelayMs, healthCheckConfig, null);
    }
    
    public MultinodeConnectionManager(List<ServerEndpoint> serverEndpoints, 
                                    int retryAttempts, long retryDelayMs,
                                    HealthCheckConfig healthCheckConfig,
                                    ConnectionTracker connectionTracker) {
        if (serverEndpoints == null || serverEndpoints.isEmpty()) {
            throw new IllegalArgumentException("Server endpoints list cannot be null or empty");
        }
        
        this.serverEndpoints = List.copyOf(serverEndpoints);
        this.channelMap = new ConcurrentHashMap<>();
        this.sessionToServerMap = new ConcurrentHashMap<>();
        this.connHashToServersMap = new ConcurrentHashMap<>();
        this.roundRobinCounter = new AtomicInteger(0);
        this.retryAttempts = retryAttempts;
        this.retryDelayMs = retryDelayMs;
        this.healthListeners = new ArrayList<>();
        this.healthCheckConfig = healthCheckConfig != null ? healthCheckConfig : HealthCheckConfig.createDefault();
        this.lastHealthCheckTimestamp = new AtomicLong(0);
        this.connectionTracker = connectionTracker != null ? connectionTracker : new ConnectionTracker(); // Legacy
        this.sessionTracker = new SessionTracker();
        this.healthCheckValidator = new HealthCheckValidator(this.healthCheckConfig, this);
        this.connectionRedistributor = new ConnectionRedistributor(this.connectionTracker, this.healthCheckConfig);
        
        // Initialize channels and stubs for all servers
        initializeConnections();
        
        // Start periodic health checker
        if (this.healthCheckConfig.isRedistributionEnabled()) {
            this.healthCheckScheduler = Executors.newSingleThreadScheduledExecutor(r -> {
                Thread t = new Thread(r, "ojp-health-checker");
                t.setDaemon(true); // Don't prevent JVM shutdown
                return t;
            });
            
            long intervalMs = this.healthCheckConfig.getHealthCheckIntervalMs();
            this.healthCheckScheduler.scheduleAtFixedRate(
                () -> {
                    try {
                        performHealthCheck();
                    } catch (Exception e) {
                        log.warn("Periodic health check failed: {}", e.getMessage());
                    }
                },
                intervalMs, // Initial delay
                intervalMs, // Period
                TimeUnit.MILLISECONDS
            );
            log.info("Started periodic health checker with interval {}ms", intervalMs);
        } else {
            this.healthCheckScheduler = null;
            log.info("Periodic health checker disabled (redistribution not enabled)");
        }
        
        log.info("MultinodeConnectionManager initialized with {} servers: {}, health check config: {}", 
                serverEndpoints.size(), serverEndpoints, this.healthCheckConfig);
    }
    
    private void initializeConnections() {
        for (ServerEndpoint endpoint : serverEndpoints) {
            try {
                createChannelAndStub(endpoint);
                endpoint.markHealthy();
                log.debug("Successfully initialized connection to {}", endpoint.getAddress());
            } catch (Exception e) {
                log.warn("Failed to initialize connection to {}: {}", endpoint.getAddress(), e.getMessage());
                endpoint.markUnhealthy();
                endpoint.setLastFailureTime(System.currentTimeMillis());
            }
        }
    }
    
    private ChannelAndStub createChannelAndStub(ServerEndpoint endpoint) {
        String target = DNS_PREFIX + endpoint.getHost() + ":" + endpoint.getPort();
        ManagedChannel channel = GrpcChannelFactory.createChannel(target);
        
        StatementServiceGrpc.StatementServiceBlockingStub blockingStub = 
                StatementServiceGrpc.newBlockingStub(channel);
        StatementServiceGrpc.StatementServiceStub asyncStub = 
                StatementServiceGrpc.newStub(channel);
        
        ChannelAndStub channelAndStub = new ChannelAndStub(channel, blockingStub, asyncStub);
        channelMap.put(endpoint, channelAndStub);
        
        return channelAndStub;
    }
    
    /**
     * Establishes a connection by calling connect() on servers.
     * 
     * UNIFIED MODE (always enabled):
     * - Both XA and non-XA connections connect to ALL servers
     * - Ensures all servers have datasource configuration for improved failover
     * - Sessions tracked via SessionTracker for accurate load balancing
     * - Each session bound directly to the ServerEndpoint it was created on
     * 
     * Returns the SessionInfo from the successful connection.
     */
    public SessionInfo connect(ConnectionDetails connectionDetails) throws SQLException {
        boolean isXA = connectionDetails.getIsXA();
        
        log.info("=== connect() called: isXA={} (unified mode always enabled) ===", isXA);
        
        // Try to trigger health check (time-based, non-blocking)
        if (healthCheckConfig.isRedistributionEnabled()) {
            tryTriggerHealthCheck();
        }
        
        // UNIFIED MODE: Both XA and non-XA connect to all servers
        log.debug("Connecting to all servers for {} connection", isXA ? "XA" : "non-XA");
        return connectToAllServers(connectionDetails);
    }
    
    /**
     * Attempts to trigger a health check if enough time has elapsed since the last check.
     * Uses compareAndSet to ensure only one thread executes the health check.
     * Non-blocking - if another thread is already doing a health check, this returns immediately.
     */
    private void tryTriggerHealthCheck() {
        long now = System.currentTimeMillis();
        long lastCheck = lastHealthCheckTimestamp.get();
        long elapsed = now - lastCheck;
        
        // Only check if interval has passed
        if (elapsed >= healthCheckConfig.getHealthCheckIntervalMs()) {
            // Atomic update - only one thread succeeds
            if (lastHealthCheckTimestamp.compareAndSet(lastCheck, now)) {
                try {
                    performHealthCheck();
                } catch (Exception e) {
                    log.warn("Health check failed: {}", e.getMessage());
                    // Don't fail the connection attempt - health check is best effort
                }
            }
        }
    }
    
    /**
     * Performs health check on all servers.
     * - For XA mode: Proactively checks healthy servers and invalidates sessions when they fail
     * - For all modes: Checks unhealthy servers to see if they've recovered
     */
    private void performHealthCheck() {
        log.debug("Performing health check on servers");
        
        // XA Mode: Proactively check healthy servers to detect failures early
        // This ensures sessions are invalidated even if no active operations are hitting the server
        if (xaConnectionRedistributor != null) {
            List<ServerEndpoint> healthyServers = serverEndpoints.stream()
                    .filter(ServerEndpoint::isHealthy)
                    .collect(Collectors.toList());
            
            for (ServerEndpoint endpoint : healthyServers) {
                if (!validateServer(endpoint)) {
                    log.info("XA Health check: Server {} has become unhealthy", endpoint.getAddress());
                    
                    // Mark server unhealthy
                    endpoint.setHealthy(false);
                    endpoint.setLastFailureTime(System.currentTimeMillis());
                    
                    // XA Mode: Immediately invalidate sessions and connections for the failed server
                    invalidateSessionsAndConnectionsForFailedServer(endpoint);
                    
                    // Notify listeners
                    notifyServerUnhealthy(endpoint, new Exception("Health check failed"));
                }
            }
        }
        
        // Check unhealthy servers to see if they've recovered
        List<ServerEndpoint> unhealthyServers = serverEndpoints.stream()
                .filter(endpoint -> !endpoint.isHealthy())
                .collect(Collectors.toList());
        
        if (unhealthyServers.isEmpty()) {
            log.debug("No unhealthy servers to check");
            return;
        }
        
        log.info("Checking {} unhealthy server(s)", unhealthyServers.size());
        
        List<ServerEndpoint> recoveredServers = new ArrayList<>();
        
        // Check each unhealthy server
        for (ServerEndpoint endpoint : unhealthyServers) {
            long timeSinceFailure = System.currentTimeMillis() - endpoint.getLastFailureTime();
            
            // Only check if enough time has passed since last failure
            if (timeSinceFailure >= healthCheckConfig.getHealthCheckThresholdMs()) {
                if (validateServer(endpoint)) {
                    log.info("Server {} has recovered", endpoint.getAddress());
                    
                    endpoint.markHealthy();
                    recoveredServers.add(endpoint);
                    notifyServerRecovered(endpoint);
                } else {
                    // Still unhealthy, update timestamp
                    endpoint.setLastFailureTime(System.currentTimeMillis());
                    log.debug("Server {} still unhealthy", endpoint.getAddress());
                }
            }
        }
        
        // For non-XA mode: trigger connection redistribution when servers recover
        // XA mode handles redistribution differently (through invalidation), so skip it
        if (!recoveredServers.isEmpty() && xaConnectionRedistributor == null && healthCheckConfig.isRedistributionEnabled()) {
            log.info("Triggering connection redistribution for {} recovered server(s)", 
                    recoveredServers.size());
            
            List<ServerEndpoint> allHealthyServers = serverEndpoints.stream()
                    .filter(ServerEndpoint::isHealthy)
                    .collect(Collectors.toList());
            
            try {
                connectionRedistributor.rebalance(recoveredServers, allHealthyServers);
            } catch (Exception e) {
                log.error("Failed to redistribute connections after server recovery: {}", 
                        e.getMessage(), e);
            }
        }
    }

    /**
     * Invalidates a specified number of connections for a server.
     * 
     * @param server The server whose connections should be invalidated
     * @param connections The list of connections for this server
     * @param count The number of connections to invalidate
     * @return The actual number of connections invalidated
     */
    private int invalidateConnectionsForServer(ServerEndpoint server, List<java.sql.Connection> connections, int count) {
        int invalidated = 0;
        int toInvalidate = Math.min(count, connections.size());
        
        for (int i = 0; i < toInvalidate; i++) {
            java.sql.Connection conn = connections.get(i);
            if (conn instanceof org.openjproxy.jdbc.Connection) {
                org.openjproxy.jdbc.Connection ojpConn = (org.openjproxy.jdbc.Connection) conn;
                ojpConn.markForceInvalid();
                try {
                    conn.close();
                    invalidated++;
                    log.debug("Invalidated and closed connection {} for server {} during rebalancing", 
                            System.identityHashCode(conn), server.getAddress());
                } catch (Exception e) {
                    log.warn("Failed to close connection {} for server {} during rebalancing: {}", 
                            System.identityHashCode(conn), server.getAddress(), e.getMessage());
                }
            }
        }
        
        log.info("Invalidated {} of {} connections for server {} during rebalancing", 
                invalidated, toInvalidate, server.getAddress());
        return invalidated;
    }
    
    /**
     * Invalidates all XA sessions and connections for a server that has become unhealthy.
     * 
     * In XA mode, when a server fails, we immediately:
     * 1. Clear client-side session bindings from sessionToServerMap
     * 2. Unregister sessions from SessionTracker
     * 3. Mark Connection objects as invalid (forceInvalid) so pools discard them
     * 4. Close connections to force pool replacement
     * 
     * This prevents attempts to use stale sessions on the failed server. When the server
     * recovers, new connections will be created with fresh sessions.
     * 
     * @param endpoint The server endpoint that has failed
     */
    private void invalidateSessionsAndConnectionsForFailedServer(ServerEndpoint endpoint) {
        log.info("Invalidating all XA sessions and connections for failed server {}", endpoint.getAddress());
        
        // Step 1: Remove all session bindings for this server
        List<String> sessionsToInvalidate = sessionToServerMap.entrySet().stream()
                .filter(entry -> entry.getValue().equals(endpoint))
                .map(Map.Entry::getKey)
                .collect(Collectors.toList());
        
        for (String sessionUUID : sessionsToInvalidate) {
            sessionToServerMap.remove(sessionUUID);
            // Unregister from SessionTracker
            sessionTracker.unregisterSession(sessionUUID);
            log.debug("Removed session binding {} for failed server {}", sessionUUID, endpoint.getAddress());
        }
        
        // Step 2: Mark all connections for this server as invalid and close them
        Map<ServerEndpoint, List<java.sql.Connection>> distribution = connectionTracker.getDistribution();
        List<java.sql.Connection> connectionsToInvalidate = distribution.get(endpoint);
        
        if (connectionsToInvalidate != null && !connectionsToInvalidate.isEmpty()) {
            int invalidatedCount = 0;
            for (java.sql.Connection conn : connectionsToInvalidate) {
                if (conn instanceof org.openjproxy.jdbc.Connection) {
                    org.openjproxy.jdbc.Connection ojpConn = (org.openjproxy.jdbc.Connection) conn;
                    ojpConn.markForceInvalid();
                    try {
                        conn.close();
                        invalidatedCount++;
                        log.debug("Invalidated and closed connection {} for failed server {}", 
                                System.identityHashCode(conn), endpoint.getAddress());
                    } catch (Exception e) {
                        log.warn("Failed to close connection {} for failed server {}: {}", 
                                System.identityHashCode(conn), endpoint.getAddress(), e.getMessage());
                    }
                }
            }
            
            log.info("Invalidated {} session(s) and {} connection(s) for failed server {}",
                    sessionsToInvalidate.size(), invalidatedCount, endpoint.getAddress());
        } else {
            log.info("Invalidated {} session(s) for failed server {} (no connections tracked)",
                    sessionsToInvalidate.size(), endpoint.getAddress());
        }
    }
    
    /**
     * Validates if a server is healthy by attempting a simple connection.
     * Returns true if server is responsive, false otherwise.
     */
    private boolean validateServer(ServerEndpoint endpoint) {
        return healthCheckValidator.validateServer(endpoint);
    }
    
    /**
     * Provides public access to create a channel and stub for an endpoint.
     * Used by HealthCheckValidator for server validation.
     */
    public ChannelAndStub createChannelAndStubForEndpoint(ServerEndpoint endpoint) {
        return createChannelAndStub(endpoint);
    }
    
    /**
     * Closes a session by its UUID.
     * Used by HealthCheckValidator to clean up test connections.
     */
    public void closeSession(String sessionUUID) throws SQLException {
        if (sessionUUID == null || sessionUUID.isEmpty()) {
            return;
        }
        
        // Remove from session map
        ServerEndpoint server = sessionToServerMap.remove(sessionUUID);
        
        if (server != null) {
            log.debug("Closed session {}, was bound to {}", sessionUUID, server.getAddress());
        }
        
        // TODO: Actually call close on the server - for now just remove from tracking
    }
    
    /**
     * Gets the connection tracker for integration with XA data sources.
     * @deprecated Use getSessionTracker() for unified tracking
     */
    @Deprecated
    public ConnectionTracker getConnectionTracker() {
        return connectionTracker;
    }
    
    /**
     * Gets the session tracker for unified session-based load balancing.
     * Replaces ConnectionTracker with lighter-weight session tracking.
     */
    public SessionTracker getSessionTracker() {
        return sessionTracker;
    }
    
    /**
     * Gets the health check configuration.
     */
    public HealthCheckConfig getHealthCheckConfig() {
        return healthCheckConfig;
    }
    
    /**
     * Connects to all servers to ensure datasource information is available on all nodes.
     * Used for both XA and non-XA connections in unified mode.
     * 
     * Each session is bound directly to the ServerEndpoint object that was just connected to.
     * Only sessions with UUIDs (actual server sessions) are bound.
     */
    private SessionInfo connectToAllServers(ConnectionDetails connectionDetails) throws SQLException {
        SessionInfo primarySessionInfo = null;
        SQLException lastException = null;
        int successfulConnections = 0;
        List<ServerEndpoint> connectedServers = new ArrayList<>();
        boolean isXA = connectionDetails.getIsXA();
        
        // Try to connect to all servers
        for (ServerEndpoint server : serverEndpoints) {
            if (!server.isHealthy()) {
                // Attempt to recover unhealthy servers if enough time has passed
                long currentTime = System.currentTimeMillis();
                if ((currentTime - server.getLastFailureTime()) > retryDelayMs) {
                    log.info("Attempting to recover unhealthy server {} during connect()", server.getAddress());
                    try {
                        createChannelAndStub(server);
                        server.setHealthy(true);
                        server.setLastFailureTime(0);
                        log.info("Successfully recovered server {} during connect()", server.getAddress());
                        // Continue to attempt connection below
                    } catch (Exception e) {
                        server.setLastFailureTime(currentTime);
                        log.debug("Server {} recovery attempt failed during connect(): {}", 
                                server.getAddress(), e.getMessage());
                        continue;  // Skip this server
                    }
                } else {
                    log.debug("Skipping unhealthy server: {} (waiting for retry delay)", server.getAddress());
                    continue;
                }
            }
            
            try {
                ChannelAndStub channelAndStub = channelMap.get(server);
                if (channelAndStub == null) {
                    channelAndStub = createChannelAndStub(server);
                }
                
                log.info("Connecting to server {} ({} connection)", server.getAddress(), isXA ? "XA" : "non-XA");
                SessionInfo sessionInfo = channelAndStub.blockingStub.connect(connectionDetails);
                
                // Mark server as healthy
                server.setHealthy(true);
                server.setLastFailureTime(0);
                
                // Bind session directly to the ServerEndpoint we just connected to
                // Only bind sessions with UUIDs (actual server sessions, not lazy allocation)
                if (sessionInfo.getSessionUUID() != null && !sessionInfo.getSessionUUID().isEmpty()) {
                    // Direct binding to ServerEndpoint object - no string matching needed
                    sessionToServerMap.put(sessionInfo.getSessionUUID(), server);
                    sessionTracker.registerSession(sessionInfo.getSessionUUID(), server);
                    
                    String connectedServerAddress = server.getHost() + ":" + server.getPort();
                    log.info("Session {} bound to server {} (direct binding to ServerEndpoint)", 
                            sessionInfo.getSessionUUID(), connectedServerAddress);
                } else {
                    // No UUID = lazy allocation = no session on server yet = don't bind
                    log.debug("No sessionUUID from server {} - skipping binding (lazy allocation)", 
                            server.getAddress());
                }
                
                log.info("Successfully connected to server {} ({} connection)", 
                        server.getAddress(), isXA ? "XA" : "non-XA");
                successfulConnections++;
                
                // Track that this server received a connect() call
                connectedServers.add(server);
                
                // Use the first successful connection as the primary
                if (primarySessionInfo == null) {
                    primarySessionInfo = sessionInfo;
                }
                
            } catch (StatusRuntimeException e) {
                try {
                    GrpcExceptionHandler.handle(e);
                    lastException = new SQLException("gRPC call failed: " + e.getMessage(), e);
                } catch (SQLException sqlEx) {
                    lastException = sqlEx;
                }
                handleServerFailure(server, e);
                
                log.warn("Connection failed to server {}: {}", 
                        server.getAddress(), lastException.getMessage());
            }
        }
        
        if (primarySessionInfo == null) {
            throw new SQLException("Failed to connect to any server. " +
                    "Last error: " + (lastException != null ? lastException.getMessage() : "No healthy servers available"));
        }
        
        // Track which servers received connect() for this connection hash
        // This is used during terminateSession() to ensure all servers are cleaned up
        if (primarySessionInfo.getConnHash() != null && !primarySessionInfo.getConnHash().isEmpty()) {
            connHashToServersMap.put(primarySessionInfo.getConnHash(), new ArrayList<>(connectedServers));
            log.info("Tracked {} servers for connection hash {}", connectedServers.size(), primarySessionInfo.getConnHash());
        }
        
        log.info("Connected to {} out of {} servers ({} connection)", 
                successfulConnections, serverEndpoints.size(), isXA ? "XA" : "non-XA");
        return primarySessionInfo;
    }
    
    
    /**
     * Gets the appropriate server based on session affinity.
     * 
     * If sessionKey is null, returns a round-robin selected server.
     * If sessionKey is not null, returns the server bound to that session.
     * 
     * @param sessionKey the session identifier (sessionUUID), or null for round-robin
     * @return the server endpoint to use
     * @throws SQLException if session exists but server is unavailable or not bound
     */
    public ServerEndpoint affinityServer(String sessionKey) throws SQLException {
        if (sessionKey == null || sessionKey.isEmpty()) {
            // No session identifier, use round-robin
            log.warn("DIAGNOSTIC: affinityServer called with NULL or EMPTY sessionKey - routing via round-robin. " +
                    "This may cause 'Connection not found' errors if queries reach wrong server. " +
                    "SessionKey value: '{}', isEmpty: {}, isNull: {}", 
                    sessionKey, 
                    sessionKey != null && sessionKey.isEmpty(),
                    sessionKey == null);
            return selectHealthyServer();
        }
        
        log.info("Looking up server for session: {}", sessionKey);
        ServerEndpoint sessionServer = sessionToServerMap.get(sessionKey);
        
        log.info("=== affinityServer lookup: sessionKey={}, found server={}, total sessions in map={} ===", 
                sessionKey, 
                sessionServer != null ? sessionServer.getAddress() : "NOT_FOUND",
                sessionToServerMap.size());
        
        // Log session distribution for debugging
        if (sessionServer != null && log.isDebugEnabled()) {
            Map<String, Long> serverDistribution = sessionToServerMap.values().stream()
                    .collect(java.util.stream.Collectors.groupingBy(
                            ServerEndpoint::getAddress,
                            java.util.stream.Collectors.counting()));
            log.debug("Session distribution across servers: {}", serverDistribution);
        }
        
        // Session must be bound - throw exception if not found
        if (sessionServer == null) {
            log.error("Session {} has no associated server. Available sessions: {}. This indicates the session binding was lost.", 
                    sessionKey, sessionToServerMap.keySet());
            throw new SQLException("Session " + sessionKey + 
                    " has no associated server. Session may have expired or server may be unavailable. " +
                    "Available bound sessions: " + sessionToServerMap.keySet());
        }
        
        log.info("Session {} is bound to server {}", sessionKey, sessionServer.getAddress());
        
        if (!sessionServer.isHealthy()) {
            // Remove from map and throw exception - do NOT fall back to round-robin
            sessionToServerMap.remove(sessionKey);
            throw new SQLException("Session " + sessionKey + 
                    " is bound to server " + sessionServer.getAddress() + 
                    " which is currently unavailable. Cannot continue with this session.");
        }
        
        return sessionServer;
    }
    
    
    /**
     * Gets the channel and stub for a specific server.
     */
    public ChannelAndStub getChannelAndStub(ServerEndpoint endpoint) {
        ChannelAndStub channelAndStub = channelMap.get(endpoint);
        if (channelAndStub == null) {
            try {
                channelAndStub = createChannelAndStub(endpoint);
            } catch (Exception e) {
                log.error("Failed to create channel for server {}: {}", endpoint.getAddress(), e.getMessage());
                return null;
            }
        }
        return channelAndStub;
    }
    
    private ServerEndpoint selectHealthyServer() {
        List<ServerEndpoint> healthyServers = serverEndpoints.stream()
                .filter(ServerEndpoint::isHealthy)
                .collect(Collectors.toList());
        
        // Only attempt recovery if NO servers are healthy (last resort)
        // Time-based health checks via tryTriggerHealthCheck() handle recovery for normal cases
        if (healthyServers.isEmpty()) {
            log.warn("No healthy servers available, attempting recovery as last resort");
            attemptServerRecovery();
            
            // Re-check healthy servers after recovery attempt
            healthyServers = serverEndpoints.stream()
                    .filter(ServerEndpoint::isHealthy)
                    .collect(Collectors.toList());
        }
        
        if (healthyServers.isEmpty()) {
            log.error("No healthy servers available after recovery attempt");
            return null;
        }
        
        // Choose selection strategy based on configuration
        if (healthCheckConfig.isLoadAwareSelectionEnabled()) {
            return selectByLeastConnections(healthyServers);
        } else {
            return selectByRoundRobin(healthyServers);
        }
    }
    
    /**
     * Selects the healthy server with the fewest active sessions.
     * This provides automatic load balancing by directing new connections
     * to the least-loaded server.
     * 
     * Uses SessionTracker for accurate session counts across both XA and non-XA connections.
     * When all servers have equal session counts, falls back to round-robin selection.
     * 
     * @param healthyServers List of healthy servers to choose from
     * @return The server with the lowest session count
     */
    private ServerEndpoint selectByLeastConnections(List<ServerEndpoint> healthyServers) {
        if (healthyServers.isEmpty()) {
            return null;
        }
        
        // Get current session counts per server from SessionTracker
        Map<ServerEndpoint, Integer> sessionCounts = sessionTracker.getSessionCounts();
