package org.openjproxy.grpc.client;

import com.openjproxy.grpc.CallResourceRequest;
import com.openjproxy.grpc.CallResourceResponse;
import com.openjproxy.grpc.ConnectionDetails;
import com.openjproxy.grpc.LobDataBlock;
import com.openjproxy.grpc.LobReference;
import com.openjproxy.grpc.OpResult;
import com.openjproxy.grpc.SessionInfo;
import com.openjproxy.grpc.StatementServiceGrpc;
import io.grpc.StatusRuntimeException;
import org.openjproxy.grpc.dto.Parameter;
import org.openjproxy.jdbc.Connection;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.SQLException;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Multinode implementation of StatementService that provides:
 * - Round-robin load balancing across multiple OJP servers
 * - Session stickiness (once a session is bound to a server, all requests for that session go to the same server)
 * - Automatic failover on connection-level errors
 * - Thread-safe concurrent request handling
 * 
 * This class delegates to StatementServiceGrpcClient instances, one per server endpoint.
 */
public class MultinodeStatementService implements StatementService {
    
    private static final Logger log = LoggerFactory.getLogger(MultinodeStatementService.class);
    
    private final MultinodeConnectionManager connectionManager;
    private final Map<ServerEndpoint, StatementServiceGrpcClient> clientMap;
    private final String originalUrl;
    
    /**
     * Creates a new MultinodeStatementService.
     * 
     * @param connectionManager The connection manager to use for server selection and session tracking
     * @param originalUrl The original multinode URL (for logging and debugging)
     */
    public MultinodeStatementService(MultinodeConnectionManager connectionManager, String originalUrl) {
        this.connectionManager = connectionManager;
        this.clientMap = new ConcurrentHashMap<>();
        this.originalUrl = originalUrl;
        
        log.info("MultinodeStatementService initialized with {} servers from URL: {}", 
                connectionManager.getServerEndpoints().size(), originalUrl);
    }
    
    /**
     * Gets or creates a StatementServiceGrpcClient for a specific server endpoint.
     * The client uses the same gRPC channel and stubs as the MultinodeConnectionManager
     * to ensure session continuity.
     * 
     * This method checks if the cached client is still using the current channel/stubs
     * from the connection manager. If not (e.g., after server recovery), it creates
     * a new client with the updated stubs.
     */
    private StatementServiceGrpcClient getClient(ServerEndpoint endpoint) {
        // Get the current channel and stubs from the connection manager
        MultinodeConnectionManager.ChannelAndStub currentChannelAndStub = 
                connectionManager.getChannelAndStub(endpoint);
        
        if (currentChannelAndStub == null) {
            log.error("Unable to get channel and stub for endpoint: {}", endpoint.getAddress());
            throw new RuntimeException("Unable to initialize client for endpoint: " + endpoint.getAddress());
        }
        
        // Check if we have a cached client
        StatementServiceGrpcClient cachedClient = clientMap.get(endpoint);
        
        if (cachedClient != null) {
            try {
                // Check if the cached client is still using the current stubs
                java.lang.reflect.Field blockingStubField = StatementServiceGrpcClient.class
                    .getDeclaredField("statemetServiceBlockingStub");
                blockingStubField.setAccessible(true);
                StatementServiceGrpc.StatementServiceBlockingStub cachedStub = 
                    (StatementServiceGrpc.StatementServiceBlockingStub) blockingStubField.get(cachedClient);
                
                // If the stub is the same as the current one, reuse the client
                if (cachedStub == currentChannelAndStub.blockingStub) {
                    return cachedClient;
                }
                
                // The stubs have changed (server was recovered), remove the old client
                log.info("Stubs changed for endpoint {} (likely after recovery), creating new client", 
                        endpoint.getAddress());
                clientMap.remove(endpoint);
                
            } catch (Exception e) {
                log.warn("Error checking cached client stubs for {}: {}", endpoint.getAddress(), e.getMessage());
                // If we can't check, remove the cached client to be safe
                clientMap.remove(endpoint);
            }
        }
        
        // Create a new client
        log.info("Creating new StatementServiceGrpcClient for endpoint: {}", endpoint.getAddress());
        
        log.info("Got channel and stub for endpoint {}: blockingStub={}, asyncStub={}", 
            endpoint.getAddress(), 
            System.identityHashCode(currentChannelAndStub.blockingStub),
            System.identityHashCode(currentChannelAndStub.asyncStub));
        
        // Create a client and inject the connection manager's stubs into it
        // This ensures all operations use the same gRPC connection where the session was created
        StatementServiceGrpcClient client = new StatementServiceGrpcClient();
        
        try {
            // Use reflection to set the stubs to the connection manager's stubs
            // This is necessary because StatementServiceGrpcClient doesn't have a constructor
            // that accepts stubs, and we need to use the same channel for session continuity
            java.lang.reflect.Field blockingStubField = StatementServiceGrpcClient.class
                .getDeclaredField("statemetServiceBlockingStub");
            blockingStubField.setAccessible(true);
            blockingStubField.set(client, currentChannelAndStub.blockingStub);
            
            java.lang.reflect.Field asyncStubField = StatementServiceGrpcClient.class
                .getDeclaredField("statemetServiceStub");
            asyncStubField.setAccessible(true);
            asyncStubField.set(client, currentChannelAndStub.asyncStub);
            
            log.info("Initialized StatementServiceGrpcClient with connection manager's stubs for endpoint: {}", 
                endpoint.getAddress());
        } catch (Exception e) {
            log.error("Failed to initialize client for endpoint {}: {}", endpoint.getAddress(), e.getMessage(), e);
            throw new RuntimeException("Failed to initialize client for endpoint: " + endpoint.getAddress(), e);
        }
        
        // Cache the new client
        clientMap.put(endpoint, client);
        
        return client;
    }
    
