#include <iostream>
#include <memory>
#include <string>
#include <sql.h>
#include <sqlext.h>

// Collects diagnostic records without without saving until SQL_NO_DATA.
void CollectDiagRecords(SQLSMALLINT hndlType, SQLHANDLE hndl)
{
	constexpr size_t msgBufferSize = 100;

	SQLINTEGER error;
	SQLSMALLINT msgSize{ 0 };
	SQLWCHAR msgBuffer[msgBufferSize];
	SQLWCHAR state[SQL_SQLSTATE_SIZE + 1];

	for(SQLSMALLINT index = 1; ; ++index)
	{
        // Crashes after calling this function with index == 332. It is dependant on msgBufferSize and SQL procedure.
		auto diagRet = SQLGetDiagRec(
			hndlType,
			hndl,
			index,
			state,
			&error,
			msgBuffer,
			(SQLSMALLINT)msgBufferSize,
			&msgSize);


		switch(diagRet)
		{
		case SQL_SUCCESS:
			//BOOST_LOG_TRIVIAL(info) << index << " " << boost::locale::conv::utf_to_utf<char>(state, state + SQL_SQLSTATE_SIZE) << " " << boost::locale::conv::utf_to_utf<char>(msgBuffer, msgBuffer + msgSize) << std::endl;
			continue;

		case SQL_SUCCESS_WITH_INFO:
		{
			auto bufferSize = msgSize + 1;
			std::unique_ptr< SQLWCHAR[] > buf{ new SQLWCHAR[bufferSize] };
			SQLGetDiagRec(hndlType, hndl, index, state, &error, buf.get(), bufferSize, &msgSize);
		}
		continue;

		case SQL_INVALID_HANDLE:
			throw std::runtime_error("CollectDiagRecords SQL_INVALID_HANDLE");

		case SQL_ERROR:
			throw std::runtime_error("CollectDiagRecords SQL_ERROR");

		case SQL_NO_DATA:
			// End collecting diag info
			return;
		}
	}
}

int main(int argc, char** argv)
{
	try
	{
		SQLHDBC hdbc = nullptr;
		SQLHENV henv = nullptr;

		// Allocate environment handle
		auto ret = SQLAllocHandle(SQL_HANDLE_ENV, SQL_NULL_HANDLE, &henv);

		if(ret != SQL_SUCCESS && ret != SQL_SUCCESS_WITH_INFO)
		{
			throw std::runtime_error("SQLAllocHandle SQL_HANDLE_ENV");
		}

		SQLSetEnvAttr(henv, SQL_ATTR_ODBC_VERSION, (void*)SQL_OV_ODBC3, 0);

		// Allocate connection handle
		ret = SQLAllocHandle(SQL_HANDLE_DBC, henv, &hdbc);

		if(ret != SQL_SUCCESS && ret != SQL_SUCCESS_WITH_INFO)
		{
			CollectDiagRecords(SQL_HANDLE_ENV, henv);
			throw std::runtime_error("SQLAllocHandle SQL_HANDLE_DBC");
		}

		// Connect
		SQLWCHAR connString[] = L"Driver={PostgreSQL Unicode(x64)};Server=127.0.0.1;Port=5433;Database=crashdb;Uid=user1;Pwd=user1;";
		ret = SQLDriverConnect(hdbc, nullptr, connString, sizeof(connString)/sizeof(SQLWCHAR), nullptr, 0, nullptr, SQL_DRIVER_NOPROMPT);

		if(ret != SQL_SUCCESS && ret != SQL_SUCCESS_WITH_INFO)
		{
			CollectDiagRecords(SQL_HANDLE_DBC, hdbc);
			throw std::runtime_error("SQLDriverConnect");
		}

		SQLHSTMT hstmt;
		ret = SQLAllocHandle(SQL_HANDLE_STMT, hdbc, &hstmt);

		if(ret != SQL_SUCCESS && ret != SQL_SUCCESS_WITH_INFO)
		{
			throw std::runtime_error("SQLAllocHandle SQL_HANDLE_STMT");
		}

		SQLWCHAR query[] = L"CALL public.crashme()";
		ret = SQLExecDirect(hstmt, query, sizeof(query)/sizeof(SQLWCHAR));

		switch(ret)
		{
		case SQL_SUCCESS:
			break;

		case SQL_SUCCESS_WITH_INFO:
            // SQL Procedure crashme() returns "NOTICE" messages, which causes SQLExecDirect return SQL_SUCCESS_WITH_INFO.
			CollectDiagRecords(SQL_HANDLE_STMT, hstmt);
			break;

		default:
			throw std::runtime_error("SQLExecDirect default");
		}

		SQLFreeHandle(SQL_HANDLE_STMT, hstmt);
		SQLDisconnect(hdbc);
		SQLFreeHandle(SQL_HANDLE_DBC, hdbc);
		SQLFreeHandle(SQL_HANDLE_ENV, henv);
	}
	catch(std::exception& e)
	{
		std::cerr << e.what() << std::endl;
		return -1;
	}

	return 0;
}