#include <stdio.h>
#include <stdlib.h>
#include <sql.h>
#include <sqlext.h>
#include <string.h>
#include <assert.h>

// compilation: gcc -g pmodbc.c -lodbc -o pmodbc
// opt: gcc -O3 pmodbc.c -lodbc -o pmodbc

static void list_drivers();
static void query_db(char* dsn, char* query, int csv);
void extract_error(char *fn, SQLHANDLE handle, SQLSMALLINT type);

int main(int argc, char** argv) {
    if (argc < 4) {
        printf("Incorrect number of arguments.\n\n");
        printf("Arguments:\n");
        printf("pmodbc [dsn] [query] [csv={0,1}]\n\n");
        printf("DSN List:\n");
        list_drivers();
    } else {
        query_db(argv[1], argv[2], *argv[3] == '1');
    }
}


static void list_drivers() {
    SQLHENV env;
    char dsn[256];
    char desc[256];
    SQLSMALLINT dsn_ret;
    SQLSMALLINT desc_ret;
    SQLUSMALLINT direction;
    SQLRETURN ret;

    SQLAllocHandle(SQL_HANDLE_ENV, SQL_NULL_HANDLE, &env);
    SQLSetEnvAttr(env, SQL_ATTR_ODBC_VERSION, (void *) SQL_OV_ODBC3, 0);

    direction = SQL_FETCH_FIRST;
    while(SQL_SUCCEEDED(ret = SQLDataSources(env, direction, dsn, sizeof(dsn), &dsn_ret, desc, sizeof(desc), &desc_ret))) {
        direction = SQL_FETCH_NEXT;
        printf("%s - %s\n", dsn, desc);
        if (ret == SQL_SUCCESS_WITH_INFO) printf("\tdata truncation\n");
    }
}


static void query_db(char* dsn_str, char* query, int csv) {
    SQLHENV env;
    SQLHDBC dbc;
    SQLHSTMT stmt;
    SQLRETURN ret; /* ODBC API return status */
    SQLSMALLINT columns; /* number of columns in result-set */
    int row = 0;
    char dsn[1000];
    snprintf(dsn, 1000, "DSN=%s;", dsn_str);

    /* Allocate an environment handle */
    SQLAllocHandle(SQL_HANDLE_ENV, SQL_NULL_HANDLE, &env);
    /* We want ODBC 3 support */
    SQLSetEnvAttr(env, SQL_ATTR_ODBC_VERSION, (void *) SQL_OV_ODBC3, 0);
    /* Allocate a connection handle */
    ret = SQLAllocHandle(SQL_HANDLE_DBC, env, &dbc);
    if (!SQL_SUCCEEDED(ret)) {
      extract_error("SQLAllocHandle for dbc", env, SQL_HANDLE_ENV);
      exit(1);
    }
    /* Connect to the DSN mydsn */
    /* You will need to change mydsn to one you have created and tested */
    SQLDriverConnect(dbc, NULL, dsn, SQL_NTS, NULL, 0, NULL, SQL_DRIVER_COMPLETE);
    /* Allocate a statement handle */
    SQLAllocHandle(SQL_HANDLE_STMT, dbc, &stmt);
    /* Retrieve a list of tables */
    //SQLTables(stmt, NULL, 0, NULL, 0, NULL, 0, "TABLE", SQL_NTS);
    ret = SQLExecDirect(stmt, query, SQL_NTS);
    if (!SQL_SUCCEEDED(ret)) {
        extract_error("SQLExecuteDirect for dbc", stmt, SQL_HANDLE_STMT);
        exit(1);
    }
    /* How many columns are there */
    SQLNumResultCols(stmt, &columns);
    if (csv == 0) {
        while (SQL_SUCCEEDED(ret = SQLFetch(stmt))) {
        }
    } else {
        char csvbuf[1000];
        /* Loop through the rows in the result-set */
        while (SQL_SUCCEEDED(ret = SQLFetch(stmt))) {
            SQLUSMALLINT i;
            SQLLEN curlen = 0;
            SQLLEN indicator;
            /* Loop through the columns */
            for (i = 1; i <= columns; i++) {
                /* retrieve column data as a string */
                ret = SQLGetData(stmt, i, SQL_C_CHAR, csvbuf + curlen, sizeof(csvbuf) - curlen, &indicator);
                if (!SQL_SUCCEEDED(ret)) {
                    extract_error("SQLExecuteDirect for dbc", stmt, SQL_HANDLE_STMT);
                    exit(1);    
                }
                curlen += indicator;
                if (i != columns) 
                    csvbuf[curlen++] = ',';
            }
            puts(csvbuf);
        }
    }
}

void extract_error(char *fn, SQLHANDLE handle, SQLSMALLINT type)
{
    SQLINTEGER   i = 0;
    SQLINTEGER   native;
    SQLCHAR  state[ 7 ];
    SQLCHAR  text[256];
    SQLSMALLINT  len;
    SQLRETURN    ret;

    fprintf(stderr,
            "\n"
            "The driver reported the following diagnostics whilst running "
            "%s\n\n",
            fn);

    do
    {
        ret = SQLGetDiagRec(type, handle, ++i, state, &native, text,
                            sizeof(text), &len );
        if (SQL_SUCCEEDED(ret))
            printf("%s:%d:%d:%s\n", state, i, native, text);
    }
    while( ret == SQL_SUCCESS );
}