    /**
     * Enhances SessionInfo with current cluster health status.
     * Creates a new SessionInfo with the clusterHealth field populated based on
     * the current health status of all servers in the cluster.
     * 
     * @param sessionInfo The original session info
     * @return A new SessionInfo with cluster health populated
     */
    private SessionInfo withClusterHealth(SessionInfo sessionInfo) {
        if (sessionInfo == null) {
            return null;
        }
        
        String clusterHealth = connectionManager.generateClusterHealth();
        
        log.info("[XA-REBALANCE-DEBUG] withClusterHealth called: connHash={}, isXA={}, original clusterHealth={}, new clusterHealth={}", 
                sessionInfo.getConnHash(), sessionInfo.getIsXA(), sessionInfo.getClusterHealth(), clusterHealth);
        
        SessionInfo enhanced = SessionInfo.newBuilder(sessionInfo)
                .setClusterHealth(clusterHealth)
                .build();
        
        return enhanced;
    }
    
    /**
     * Checks if a session was created (sessionUUID went from empty to non-empty) and binds it to the server.
     * Also ensures existing sessions remain bound to the correct server.
     * This is called after operations that may create a session (e.g., startTransaction).
     * 
     * @param requestSessionInfo The SessionInfo sent in the request
     * @param responseSessionInfo The SessionInfo received in the response
     * @param server The server that handled the request
     */
    private void checkAndBindSession(SessionInfo requestSessionInfo, SessionInfo responseSessionInfo, ServerEndpoint server) {
        // Check if session was created in this call (request had no UUID, response has UUID)
        boolean requestHadNoSession = requestSessionInfo == null || 
                                      requestSessionInfo.getSessionUUID() == null || 
                                      requestSessionInfo.getSessionUUID().isEmpty();
        boolean responseHasSession = responseSessionInfo != null && 
                                      responseSessionInfo.getSessionUUID() != null && 
                                      !responseSessionInfo.getSessionUUID().isEmpty();
        
        if (responseHasSession) {
            String sessionUUID = responseSessionInfo.getSessionUUID();
            String targetServer = responseSessionInfo.getTargetServer();
            
            // Check if session is already bound
            String currentBinding = connectionManager.getBoundTargetServer(sessionUUID);
            
            if (currentBinding == null) {
                // Session not bound yet - bind it now
                if (targetServer != null && !targetServer.isEmpty()) {
                    connectionManager.bindSession(sessionUUID, targetServer);
                    log.info("Session {} bound to target server {} (was unbound)", sessionUUID, targetServer);
                } else {
                    // Fallback: use server address if targetServer not provided
                    String serverAddress = server.getHost() + ":" + server.getPort();
                    connectionManager.bindSession(sessionUUID, serverAddress);
                    log.warn("Session {} bound to server {} (no targetServer in response, was unbound)", 
                            sessionUUID, serverAddress);
                }
            } else {
                // Session already bound - verify it matches
                String expectedServer = (targetServer != null && !targetServer.isEmpty()) 
                    ? targetServer 
                    : (server.getHost() + ":" + server.getPort());
                
                if (!currentBinding.equals(expectedServer)) {
                    log.error("Session {} binding mismatch - currently bound to {}, response indicates {}. Re-binding to response server.", 
                            sessionUUID, currentBinding, expectedServer);
                    // Re-bind to the server that actually handled the request
                    if (targetServer != null && !targetServer.isEmpty()) {
                        connectionManager.bindSession(sessionUUID, targetServer);
                    } else {
                        connectionManager.bindSession(sessionUUID, expectedServer);
                    }
                } else {
                    log.debug("Session {} binding verified: {}", sessionUUID, currentBinding);
                }
            }
        }
    }
    
