#include "sm_globals.h"

#include "HandleSys.h"

#include "Database.h"

#include "ExtensionSys.h"

#include "PluginSys.h"

HandleType_t hQueryType;

HandleType_t hStmtType;

class DatabaseHelpers : 

	public SMGlobalClass,

	public IHandleTypeDispatch

{

public:

	virtual void OnSourceModAllInitialized()

	{

		HandleAccess acc;

		/* Disable cloning */

		g_HandleSys.InitAccessDefaults(NULL, &acc);

		acc.access[HandleAccess_Clone] = HANDLE_RESTRICT_OWNER|HANDLE_RESTRICT_IDENTITY;

		TypeAccess tacc;

		g_HandleSys.InitAccessDefaults(&tacc, NULL);

		tacc.ident = g_pCoreIdent;

		hQueryType = g_HandleSys.CreateType("IQuery", this, 0, &tacc, &acc, g_pCoreIdent, NULL);

		hStmtType = g_HandleSys.CreateType("IPreparedQuery", this, hQueryType, &tacc, &acc, g_pCoreIdent, NULL);

	}

	virtual void OnSourceModShutdown()

	{

		g_HandleSys.RemoveType(hStmtType, g_pCoreIdent);

		g_HandleSys.RemoveType(hQueryType, g_pCoreIdent);

	}

	virtual void OnHandleDestroy(HandleType_t type, void *object)

	{

		if (type == hQueryType)

		{

			IQuery *query = (IQuery *)object;

			query->Destroy();

		} else if (type == hStmtType) {

			IPreparedQuery *query = (IPreparedQuery *)object;

			query->Destroy();

		}

	}

} s_DatabaseNativeHelpers;

//is this safe for stmt handles? i think since it's single inheritance, it always will be.

inline HandleError ReadQueryHndl(Handle_t hndl, IPluginContext *pContext, IQuery **query)

{

	HandleSecurity sec;

	sec.pOwner = pContext->GetIdentity();

	sec.pIdentity = g_pCoreIdent;

	return g_HandleSys.ReadHandle(hndl, hQueryType, &sec, (void **)query);

}

inline HandleError ReadStmtHndl(Handle_t hndl, IPluginContext *pContext, IPreparedQuery **query)

{

	HandleSecurity sec;

	sec.pOwner = pContext->GetIdentity();

	sec.pIdentity = g_pCoreIdent;

	return g_HandleSys.ReadHandle(hndl, hStmtType, &sec, (void **)query);

}

inline HandleError ReadDbOrStmtHndl(Handle_t hndl, IPluginContext *pContext, IDatabase **db, IPreparedQuery **query)

{

	HandleError err;

	if ((err = g_DBMan.ReadHandle(hndl, DBHandle_Database, (void **)db)) == HandleError_Type)

	{

		*db = NULL;

		return ReadStmtHndl(hndl, pContext, query);

	}

	return err;

}

static cell_t SQL_Connect(IPluginContext *pContext, const cell_t *params)

{

	char *conf, *err;

	size_t maxlength = (size_t)params[4];

	bool persistent = params[2] ? true : false;

	pContext->LocalToString(params[1], &conf);

	pContext->LocalToString(params[3], &err);

	IDBDriver *driver;

	IDatabase *db;

	if (!g_DBMan.Connect(conf, &driver, &db, persistent, err, maxlength))

	{

		return BAD_HANDLE;

	}

	/* HACK! Add us to the dependency list */

	CExtension *pExt = g_Extensions.GetExtensionFromIdent(driver->GetIdentity());

	if (pExt)

	{

		g_Extensions.BindChildPlugin(pExt, g_PluginSys.FindPluginByContext(pContext->GetContext()));

	}

	return db->GetHandle();

}

static cell_t SQL_ConnectEx(IPluginContext *pContext, const cell_t *params)

