/**
 * vim: set ts=4 :
 * =============================================================================
 * SourceMod
 * Copyright (C) 2004-2008 AlliedModders LLC.  All rights reserved.
 * =============================================================================
 *
 * This program is free software; you can redistribute it and/or modify it under
 * the terms of the GNU General Public License, version 3.0, as published by the
 * Free Software Foundation.
 * 
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * As a special exception, AlliedModders LLC gives you permission to link the
 * code of this program (as well as its derivative works) to "Half-Life 2," the
 * "Source Engine," the "SourcePawn JIT," and any Game MODs that run on software
 * by the Valve Corporation.  You must obey the GNU General Public License in
 * all respects for all other code used.  Additionally, AlliedModders LLC grants
 * this exception to all derivative works.  AlliedModders LLC defines further
 * exceptions, found in LICENSE.txt (as of this writing, version JULY-31-2007),
 * or <http://www.sourcemod.net/license.php>.
 *
 * Version: $Id$
 */

#include <malloc.h>
#include "sm_globals.h"
#include "sm_stringutil.h"
#include "HandleSys.h"
#include "sm_trie_tpl.h"

HandleType_t htCellTrie;

enum TrieNodeType
{
	TrieNode_Cell,
	TrieNode_CellArray,
	TrieNode_String,
};

struct SmartTrieNode
{
	SmartTrieNode()
	{
		ptr = NULL;
		type = TrieNode_Cell;
	}
	SmartTrieNode(const SmartTrieNode &obj)
	{
		type = obj.type;
		ptr = obj.ptr;
		data = obj.data;
		data_len = obj.data_len;
	}
	SmartTrieNode & operator =(const SmartTrieNode &src)
	{
		type = src.type;
		ptr = src.ptr;
		data = src.data;
		data_len = src.data_len;
		return *this;
	}
	TrieNodeType type;
	cell_t *ptr;
	cell_t data;
	cell_t data_len;
};

struct CellTrie
{
	KTrie<SmartTrieNode> trie;
	cell_t mem_usage;
};

class TrieHelpers : 
	public SMGlobalClass,
	public IHandleTypeDispatch
{
public: //SMGlobalClass
	void OnSourceModAllInitialized()
	{
		htCellTrie = g_HandleSys.CreateType("Trie", this, 0, NULL, NULL, g_pCoreIdent, NULL);
	}
	void OnSourceModShutdown()
	{
		g_HandleSys.RemoveType(htCellTrie, g_pCoreIdent);
	}
public: //IHandleTypeDispatch
	static void DestroySmartTrieNode(SmartTrieNode *pNode)
	{
		free(pNode->ptr);
	}
	void OnHandleDestroy(HandleType_t type, void *object)
	{
		CellTrie *pTrie = (CellTrie *)object;

		pTrie->trie.run_destructor(DestroySmartTrieNode);

		delete pTrie;
	}
	bool GetHandleApproxSize(HandleType_t type, void *object, unsigned int *pSize)
	{
		CellTrie *pArray = (CellTrie *)object;
		*pSize = sizeof(CellTrie) + pArray->mem_usage + pArray->trie.mem_usage();
		return true;
	}
} s_CellTrieHelpers;

static cell_t CreateTrie(IPluginContext *pContext, const cell_t *params)
{
	CellTrie *pTrie = new CellTrie;
	Handle_t hndl;

	pTrie->mem_usage = 0;

	if ((hndl = g_HandleSys.CreateHandle(htCellTrie, pTrie, pContext->GetIdentity(), g_pCoreIdent, NULL))
		== BAD_HANDLE)
	{
		delete pTrie;
		return BAD_HANDLE;
	}

	return hndl;
}

static void UpdateNodeCells(CellTrie *pTrie, SmartTrieNode *pData, const cell_t *cells, cell_t num_cells)
{
	if (num_cells == 1)
	{
		pData->data = *cells;
		pData->type = TrieNode_Cell;
	}
	else
	{
		pData->type = TrieNode_CellArray;
		if (pData->ptr == NULL)
		{
			pData->ptr = (cell_t *)malloc(num_cells * sizeof(cell_t));
			pData->data_len = num_cells;
			pTrie->mem_usage += (pData->data_len * sizeof(cell_t));
		}
		else if (pData->data_len < num_cells)
		{
			pData->ptr = (cell_t *)realloc(pData->ptr, num_cells * sizeof(cell_t));
			pTrie->mem_usage += (num_cells - pData->data_len) * sizeof(cell_t);
			pData->data_len = num_cells;
		}
		if (num_cells != 0)
		{
			memcpy(pData->ptr, cells, sizeof(cell_t) * num_cells);
		}
		pData->data = num_cells;
	}
}