    /**
     * Executes an operation that returns OpResult with session stickiness and binding check.
     * This wrapper handles binding newly-created sessions to servers.
     * 
     * @param requestSessionInfo The session info for determining which server to use
     * @param operation The operation to execute
     * @return The OpResult result
     * @throws SQLException if the operation fails
     */
    private OpResult executeOpResultWithSessionStickinessAndBinding(SessionInfo requestSessionInfo, 
                                                                      ThrowingFunction<StatementServiceGrpcClient, OpResult> operation) 
            throws SQLException {
        // Get the appropriate server based on session binding or round-robin
        String sessionKey = (requestSessionInfo != null && requestSessionInfo.getSessionUUID() != null && !requestSessionInfo.getSessionUUID().isEmpty()) 
                ? requestSessionInfo.getSessionUUID() : null;
        ServerEndpoint server = connectionManager.affinityServer(sessionKey);
        
        log.debug("executeOpResultWithSessionStickinessAndBinding: session={}, server={}", 
            requestSessionInfo != null ? requestSessionInfo.getSessionUUID() : "null", 
            server != null ? server.getAddress() : "null");
        
        try {
            // Get the channel and stub for the selected server
            MultinodeConnectionManager.ChannelAndStub channelAndStub = 
                    connectionManager.getChannelAndStub(server);
            
            if (channelAndStub == null) {
                throw new SQLException("Unable to get channel for server: " + server.getAddress());
            }
            
            // Get or create the client for this endpoint
            StatementServiceGrpcClient client = getClient(server);
            
            // Execute the operation
            OpResult result = operation.apply(client);
            
            // Check if result contains a session and bind it
            if (result != null && result.hasSession()) {
                SessionInfo responseSessionInfo = result.getSession();
                checkAndBindSession(requestSessionInfo, responseSessionInfo, server);
            }
            
            return result;
            
        } catch (StatusRuntimeException e) {
            // Let GrpcExceptionHandler convert the exception
            SQLException sqlEx;
            try {
                throw GrpcExceptionHandler.handle(e);
            } catch (SQLException ex) {
                sqlEx = ex;
            }
            
            // Only mark server unhealthy for connection-level errors
            if (connectionManager.isConnectionLevelError(e)) {
                log.warn("Connection-level error on server {}: {}", server.getAddress(), sqlEx.getMessage());
                // Notify connection manager to mark server unhealthy and invalidate sessions/connections
                connectionManager.handleServerFailure(server, e);
            } else {
                log.debug("Database-level error on server {}: {}", server.getAddress(), sqlEx.getMessage());
            }
            
            throw sqlEx;
        } catch (SQLException e) {
            throw e;
        } catch (Exception e) {
            throw new SQLException("Unexpected error executing operation: " + e.getMessage(), e);
        }
    }
    
