#include "sm_globals.h"

#include "HandleSys.h"

#include "Database.h"

#include "ExtensionSys.h"

#include "PluginSys.h"

#include "sm_stringutil.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;

}

class TQueryOp : public IDBThreadOperation

{

public:

	TQueryOp(IDatabase *db, IPluginFunction *pf, const char *query, cell_t data) : 

	  m_pDatabase(db), m_pFunction(pf), m_Query(query), m_Data(data),

	  me(g_PluginSys.GetPluginByCtx(pf->GetParentContext()->GetContext())),

	  m_pQuery(NULL)

	{

		/* We always increase the reference count because this is potentially

		 * asynchronous.  Otherwise the original handle could be closed while 

		 * we're still latched onto it.

		 */

		m_pDatabase->IncReferenceCount();

		/* Now create our own Handle such that it can only be closed by us.

		 * We allow cloning just in case someone wants to hold onto it.

		 */

		HandleSecurity sec(me->GetIdentity(), g_pCoreIdent);

		HandleAccess access;

		g_HandleSys.InitAccessDefaults(NULL, &access);

		access.access[HandleAccess_Delete] = HANDLE_RESTRICT_IDENTITY|HANDLE_RESTRICT_OWNER;

		m_MyHandle = g_HandleSys.CreateHandleEx(g_DBMan.GetDatabaseType(),

			db,

			&sec,

			&access,

			NULL);

	}

	~TQueryOp()

	{

		if (m_pQuery)

		{

			m_pQuery->Destroy();

		}

		/* Close our Handle if it's valid. */

		if (m_MyHandle != BAD_HANDLE)

		{

			HandleSecurity sec(me->GetIdentity(), g_pCoreIdent);

			g_HandleSys.FreeHandle(m_MyHandle, &sec);

		} else {

			/* Otherwise, there is an open ref to the db */

			m_pDatabase->Close();

		}

	}

	CPlugin *GetPlugin()

	{

		return me;

	}

	IDBDriver *GetDriver()

	{

		return m_pDatabase->GetDriver();

	}

	void RunThreadPart()

	{

		m_pDatabase->LockForFullAtomicOperation();

		m_pQuery = m_pDatabase->DoQuery(m_Query.c_str());

		if (!m_pQuery)

		{

			UTIL_Format(error, sizeof(error), "%s", m_pDatabase->GetError());

		}

		m_pDatabase->UnlockFromFullAtomicOperation();

	}

	void CancelThinkPart()

	{

		m_pFunction->PushCell(BAD_HANDLE);

		m_pFunction->PushCell(BAD_HANDLE);

		m_pFunction->PushString("Driver is unloading");

		m_pFunction->PushCell(m_Data);

		m_pFunction->Execute(NULL);

		delete this;

	}

	void RunThinkPart()

	{

		/* Create a Handle for our query */

		HandleSecurity sec(me->GetIdentity(), g_pCoreIdent);

		HandleAccess access;

		g_HandleSys.InitAccessDefaults(NULL, &access);

		access.access[HandleAccess_Delete] = HANDLE_RESTRICT_IDENTITY|HANDLE_RESTRICT_OWNER;

		Handle_t qh = BAD_HANDLE;

		if (m_pQuery)

		{

			qh = g_HandleSys.CreateHandle(hQueryType, m_pQuery, me->GetIdentity(), g_pCoreIdent, NULL);

			if (qh != BAD_HANDLE)

			{

				m_pQuery = NULL;

			} else {

				UTIL_Format(error, sizeof(error), "Could not alloc handle");

			}

		}

		m_pFunction->PushCell(m_MyHandle);

		m_pFunction->PushCell(qh);

		m_pFunction->PushString(qh == BAD_HANDLE ? error : "");

		m_pFunction->PushCell(m_Data);

		m_pFunction->Execute(NULL);

		if (qh != BAD_HANDLE)

		{

			g_HandleSys.FreeHandle(qh, &sec);

		}

		delete this;

	}

private:

	IDatabase *m_pDatabase;

	IPluginFunction *m_pFunction;

	String m_Query;

	cell_t m_Data;

	CPlugin *me;

	IQuery *m_pQuery;

	char error[255];

	Handle_t m_MyHandle;

};

class TConnectOp : public IDBThreadOperation

{

public:

	TConnectOp(IPluginFunction *func, IDBDriver *driver, const char *_dbname)