static void UpdateNodeString(CellTrie *pTrie, SmartTrieNode *pData, const char *str)
{
	size_t len = strlen(str);
	cell_t num_cells = (len + sizeof(cell_t)) / sizeof(cell_t);

	if (pData->ptr == NULL)
	{
		pData->ptr = (cell_t *)malloc(num_cells * sizeof(cell_t));
		pData->data_len = num_cells;
		pTrie->mem_usage += (pData->data_len * sizeof(cell_t));
	}
	else if (pData->data_len < num_cells)
	{
		pData->ptr = (cell_t *)realloc(pData->ptr, num_cells * sizeof(cell_t));
		pTrie->mem_usage += (num_cells - pData->data_len) * sizeof(cell_t);
		pData->data_len = num_cells;
	}

	strcpy((char *)pData->ptr, str);
	pData->data = len;
	pData->type = TrieNode_String;
}

static cell_t SetTrieValue(IPluginContext *pContext, const cell_t *params)
{
	Handle_t hndl;
	CellTrie *pTrie;
	HandleError err;
	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);

	hndl = params[1];

	if ((err = g_HandleSys.ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
		!= HandleError_None)
	{
		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
	}

	char *key;
	pContext->LocalToString(params[2], &key);

	SmartTrieNode *pNode;
	if ((pNode = pTrie->trie.retrieve(key)) == NULL)
	{
		SmartTrieNode node;
		UpdateNodeCells(pTrie, &node, &params[3], 1);
		return pTrie->trie.insert(key, node) ? 1 : 0;
	}

	if (!params[4])
	{
		return 0;
	}

	UpdateNodeCells(pTrie, pNode, &params[3], 1);

	return 1;
}

static cell_t SetTrieArray(IPluginContext *pContext, const cell_t *params)
{
	Handle_t hndl;
	CellTrie *pTrie;
	HandleError err;
	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);

	hndl = params[1];

	if ((err = g_HandleSys.ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
		!= HandleError_None)
	{
		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
	}

	if (params[4] < 0)
	{
		return pContext->ThrowNativeError("Invalid array size: %d", params[4]);
	}

	char *key;
	cell_t *array;
	pContext->LocalToString(params[2], &key);
	pContext->LocalToPhysAddr(params[3], &array);

	SmartTrieNode *pNode;
	if ((pNode = pTrie->trie.retrieve(key)) == NULL)
	{
		SmartTrieNode node;
		UpdateNodeCells(pTrie, &node, array, params[4]);
		if (!pTrie->trie.insert(key, node))
		{
			free(node.ptr);
			return 0;
		}
		return 1;
	}

	if (!params[4])
	{
		return 0;
	}

	UpdateNodeCells(pTrie, pNode, array, params[4]);

	return 1;
}

static cell_t SetTrieString(IPluginContext *pContext, const cell_t *params)
{
	Handle_t hndl;
	CellTrie *pTrie;
	HandleError err;
	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);

	hndl = params[1];

	if ((err = g_HandleSys.ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
		!= HandleError_None)
	{
		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
	}

	char *key, *val;
	pContext->LocalToString(params[2], &key);
	pContext->LocalToString(params[3], &val);

	SmartTrieNode *pNode;
	if ((pNode = pTrie->trie.retrieve(key)) == NULL)
	{
		SmartTrieNode node;
		UpdateNodeString(pTrie, &node, val);
		if (!pTrie->trie.insert(key, node))
		{
			free(node.ptr);
			return 0;
		}
		return 1;
	}

	if (!params[4])
	{
		return 0;
	}

	UpdateNodeString(pTrie, pNode, val);

	return 1;
}

static cell_t RemoveFromTrie(IPluginContext *pContext, const cell_t *params)
{
	Handle_t hndl;
	CellTrie *pTrie;
	HandleError err;
	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);

	hndl = params[1];

	if ((err = g_HandleSys.ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
		!= HandleError_None)
	{
		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
	}

	char *key;
	pContext->LocalToString(params[2], &key);

	SmartTrieNode *pNode;
	if ((pNode = pTrie->trie.retrieve(key)) == NULL)
	{
		return 0;
	}

	free(pNode->ptr);
	pNode->ptr = NULL;

	return pTrie->trie.remove(key) ? 1 : 0;
}