{

	IDBDriver *driver;

	if (params[1] == BAD_HANDLE)

	{

		if ((driver = g_DBMan.GetDefaultDriver()) == NULL)

		{

			return pContext->ThrowNativeError("Could not find any default driver");

		}

	} else {

		HandleError err;

		if ((err = g_DBMan.ReadHandle(params[1], DBHandle_Driver, (void **)&driver))

			!= HandleError_None)

		{

			return pContext->ThrowNativeError("Invalid driver Handle %x (error: %d)", params[1], err);

		}

	}

	char *host, *user, *pass, *database, *error;

	size_t maxlength = (size_t)params[7];

	bool persistent = params[8] ? true : false;

	unsigned int port = params[9];

	unsigned int maxTimeout = params[10];

	pContext->LocalToString(params[2], &host);

	pContext->LocalToString(params[3], &user);

	pContext->LocalToString(params[4], &pass);

	pContext->LocalToString(params[5], &database);

	pContext->LocalToString(params[6], &error);

	DatabaseInfo info;

	info.database = database;

	info.driver = driver->GetIdentifier();

	info.host = host;

	info.maxTimeout = maxTimeout;

	info.pass = pass;

	info.port = port;

	info.user = user;

	IDatabase *db = driver->Connect(&info, persistent, error, maxlength);

	if (db)

	{

		/* HACK! Add us to the dependency list */

		CExtension *pExt = g_Extensions.GetExtensionFromIdent(driver->GetIdentity());

		if (pExt)

		{

			g_Extensions.BindChildPlugin(pExt, g_PluginSys.FindPluginByContext(pContext->GetContext()));

		}

		return db->GetHandle();

	}

	return BAD_HANDLE;

}

static cell_t SQL_GetDriverIdent(IPluginContext *pContext, const cell_t *params)

{

	IDBDriver *driver;

	if (params[1] == BAD_HANDLE)

	{

		if ((driver = g_DBMan.GetDefaultDriver()) == NULL)

		{

			return pContext->ThrowNativeError("Could not find any default driver");

		}

	} else {

		HandleError err;

		if ((err = g_DBMan.ReadHandle(params[1], DBHandle_Driver, (void **)&driver))

			!= HandleError_None)

		{

			return pContext->ThrowNativeError("Invalid driver Handle %x (error: %d)", params[1], err);

		}

	}

	pContext->StringToLocalUTF8(params[2], params[3], driver->GetIdentifier(), NULL);

	return 1;

}

static cell_t SQL_GetDriver(IPluginContext *pContext, const cell_t *params)

{

	char *name;

	pContext->LocalToString(params[1], &name);

	IDBDriver *driver = NULL;

	if (name[0] == '\0')

	{

		driver = g_DBMan.GetDefaultDriver();

	} else {

		driver = g_DBMan.FindOrLoadDriver(name);

	}

	return (driver != NULL) ? driver->GetHandle() : BAD_HANDLE;

}

static cell_t SQL_GetDriverProduct(IPluginContext *pContext, const cell_t *params)

{

	IDBDriver *driver;

	if (params[1] == BAD_HANDLE)

	{

		if ((driver = g_DBMan.GetDefaultDriver()) == NULL)

		{

			return pContext->ThrowNativeError("Could not find any default driver");

		}

	} else {

		HandleError err;

		if ((err = g_DBMan.ReadHandle(params[1], DBHandle_Driver, (void **)&driver))

			!= HandleError_None)

		{

			return pContext->ThrowNativeError("Invalid driver Handle %x (error: %d)", params[1], err);

		}

	}

	pContext->StringToLocalUTF8(params[2], params[3], driver->GetProductName(), NULL);

	return 1;

}

static cell_t SQL_GetAffectedRows(IPluginContext *pContext, const cell_t *params)