    /**
     * Executes an operation that returns Iterator<OpResult> with session stickiness and binding check.
     * This wrapper handles binding newly-created sessions to servers by checking the first result.
     * 
     * @param requestSessionInfo The session info for determining which server to use
     * @param operation The operation to execute
     * @return The Iterator<OpResult> result
     * @throws SQLException if the operation fails
     */
    private Iterator<OpResult> executeIteratorWithSessionStickinessAndBinding(SessionInfo requestSessionInfo, 
                                                                                ThrowingFunction<StatementServiceGrpcClient, Iterator<OpResult>> operation) 
            throws SQLException {
        // Get the appropriate server based on session binding or round-robin
        String sessionKey = (requestSessionInfo != null && requestSessionInfo.getSessionUUID() != null && !requestSessionInfo.getSessionUUID().isEmpty()) 
                ? requestSessionInfo.getSessionUUID() : null;
        ServerEndpoint server = connectionManager.affinityServer(sessionKey);
        
        log.debug("executeIteratorWithSessionStickinessAndBinding: session={}, server={}", 
            requestSessionInfo != null ? requestSessionInfo.getSessionUUID() : "null", 
            server != null ? server.getAddress() : "null");
        
        try {
            // Get the channel and stub for the selected server
            MultinodeConnectionManager.ChannelAndStub channelAndStub = 
                    connectionManager.getChannelAndStub(server);
            
            if (channelAndStub == null) {
                throw new SQLException("Unable to get channel for server: " + server.getAddress());
            }
            
            // Get or create the client for this endpoint
            StatementServiceGrpcClient client = getClient(server);
            
            // Execute the operation
            Iterator<OpResult> resultIterator = operation.apply(client);
            
            // Wrap the iterator to check and bind session from the first result
            return new Iterator<OpResult>() {
                private boolean firstResultProcessed = false;
                
                @Override
                public boolean hasNext() {
                    return resultIterator.hasNext();
                }
                
                @Override
                public OpResult next() {
                    OpResult result = resultIterator.next();
                    
                    // Check and bind session from the first result
                    if (!firstResultProcessed && result != null && result.hasSession()) {
                        SessionInfo responseSessionInfo = result.getSession();
                        try {
                            checkAndBindSession(requestSessionInfo, responseSessionInfo, server);
                        } catch (Exception e) {
                            log.warn("Failed to check/bind session from query result: {}", e.getMessage());
                        }
                        firstResultProcessed = true;
                    }
                    
                    return result;
                }
            };
            
        } catch (StatusRuntimeException e) {
            // Let GrpcExceptionHandler convert the exception
            SQLException sqlEx;
            try {
                throw GrpcExceptionHandler.handle(e);
            } catch (SQLException ex) {
                sqlEx = ex;
            }
            
            // Only mark server unhealthy for connection-level errors
            if (connectionManager.isConnectionLevelError(e)) {
                log.warn("Connection-level error on server {}: {}", server.getAddress(), sqlEx.getMessage());
                // Notify connection manager to mark server unhealthy and invalidate sessions/connections
                connectionManager.handleServerFailure(server, e);
            } else {
                log.debug("Database-level error on server {}: {}", server.getAddress(), sqlEx.getMessage());
            }
            
            throw sqlEx;
        } catch (SQLException e) {
            throw e;
        } catch (Exception e) {
            throw new SQLException("Unexpected error executing operation: " + e.getMessage(), e);
        }
    }
    