static cell_t ClearTrie(IPluginContext *pContext, const cell_t *params)
{
	Handle_t hndl;
	CellTrie *pTrie;
	HandleError err;
	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);

	hndl = params[1];

	if ((err = g_HandleSys.ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
		!= HandleError_None)
	{
		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
	}

	pTrie->trie.run_destructor(TrieHelpers::DestroySmartTrieNode);
	pTrie->trie.clear();

	return 1;
}

static cell_t GetTrieValue(IPluginContext *pContext, const cell_t *params)
{
	Handle_t hndl;
	CellTrie *pTrie;
	HandleError err;
	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);

	hndl = params[1];

	if ((err = g_HandleSys.ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
		!= HandleError_None)
	{
		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
	}

	char *key;
	cell_t *pValue;
	pContext->LocalToString(params[2], &key);
	pContext->LocalToPhysAddr(params[3], &pValue);

	SmartTrieNode *pNode;
	if ((pNode = pTrie->trie.retrieve(key)) == NULL)
	{
		return 0;
	}

	if (pNode->type == TrieNode_Cell)
	{
		*pValue = pNode->data;
		return 1;
	}
	
	return 0;
}

static cell_t GetTrieArray(IPluginContext *pContext, const cell_t *params)
{
	Handle_t hndl;
	CellTrie *pTrie;
	HandleError err;
	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);

	hndl = params[1];

	if ((err = g_HandleSys.ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
		!= HandleError_None)
	{
		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
	}

	if (params[4] < 0)
	{
		return pContext->ThrowNativeError("Invalid array size: %d", params[4]);
	}

	char *key;
	cell_t *pValue, *pSize;
	pContext->LocalToString(params[2], &key);
	pContext->LocalToPhysAddr(params[3], &pValue);
	pContext->LocalToPhysAddr(params[5], &pSize);

	SmartTrieNode *pNode;
	if ((pNode = pTrie->trie.retrieve(key)) == NULL
		|| pNode->type != TrieNode_CellArray)
	{
		return 0;
	}

	if (pNode->ptr == NULL)
	{
		*pSize = 0;
		return 1;
	}

	if (pNode->data > params[4])
	{
		*pSize = params[4];
	}
	else if (params[4] != 0)
	{
		*pSize = pNode->data;
	}
	else
	{
		return 1;
	}

	memcpy(pValue, pNode->ptr, sizeof(cell_t) * pSize[0]);

	return 1;
}

static cell_t GetTrieString(IPluginContext *pContext, const cell_t *params)
{
	Handle_t hndl;
	CellTrie *pTrie;
	HandleError err;
	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);

	hndl = params[1];

	if ((err = g_HandleSys.ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
		!= HandleError_None)
	{
		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
	}

	if (params[4] < 0)
	{
		return pContext->ThrowNativeError("Invalid buffer size: %d", params[4]);
	}

	char *key;
	cell_t *pSize;
	pContext->LocalToString(params[2], &key);
	pContext->LocalToPhysAddr(params[5], &pSize);

	SmartTrieNode *pNode;
	if ((pNode = pTrie->trie.retrieve(key)) == NULL
		|| pNode->type != TrieNode_String)
	{
		return 0;
	}

	if (pNode->ptr == NULL)
	{
		*pSize = 0;
		pContext->StringToLocal(params[3], params[4], "");
		return 1;
	}

	size_t written;
	pContext->StringToLocalUTF8(params[3], params[4], (char *)pNode->ptr, &written);

	*pSize = (cell_t)written;

	return 1;
}

static cell_t GetTrieSize(IPluginContext *pContext, const cell_t *params)
{
	Handle_t hndl;
	CellTrie *pTrie;
	HandleError err;
	HandleSecurity sec = HandleSecurity(pContext->GetIdentity(), g_pCoreIdent);

	hndl = params[1];

	if ((err = g_HandleSys.ReadHandle(hndl, htCellTrie, &sec, (void **)&pTrie))
		!= HandleError_None)
	{
		return pContext->ThrowNativeError("Invalid Handle %x (error %d)", hndl, err);
	}

	return pTrie->trie.size();
}

REGISTER_NATIVES(trieNatives)
{
	{"ClearTrie",				ClearTrie},
	{"CreateTrie",				CreateTrie},
	{"GetTrieArray",			GetTrieArray},
	{"GetTrieString",			GetTrieString},
	{"GetTrieValue",			GetTrieValue},
	{"RemoveFromTrie",			RemoveFromTrie},
	{"SetTrieArray",			SetTrieArray},
	{"SetTrieString",			SetTrieString},
	{"SetTrieValue",			SetTrieValue},
	{"GetTrieSize",				GetTrieSize},
	{NULL,						NULL},
};