	{

		m_pFunction = func;

		m_pDriver = driver;

		m_pDatabase = NULL;

		error[0] = '\0';

		strncopy(dbname, _dbname, sizeof(dbname));

		me = g_PluginSys.GetPluginByCtx(m_pFunction->GetParentContext()->GetContext());

	}

	CPlugin *GetPlugin()

	{

		return me;

	}

	IDBDriver *GetDriver()

	{

		return m_pDriver;

	}

	void RunThreadPart()

	{

		g_DBMan.LockConfig();

		const DatabaseInfo *pInfo = g_DBMan.FindDatabaseConf(dbname);

		if (!pInfo)

		{

			UTIL_Format(error, sizeof(error), "Could not find database config \"%s\"", dbname);

		} else {

			m_pDatabase = m_pDriver->Connect(pInfo, false, error, sizeof(error));

		}

		g_DBMan.UnlockConfig();

	}

	void CancelThinkPart()

	{

		if (m_pDatabase)

		{

			m_pDatabase->Close();

		}

		m_pFunction->PushCell(BAD_HANDLE);

		m_pFunction->PushCell(BAD_HANDLE);

		m_pFunction->PushString("Driver is unloading");

		m_pFunction->PushCell(0);

		m_pFunction->Execute(NULL);

		delete this;

	}

	void RunThinkPart()

	{

		Handle_t hndl = BAD_HANDLE;

		if (m_pDatabase)

		{

			if ((hndl = g_DBMan.CreateHandle(DBHandle_Database, m_pDatabase, me->GetIdentity()))

				== BAD_HANDLE)

			{

				m_pDatabase->Close();

				UTIL_Format(error, sizeof(error), "Unable to allocate Handle");

			}

		}

		m_pFunction->PushCell(m_pDriver->GetHandle());

		m_pFunction->PushCell(hndl);

		m_pFunction->PushString(hndl == BAD_HANDLE ? error : "");

		m_pFunction->PushCell(0);

		m_pFunction->Execute(NULL);

		delete this;

	}

private:

	CPlugin *me;

	IPluginFunction *m_pFunction;

	IDBDriver *m_pDriver;

	IDatabase *m_pDatabase;

	char dbname[64];

	char error[255];

};

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;

	}

	Handle_t hndl = g_DBMan.CreateHandle(DBHandle_Database, db, pContext->GetIdentity());

	if (!hndl)

	{

		db->Close();

		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 hndl;

}

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

{

	IPluginFunction *pf = pContext->GetFunctionById(params[1]);

	if (!pf)

	{

		return pContext->ThrowNativeError("Function id %x is invalid", params[1]);

	}

	char *conf;

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

	IDBDriver *driver = NULL;

	const DatabaseInfo *pInfo = g_DBMan.FindDatabaseConf(conf);

	char error[255];

	if (pInfo != NULL)

	{

		if (pInfo->driver[0] == '\0')

		{

			driver = g_DBMan.GetDefaultDriver();

		} else {

			driver = g_DBMan.FindOrLoadDriver(pInfo->driver);

		}

		if (!driver)

		{

			UTIL_Format(error, 

				sizeof(error), 

				"Could not find driver \"%s\"", 

				pInfo->driver[0] == '\0' ? g_DBMan.GetDefaultDriverName() : pInfo->driver);

		} else if (!driver->IsThreadSafe()) {

			UTIL_Format(error,

				sizeof(error),

				"Driver \"%s\" is not thread safe!",

				driver->GetIdentifier());

		}

	} else {

		UTIL_Format(error, sizeof(error), "Could not find database conf \"%s\"", conf);

	}

	if (!pInfo || !driver)

	{

		pf->PushCell(BAD_HANDLE);

		pf->PushCell(BAD_HANDLE);

		pf->PushString(error);

		pf->PushCell(0);

		pf->Execute(NULL);

		return 0;

	}

	/* 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()));

	}

	/* Finally, add to the thread if we can */

	TConnectOp *op = new TConnectOp(pf, driver, conf);

	CPlugin *pPlugin = g_PluginSys.GetPluginByCtx(pContext->GetContext());

	if (pPlugin->GetProperty("DisallowDBThreads", NULL)

		|| !g_DBMan.AddToThreadQueue(op, PrioQueue_High))

	{

		/* Do everything right now */

		op->RunThreadPart();

		op->RunThinkPart();

	}

	return 1;

}

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)

	{

		Handle_t hndl = g_DBMan.CreateHandle(DBHandle_Database, db, pContext->GetIdentity());

		if (!hndl)

		{

			db->Close();

			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 hndl;

	}

	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;