{

	IDatabase *db = NULL;

	IPreparedQuery *stmt = NULL;

	HandleError err;

	if ((err = ReadDbOrStmtHndl(params[1], pContext, &db, &stmt)) != HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid statement or db Handle %x (error: %d)", params[1], err);

	}

	if (db)

	{

		return db->GetAffectedRows();

	} else if (stmt) {

		return stmt->GetAffectedRows();

	}

	return pContext->ThrowNativeError("Unknown error reading db/stmt handles");

}

static cell_t SQL_GetInsertId(IPluginContext *pContext, const cell_t *params)

{

	IDatabase *db = NULL;

	IPreparedQuery *stmt = NULL;

	HandleError err;

	if ((err = ReadDbOrStmtHndl(params[1], pContext, &db, &stmt)) != HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid statement or db Handle %x (error: %d)", params[1], err);

	}

	if (db)

	{

		return db->GetInsertID();

	} else if (stmt) {

		return stmt->GetInsertID();

	}

	return pContext->ThrowNativeError("Unknown error reading db/stmt handles");

}

static cell_t SQL_GetError(IPluginContext *pContext, const cell_t *params)

{

	IDatabase *db = NULL;

	IPreparedQuery *stmt = NULL;

	HandleError err;

	if ((err = ReadDbOrStmtHndl(params[1], pContext, &db, &stmt)) != HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid statement or db Handle %x (error: %d)", params[1], err);

	}

	const char *error = "";

	if (db)

	{

		error = db->GetError();

	} else if (stmt) {

		error = stmt->GetError();

	}

	if (error[0] == '\0')

	{

		return false;

	}

	pContext->StringToLocalUTF8(params[2], params[3], error, NULL);

	return 1;

}

static cell_t SQL_QuoteString(IPluginContext *pContext, const cell_t *params)

{

	IDatabase *db = NULL;

	HandleError err;

	if ((err = g_DBMan.ReadHandle(params[1], DBHandle_Database, (void **)&db))

		!= HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid database Handle %x (error: %d)", params[1], err);

	}

	char *input, *output;

	size_t maxlength = (size_t)params[4];

	pContext->LocalToString(params[2], &input);

	pContext->LocalToString(params[3], &output);

	size_t written;

	bool s = db->QuoteString(input, output, maxlength, &written);

	cell_t *addr;

	pContext->LocalToPhysAddr(params[5], &addr);

	*addr = (cell_t)written;

	return s ? 1 : 0;

}

static cell_t SQL_FastQuery(IPluginContext *pContext, const cell_t *params)

{

	IDatabase *db = NULL;

	HandleError err;

	if ((err = g_DBMan.ReadHandle(params[1], DBHandle_Database, (void **)&db))

		!= HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid database Handle %x (error: %d)", params[1], err);

	}

	char *query;

	pContext->LocalToString(params[2], &query);

	return db->DoSimpleQuery(query) ? 1 : 0;

}

static cell_t SQL_Query(IPluginContext *pContext, const cell_t *params)

{

	IDatabase *db = NULL;

	HandleError err;

	if ((err = g_DBMan.ReadHandle(params[1], DBHandle_Database, (void **)&db))

		!= HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid database Handle %x (error: %d)", params[1], err);

	}

	char *query;

	pContext->LocalToString(params[2], &query);

	IQuery *qr = db->DoQuery(query);

	if (!qr)

	{

		return BAD_HANDLE;

	}

	Handle_t hndl = g_HandleSys.CreateHandle(hQueryType, qr, pContext->GetIdentity(), g_pCoreIdent, NULL);

	if (hndl == BAD_HANDLE)

	{

		qr->Destroy();

		return BAD_HANDLE;

	}

	return hndl;

}

static cell_t SQL_PrepareQuery(IPluginContext *pContext, const cell_t *params)

