#include <postgres.h>
#include <utils/elog.h>
#include <access/xact.h>

/** eternal variables */
extern bool dml_in_transaction;

/** external prototypes */
extern void         db2EndTransaction         (void* arg, int is_commit, int noerror);
extern void         db2EndSubtransaction      (void* arg, int nest_level, int is_commit);
extern void         db2Debug1                 (const char* message, ...);

/** local prototypes */
void db2RegisterCallback   (void* arg);
void db2UnregisterCallback (void* arg);
void transactionCallback   (XactEvent event, void *arg);
void subtransactionCallback(SubXactEvent event, SubTransactionId mySubid, SubTransactionId parentSubid, void* arg);

/** db2RegisterCallback
 *   Register a callback for PostgreSQL transaction events.
 */
void db2RegisterCallback (void *arg) {
  RegisterXactCallback (transactionCallback, arg);
  RegisterSubXactCallback (subtransactionCallback, arg);
}

/** db2UnregisterCallback
 *   Unregister a callback for PostgreSQL transaction events.
 */
void db2UnregisterCallback (void *arg) {
  UnregisterXactCallback (transactionCallback, arg);
  UnregisterSubXactCallback (subtransactionCallback, arg);
}

/** transactionCallback
 *   Commit or rollback DB2 transactions when appropriate.
 */
void transactionCallback (XactEvent event, void *arg) {
  db2Debug1("> transactionCallback");
  switch (event) {
    case XACT_EVENT_PRE_COMMIT:
    case XACT_EVENT_PARALLEL_PRE_COMMIT:
      /* remote commit */
      db2EndTransaction (arg, 1, 0);
    break;
    case XACT_EVENT_PRE_PREPARE:
      ereport (ERROR, (errcode (ERRCODE_FDW_UNABLE_TO_CREATE_EXECUTION), errmsg ("cannot prepare a transaction that used remote tables")));
    break;
    case XACT_EVENT_COMMIT:
    case XACT_EVENT_PREPARE:
    case XACT_EVENT_PARALLEL_COMMIT:
      /*
       * Commit the remote transaction ignoring errors.
       * In 9.3 or higher, the transaction must already be closed, so this does nothing.
       * In 9.2 or lower, this is ok since nothing can have been modified remotely.
       */
      db2EndTransaction (arg, 1, 1);
    break;
    case XACT_EVENT_ABORT:
    case XACT_EVENT_PARALLEL_ABORT:
      /* remote rollback */
      db2EndTransaction (arg, 0, 1);
    break;
  }
  dml_in_transaction = false;
  db2Debug1("< transactionCallback");
}

/** subtransactionCallback
 *   Set or rollback to DB2 savepoints when appropriate.
 */
void subtransactionCallback (SubXactEvent event, SubTransactionId mySubid, SubTransactionId parentSubid, void *arg) {
    db2Debug1("> subtransactionCallback");
  /* rollback to the appropriate savepoint on subtransaction abort */
  if (event == SUBXACT_EVENT_ABORT_SUB || event == SUBXACT_EVENT_PRE_COMMIT_SUB)
    db2EndSubtransaction (arg, GetCurrentTransactionNestLevel (), event == SUBXACT_EVENT_PRE_COMMIT_SUB);
  db2Debug1("< subtransactionCallback");
}