    /**
     * Executes an operation that returns SessionInfo with session stickiness and binding check.
     * This wrapper handles binding newly-created sessions to servers.
     * 
     * @param requestSessionInfo The session info for determining which server to use
     * @param operation The operation to execute
     * @return The SessionInfo result
     * @throws SQLException if the operation fails
     */
    private SessionInfo executeWithSessionStickinessAndBinding(SessionInfo requestSessionInfo, 
                                                                ThrowingFunction<StatementServiceGrpcClient, SessionInfo> operation) 
            throws SQLException {
        // Get the appropriate server based on session binding or round-robin
        String sessionKey = (requestSessionInfo != null && requestSessionInfo.getSessionUUID() != null && !requestSessionInfo.getSessionUUID().isEmpty()) 
                ? requestSessionInfo.getSessionUUID() : null;
        ServerEndpoint server = connectionManager.affinityServer(sessionKey);
        
        log.info("executeWithSessionStickinessAndBinding: session={}, server={}", 
            requestSessionInfo != null ? requestSessionInfo.getSessionUUID() : "null", 
            server != null ? server.getAddress() : "null");
        
        try {
            // Get the channel and stub for the selected server
            MultinodeConnectionManager.ChannelAndStub channelAndStub = 
                    connectionManager.getChannelAndStub(server);
            
            if (channelAndStub == null) {
                throw new SQLException("Unable to get channel for server: " + server.getAddress());
            }
            
            // Get or create the client for this endpoint
            StatementServiceGrpcClient client = getClient(server);
            
            // Execute the operation
            SessionInfo responseSessionInfo = operation.apply(client);
            
            // Check if a session was created and bind it
            checkAndBindSession(requestSessionInfo, responseSessionInfo, server);
            
            return responseSessionInfo;
            
        } catch (StatusRuntimeException e) {
            // Let GrpcExceptionHandler convert the exception
            SQLException sqlEx;
            try {
                throw GrpcExceptionHandler.handle(e);
            } catch (SQLException ex) {
                sqlEx = ex;
            }

            // Only mark server unhealthy for connection-level errors
            if (connectionManager.isConnectionLevelError(e)) {
                log.warn("Connection-level error on server {}: {}", server.getAddress(), sqlEx.getMessage());
                // Notify connection manager to mark server unhealthy and invalidate sessions/connections
                connectionManager.handleServerFailure(server, e);
            } else {
                log.debug("Database-level error on server {}: {}", server.getAddress(), sqlEx.getMessage());
            }
            
            throw sqlEx;
            
        } catch (SQLException e) {
            throw e;
        } catch (Exception e) {
            throw new SQLException("Unexpected error executing operation on server " + 
                    server.getAddress() + ": " + e.getMessage(), e);
        }
    }
    
    @Override
    public SessionInfo connect(ConnectionDetails connectionDetails) throws SQLException {
        // Use the connection manager to handle connection with retry and failover logic
        return connectionManager.connect(connectionDetails);
    }
    
    @Override
    public OpResult executeUpdate(SessionInfo sessionInfo, String sql, List<Parameter> params, 
                                  Map<String, Object> properties) throws SQLException {
        return executeUpdate(sessionInfo, sql, params, "", properties);
    }
    
    @Override
    public OpResult executeUpdate(SessionInfo sessionInfo, String sql, List<Parameter> params, 
                                  String statementUUID, Map<String, Object> properties) throws SQLException {
        SessionInfo enhancedSessionInfo = withClusterHealth(sessionInfo);
        return executeOpResultWithSessionStickinessAndBinding(enhancedSessionInfo, client -> 
            client.executeUpdate(enhancedSessionInfo, sql, params, statementUUID, properties)
        );
    }
    
    @Override
    public Iterator<OpResult> executeQuery(SessionInfo sessionInfo, String sql, List<Parameter> params, 
                                           Map<String, Object> properties) throws SQLException {
        return executeQuery(sessionInfo, sql, params, "", properties);
    }
    
