#include <string.h>
#include <stdio.h>
#include <sqlcli1.h>
#include <postgres_ext.h>
#include "db2_fdw.h"

/** external variables */
extern char      db2Message[ERRBUFSIZE];/* contains DB2 error messages, set by db2CheckErr()             */

/** external prototypes */
extern void      db2Debug1            (const char* message, ...);
extern void      db2Debug2            (const char* message, ...);
extern void      db2Debug3            (const char* message, ...);
extern void      db2Error_d           (db2error sqlstate, const char* message, const char* detail, ...);
extern void      db2RegisterCallback  (void* arg);
extern SQLRETURN db2CheckErr          (SQLRETURN status, SQLHANDLE handle, SQLSMALLINT handleType, int line, char* file);
extern void      db2FreeEnvHdl        (DB2EnvEntry* envp, const char* nls_lang);
extern char*     db2strdup            (const char* p);

/** local prototypes */
DB2ConnEntry*    db2AllocConnHdl      (DB2EnvEntry* envp,const char* srvname, char* user, char* password, char* jwt_token, const char* nls_lang);
DB2ConnEntry*    findconnEntry        (DB2ConnEntry* start, const char* srvname, const char* user);
DB2ConnEntry*    insertconnEntry      (DB2ConnEntry* start, const char* srvname, const char* uid, const char* pwd, const char* jwt_token, SQLHDBC hdbc);

/** db2AllocConnHdl
 *
 */
DB2ConnEntry* db2AllocConnHdl(DB2EnvEntry* envp,const char* srvname, char* user, char* password, char* jwt_token, const char* nls_lang) {
  DB2ConnEntry* connp   = NULL;
  SQLRETURN     rc      = 0;
  SQLHDBC       hdbc    = SQL_NULL_HDBC;

  db2Debug1("> db2AllocConnHdl(envp: %x, srvname: %s, user: %s, password: %s, jwt_token: %s, nls_lang: %s)", envp, srvname,user, password, jwt_token ? "***" : "NULL", nls_lang);
  if (nls_lang != NULL) {
    rc = SQLAllocHandle(SQL_HANDLE_DBC, envp->henv, &hdbc);
    db2Debug3("  alloc dbc handle - rc: %d, henv: %d, hdbc: %d",rc, envp->henv, hdbc);
    rc = db2CheckErr(rc, hdbc, SQL_HANDLE_DBC, __LINE__, __FILE__);
    if (rc != SQL_SUCCESS) {
        db2Error_d (FDW_UNABLE_TO_ESTABLISH_CONNECTION, "error connecting to DB2: SQLAllocHandle failed to allocate hdbc handle", db2Message);
        db2FreeEnvHdl(envp,nls_lang);
        envp = NULL;
    }
    // envp->connlist = connp = insertconnEntry (envp->connlist, srvname, user, password, hdbc);
  } else {
    /* search user session for this server in cache */
    connp = findconnEntry(envp->connlist, srvname, user);
    if (connp == NULL) {
      /* Declare all variables at beginning for C90 compatibility */
      char connStr[4096];
      int connStrLen;
      SQLCHAR outConnStr[1024];
      SQLSMALLINT outConnStrLen;

      /* create connection handle */
      rc = SQLAllocHandle(SQL_HANDLE_DBC, envp->henv, &hdbc);
      db2Debug3("  alloc dbc handle - rc: %d, henv: %d, hdbc: %d",rc, envp->henv, hdbc);
      rc = db2CheckErr(rc, envp->henv, SQL_HANDLE_ENV, __LINE__, __FILE__);
      if (rc  != SQL_SUCCESS) {
        db2Error_d (FDW_UNABLE_TO_ESTABLISH_CONNECTION, "error connecting to DB2: SQLAllochHandle failed to allocate hdbc handle", db2Message);
      }

      /* Check if JWT token authentication is used */
      if (jwt_token != NULL && jwt_token[0] != '\0') {
        /* JWT token authentication */
        db2Debug1("  using JWT token authentication");

        /* For DB2 11.5.4+ with JWT, use SQLDriverConnect with AUTHENTICATION=TOKEN */
        /* Requires: DB2 client 11.5.4+, server configured with db2token.cfg */
        /* and SRVCON_AUTH set to SERVER_ENCRYPT_TOKEN or similar */

        /* Build connection string with JWT token using DB2 TOKEN authentication keywords */
        /* JWT tokens contain identity information, so UID is typically not required */
        connStrLen = snprintf(connStr, sizeof(connStr),
                             "DSN=%s;AUTHENTICATION=TOKEN;ACCESSTOKEN=%s;ACCESSTOKENTYPE=JWT;",
                             srvname, jwt_token);

        if (connStrLen >= sizeof(connStr)) {
          db2Error_d (FDW_UNABLE_TO_ESTABLISH_CONNECTION, "connection string too long", " connection to foreign DB2 server");
        }

        db2Debug1("  connecting with connection string (token hidden)");

        /* Use SQLDriverConnect instead of SQLConnect */
        rc = SQLDriverConnect(hdbc, NULL, (SQLCHAR*)connStr, SQL_NTS,
                             outConnStr, sizeof(outConnStr), &outConnStrLen,
                             SQL_DRIVER_NOPROMPT);

        db2Debug1("  connect to database(%s) with JWT token - rc: %d, hdbc: %d", srvname, rc, hdbc);
        rc = db2CheckErr(rc, hdbc, SQL_HANDLE_DBC, __LINE__, __FILE__);
        if (rc != SQL_SUCCESS) {
          db2Error_d (FDW_UNABLE_TO_ESTABLISH_CONNECTION, "cannot authenticate with JWT token", " connection connectstring: %s ,%s", srvname, db2Message);
          db2Error_d (FDW_UNABLE_TO_ESTABLISH_CONNECTION, "cannot authenticate", " connection to foreign DB2 server,%s", db2Message);
        }
      } else {
        /* Traditional user/password authentication */
        db2Debug1("  using user/password authentication");
        rc = SQLConnect(hdbc, (SQLCHAR*)srvname, SQL_NTS, (SQLCHAR*)user, SQL_NTS, (SQLCHAR*)password, SQL_NTS);
        db2Debug1("  connect to database(%s) - rc: %d, hdbc: %d",srvname, rc, hdbc);
        rc = db2CheckErr(rc, hdbc, SQL_HANDLE_DBC, __LINE__, __FILE__);
        if (rc != SQL_SUCCESS) {
          db2Error_d (FDW_UNABLE_TO_ESTABLISH_CONNECTION, "cannot authenticate"," connection User: %s ,%s"            , user    , db2Message);
          db2Error_d (FDW_UNABLE_TO_ESTABLISH_CONNECTION, "cannot authenticate"," connection password: %s ,%s"        , password, db2Message);
          db2Error_d (FDW_UNABLE_TO_ESTABLISH_CONNECTION, "cannot authenticate"," connection connectstring: %s ,%s"   , srvname , db2Message);
          db2Error_d (FDW_UNABLE_TO_ESTABLISH_CONNECTION, "cannot authenticate"," connection to foreign DB2 server,%s"          , db2Message);
        }
      }

      if (rc == SQL_SUCCESS) {
        /* add session handle to cache */
        envp->connlist = connp = insertconnEntry (envp->connlist, srvname, user, password, jwt_token, hdbc);

        /* set Autocommit off */
        rc = SQLSetConnectAttr(hdbc, SQL_ATTR_AUTOCOMMIT, (SQLPOINTER)SQL_AUTOCOMMIT_OFF, SQL_IS_UINTEGER);
        rc = db2CheckErr(rc, hdbc, SQL_HANDLE_DBC, __LINE__, __FILE__);
        if (rc != SQL_SUCCESS) {
          db2Error_d (FDW_UNABLE_TO_ESTABLISH_CONNECTION, "failed to set autocommit=off"," connection to foreign DB2 server,%s", db2Message);
        }
        /* register callback for PostgreSQL transaction events */
        db2RegisterCallback (connp);
      }
    }
  }
  db2Debug1("< db2AllocConnHdl - returns: %x",connp);
  return connp;
}