{

	IDatabase *db = NULL;

	HandleError err;

	if ((err = g_DBMan.ReadHandle(params[1], DBHandle_Database, (void **)&db))

		!= HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid database Handle %x (error: %d)", params[1], err);

	}

	char *query, *error;

	size_t maxlength = (size_t)params[4];

	pContext->LocalToString(params[2], &query);

	pContext->LocalToString(params[3], &error);

	IPreparedQuery *qr = db->PrepareQuery(query, error, maxlength);

	if (!qr)

	{

		return BAD_HANDLE;

	}

	Handle_t hndl = g_HandleSys.CreateHandle(hStmtType, qr, pContext->GetIdentity(), g_pCoreIdent, NULL);

	if (hndl == BAD_HANDLE)

	{

		qr->Destroy();

		return BAD_HANDLE;

	}

	return hndl;

}

static cell_t SQL_FetchMoreResults(IPluginContext *pContext, const cell_t *params)

{

	IQuery *query;

	HandleError err;

	if ((err = ReadQueryHndl(params[1], pContext, &query)) != HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid query Handle %x (error: %d)", params[1], err);

	}

	return query->FetchMoreResults() ? 1 : 0;

}

static cell_t SQL_HasResultSet(IPluginContext *pContext, const cell_t *params)

{

	IQuery *query;

	HandleError err;

	if ((err = ReadQueryHndl(params[1], pContext, &query)) != HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid query Handle %x (error: %d)", params[1], err);

	}

	return query->GetResultSet() != NULL ? true : false;

}

static cell_t SQL_GetRowCount(IPluginContext *pContext, const cell_t *params)

{

	IQuery *query;

	HandleError err;

	if ((err = ReadQueryHndl(params[1], pContext, &query)) != HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid query Handle %x (error: %d)", params[1], err);

	}

	IResultSet *rs = query->GetResultSet();

	if (!rs)

	{

		return 0;

	}

	return rs->GetRowCount();

}

static cell_t SQL_GetFieldCount(IPluginContext *pContext, const cell_t *params)

{

	IQuery *query;

	HandleError err;

	if ((err = ReadQueryHndl(params[1], pContext, &query)) != HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid query Handle %x (error: %d)", params[1], err);

	}

	IResultSet *rs = query->GetResultSet();

	if (!rs)

	{

		return 0;

	}

	return rs->GetFieldCount();

}

static cell_t SQL_FieldNumToName(IPluginContext *pContext, const cell_t *params)

{

	IQuery *query;

	HandleError err;

	if ((err = ReadQueryHndl(params[1], pContext, &query)) != HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid query Handle %x (error: %d)", params[1], err);

	}

	IResultSet *rs = query->GetResultSet();

	if (!rs)

	{

		return pContext->ThrowNativeError("No current result set");

	}

	unsigned int field = params[2];

	const char *fldname;

	if ((fldname = rs->FieldNumToName(field)) == NULL)

	{

		return pContext->ThrowNativeError("Invalid field index %d", field);

	}

	pContext->StringToLocalUTF8(params[3], params[4], fldname, NULL);

	return 1;

}

static cell_t SQL_FieldNameToNum(IPluginContext *pContext, const cell_t *params)

{

	IQuery *query;

	HandleError err;

	if ((err = ReadQueryHndl(params[1], pContext, &query)) != HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid query Handle %x (error: %d)", params[1], err);

	}

	IResultSet *rs = query->GetResultSet();

	if (!rs)

	{

		return pContext->ThrowNativeError("No current result set");

	}

	char *field;

	pContext->LocalToString(params[2], &field);

	cell_t *num;

	pContext->LocalToPhysAddr(params[3], &num);

	return rs->FieldNameToNum(field, (unsigned int *)num) ? 1 : 0;

}

static cell_t SQL_FetchRow(IPluginContext *pContext, const cell_t *params)

{

	IQuery *query;

	HandleError err;

	if ((err = ReadQueryHndl(params[1], pContext, &query)) != HandleError_None)

	{

		return pContext->ThrowNativeError("Invalid query Handle %x (error: %d)", params[1], err);