    @Override
    public Iterator<OpResult> executeQuery(SessionInfo sessionInfo, String sql, List<Parameter> params, 
                                           String statementUUID, Map<String, Object> properties) throws SQLException {
        // For executeQuery, we execute with binding check and wrap the iterator to check subsequent results
        SessionInfo enhancedSessionInfo = withClusterHealth(sessionInfo);
        return executeIteratorWithSessionStickinessAndBinding(enhancedSessionInfo, client -> 
            client.executeQuery(enhancedSessionInfo, sql, params, statementUUID, properties)
        );
    }
    
    @Override
    public OpResult fetchNextRows(SessionInfo sessionInfo, String resultSetUUID, int size) throws SQLException {
        SessionInfo enhancedSessionInfo = withClusterHealth(sessionInfo);
        return executeWithSessionStickiness(enhancedSessionInfo, client -> 
            client.fetchNextRows(enhancedSessionInfo, resultSetUUID, size)
        );
    }
    
    @Override
    public LobReference createLob(Connection connection, Iterator<LobDataBlock> lobDataBlock) throws SQLException {
        SessionInfo sessionInfo = connection.getSession();
        return executeWithSessionStickiness(sessionInfo, client -> 
            client.createLob(connection, lobDataBlock)
        );
    }
    
    @Override
    public Iterator<LobDataBlock> readLob(LobReference lobReference, long pos, int length) throws SQLException {
        SessionInfo sessionInfo = lobReference.getSession();
        return executeWithSessionStickiness(sessionInfo, client -> 
            client.readLob(lobReference, pos, length)
        );
    }
    
    @Override
    public void terminateSession(SessionInfo session) {
        try {
            // For sessions with a UUID, terminate only on the bound server
            if (session != null && session.getSessionUUID() != null && !session.getSessionUUID().isEmpty()) {
                String boundServer = connectionManager.getBoundTargetServer(session.getSessionUUID());
                if (boundServer != null) {
                    // Session is bound - terminate on that specific server
                    log.info("Terminating session {} on bound server {}", session.getSessionUUID(), boundServer);
                    try {
                        ServerEndpoint server = connectionManager.affinityServer(session.getSessionUUID());
                        StatementServiceGrpcClient client = getClient(server);
                        client.terminateSession(session);
                        log.info("Successfully terminated session {} on server {}", 
                                session.getSessionUUID(), server.getAddress());
                    } catch (SQLException e) {
                        log.warn("Error terminating session {} on bound server: {}", 
                                session.getSessionUUID(), e.getMessage());
                    }
                } else {
                    // Session not bound - try all servers that received connect()
                    log.info("Session {} not bound, attempting termination on all connected servers", 
                            session.getSessionUUID());
                    List<ServerEndpoint> serversToTerminate = connectionManager.getServersForConnHash(
                            session.getConnHash());
                    if (serversToTerminate != null && !serversToTerminate.isEmpty()) {
                        for (ServerEndpoint server : serversToTerminate) {
                            try {
                                StatementServiceGrpcClient client = getClient(server);
                                client.terminateSession(session);
                                log.debug("Terminated session on server {}", server.getAddress());
                            } catch (Exception e) {
                                log.warn("Error terminating session on server {}: {}", 
                                        server.getAddress(), e.getMessage());
                            }
                        }
                    }
                }
            } else {
                // No session UUID - try terminating on all servers that received connect()
                log.info("No sessionUUID, attempting termination on all connected servers");
                List<ServerEndpoint> serversToTerminate = connectionManager.getServersForConnHash(
                        session != null ? session.getConnHash() : null);
                if (serversToTerminate != null && !serversToTerminate.isEmpty()) {
                    for (ServerEndpoint server : serversToTerminate) {
                        try {
                            StatementServiceGrpcClient client = getClient(server);
                            client.terminateSession(session);
                            log.debug("Terminated session on server {}", server.getAddress());
                        } catch (Exception e) {
                            log.warn("Error terminating session on server {}: {}", 
                                    server.getAddress(), e.getMessage());
                        }
                    }
                }
            }
            
            // Clean up connection manager's session tracking
            connectionManager.terminateSession(session);
        } catch (Exception e) {
            log.warn("Error terminating session {}: {}", 
                    session != null ? session.getSessionUUID() : "null", e.getMessage());
            // Best effort - don't throw exception for terminate
        }
    }
    