/** findconnEntry
 * 
 */
DB2ConnEntry* findconnEntry(DB2ConnEntry* start, const char* srvname, const char* user) {
  DB2ConnEntry* step = NULL;
  db2Debug2("  > findconnEntry");
  for (step = start; step != NULL; step = step->right){
    if (strcmp(step->srvname, srvname) == 0 && strcmp(step->uid, user) == 0) {
      break;
    }
  }
  db2Debug2("  < findconnEntry - returns: %x", step);
  return step;
}

/** insertconnEntry
 *
 */
DB2ConnEntry* insertconnEntry(DB2ConnEntry* start, const char* srvname, const char* uid, const char* pwd, const char* jwt_token, SQLHDBC hdbc) {
  DB2ConnEntry* step = NULL;
  DB2ConnEntry* new  = NULL;

  db2Debug2("  > insertconnEntry");
  new = malloc(sizeof(DB2ConnEntry));
  if (start == NULL){ /* first entry in list */
    new->right = new->left = NULL;
  } else {
    for (step = start; step->right != NULL; step = step->right){ }
    step->right = new;
    new->left = step;
    new->right = NULL;
  }
  // generate a deep copy
  new->srvname    = (srvname   && srvname[0]   != '\0') ? db2strdup(srvname)   : NULL;
  new->uid        = (uid       && uid[0]       != '\0') ? db2strdup(uid)       : NULL;
  new->pwd        = (pwd       && pwd[0]       != '\0') ? db2strdup(pwd)       : NULL;
  new->jwt_token  = (jwt_token && jwt_token[0] != '\0') ? db2strdup(jwt_token) : NULL;
  new->handlelist = NULL;
  new->hdbc       = hdbc;
  new->xact_level = 0;
  db2Debug2("  < insertconnEntry - returns: %x",new);
  return new;
}