/*
    * Licensed to the Apache Software Foundation (ASF) under one or more
    * contributor license agreements.  See the NOTICE file distributed with
    * this work for additional information regarding copyright ownership.
    * The ASF licenses this file to You under the Apache License, Version 2.0
    * (the "License"); you may not use this file except in compliance with
    * the License.  You may obtain a copy of the License at
    *
    *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package org.apache.shardingsphere.proxy.backend.handler.transaction;
  
  import lombok.RequiredArgsConstructor;
  import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
  import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
  import org.apache.shardingsphere.infra.hint.HintValueContext;
  import org.apache.shardingsphere.infra.session.query.QueryContext;
  import org.apache.shardingsphere.proxy.backend.connector.DatabaseConnector;
  import org.apache.shardingsphere.proxy.backend.connector.DatabaseConnectorFactory;
  import org.apache.shardingsphere.proxy.backend.handler.ProxyBackendHandler;
  import org.apache.shardingsphere.proxy.backend.response.data.QueryResponseRow;
  import org.apache.shardingsphere.proxy.backend.response.header.ResponseHeader;
  import org.apache.shardingsphere.proxy.backend.session.ConnectionSession;
  import org.apache.shardingsphere.proxy.backend.util.TransactionUtils;
  import org.apache.shardingsphere.sql.parser.sql.common.statement.tcl.xa.XABeginStatement;
  import org.apache.shardingsphere.sql.parser.sql.common.statement.tcl.xa.XACommitStatement;
  import org.apache.shardingsphere.sql.parser.sql.common.statement.tcl.xa.XARecoveryStatement;
  import org.apache.shardingsphere.sql.parser.sql.common.statement.tcl.xa.XARollbackStatement;
  import org.apache.shardingsphere.sql.parser.sql.common.statement.tcl.xa.XAStatement;
  import org.apache.shardingsphere.transaction.api.TransactionType;
  import org.apache.shardingsphere.transaction.xa.jta.exception.XATransactionNestedBeginException;
  
  import java.sql.SQLException;
  import java.util.Collections;
  
  /**
   * XA transaction handler.
   */
  // TODO Currently XA transaction started with `XA START` doesn't support for database with multiple datasource, a flag should be added for this both in init progress and add datasource from DistSQL.
  @RequiredArgsConstructor
  public final class TransactionXAHandler implements ProxyBackendHandler {
      
      private final XAStatement xaStatement;
      
      private final ConnectionSession connectionSession;
      
      private final DatabaseConnector backendHandler;
      
      public TransactionXAHandler(final SQLStatementContext sqlStatementContext, final String sql, final ConnectionSession connectionSession) {
          xaStatement = (XAStatement) sqlStatementContext.getSqlStatement();
          this.connectionSession = connectionSession;
          backendHandler = DatabaseConnectorFactory.getInstance().newInstance(
                  new QueryContext(sqlStatementContext, sql, Collections.emptyList(), new HintValueContext()), connectionSession.getDatabaseConnectionManager(), false);
      }
      
      @Override
      public boolean next() throws SQLException {
          return xaStatement instanceof XARecoveryStatement && backendHandler.next();
      }
      
      @Override
      public QueryResponseRow getRowData() throws SQLException {
          return xaStatement instanceof XARecoveryStatement ? backendHandler.getRowData() : new QueryResponseRow(Collections.emptyList());
      }
      
      @Override
      public ResponseHeader execute() throws SQLException {
          if (xaStatement instanceof XABeginStatement) {
              return begin();
          }
          if (xaStatement instanceof XACommitStatement || xaStatement instanceof XARollbackStatement) {
              return finish();
          }
          return backendHandler.execute();
      }
      
      /*
       * We have to let session occupy the thread when doing xa transaction. According to https://dev.mysql.com/doc/refman/5.7/en/xa-states.html XA and local transactions are mutually exclusive.
       */
      private ResponseHeader begin() throws SQLException {
          ShardingSpherePreconditions.checkState(!connectionSession.getTransactionStatus().isInTransaction(), XATransactionNestedBeginException::new);
          ResponseHeader result = backendHandler.execute();
          TransactionType transactionType = TransactionUtils.getTransactionType(connectionSession.getConnectionContext().getTransactionContext());
          connectionSession.getConnectionContext().getTransactionContext().beginTransaction(String.valueOf(transactionType));
          return result;
      }
      
      private ResponseHeader finish() throws SQLException {
          try {
              return backendHandler.execute();
          } finally {
              connectionSession.getConnectionContext().clearTransactionContext();
             connectionSession.getConnectionContext().clearCursorContext();
         }
     }
 }