    @Override
    public SessionInfo startTransaction(SessionInfo session) throws SQLException {
        SessionInfo enhancedSessionInfo = withClusterHealth(session);
        return executeWithSessionStickinessAndBinding(enhancedSessionInfo, client -> 
            client.startTransaction(enhancedSessionInfo)
        );
    }
    
    @Override
    public SessionInfo commitTransaction(SessionInfo session) throws SQLException {
        SessionInfo enhancedSessionInfo = withClusterHealth(session);
        return executeWithSessionStickinessAndBinding(enhancedSessionInfo, client -> 
            client.commitTransaction(enhancedSessionInfo)
        );
    }
    
    @Override
    public SessionInfo rollbackTransaction(SessionInfo session) throws SQLException {
        SessionInfo enhancedSessionInfo = withClusterHealth(session);
        return executeWithSessionStickinessAndBinding(enhancedSessionInfo, client -> 
            client.rollbackTransaction(enhancedSessionInfo)
        );
    }
    
    @Override
    public CallResourceResponse callResource(CallResourceRequest request) throws SQLException {
        SessionInfo sessionInfo = request.getSession();
        SessionInfo enhancedSessionInfo = withClusterHealth(sessionInfo);
        CallResourceRequest enhancedRequest = CallResourceRequest.newBuilder(request)
                .setSession(enhancedSessionInfo)
                .build();
        return executeWithSessionStickiness(enhancedSessionInfo, client -> 
            client.callResource(enhancedRequest)
        );
    }
    
    // XA Transaction Operations
    @Override
    public com.openjproxy.grpc.XaResponse xaStart(com.openjproxy.grpc.XaStartRequest request) throws SQLException {
        SessionInfo sessionInfo = request.getSession();
        SessionInfo enhancedSessionInfo = withClusterHealth(sessionInfo);
        com.openjproxy.grpc.XaStartRequest enhancedRequest = com.openjproxy.grpc.XaStartRequest.newBuilder(request)
                .setSession(enhancedSessionInfo)
                .build();
        return executeWithSessionStickiness(enhancedSessionInfo, client -> 
            client.xaStart(enhancedRequest)
        );
    }
    
    @Override
    public com.openjproxy.grpc.XaResponse xaEnd(com.openjproxy.grpc.XaEndRequest request) throws SQLException {
        SessionInfo sessionInfo = request.getSession();
        SessionInfo enhancedSessionInfo = withClusterHealth(sessionInfo);
        com.openjproxy.grpc.XaEndRequest enhancedRequest = com.openjproxy.grpc.XaEndRequest.newBuilder(request)
                .setSession(enhancedSessionInfo)
                .build();
        return executeWithSessionStickiness(enhancedSessionInfo, client -> 
            client.xaEnd(enhancedRequest)
        );
    }
    
    @Override
    public com.openjproxy.grpc.XaPrepareResponse xaPrepare(com.openjproxy.grpc.XaPrepareRequest request) throws SQLException {
        SessionInfo sessionInfo = request.getSession();
        SessionInfo enhancedSessionInfo = withClusterHealth(sessionInfo);
        com.openjproxy.grpc.XaPrepareRequest enhancedRequest = com.openjproxy.grpc.XaPrepareRequest.newBuilder(request)
                .setSession(enhancedSessionInfo)
                .build();
        return executeWithSessionStickiness(enhancedSessionInfo, client -> 
