/*
    Copyright (C) 2008-2024, Millistream Market Data <support@millistream.com>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Lesser General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    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 Lesser General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.

*/

#include "common.h"
#include "decode.h"
#include "encode.h"
#include "templates.h"

static int compare_tags (const void *A, const void *B)
{
	return (((const struct msg_field *)A)->tag - ((const struct msg_field *)B)->tag);
}

static void templates_rebalance (struct templates *templates)
{
	for (int i = 0; i < templates->num; i++) {
		free (templates->templates[i].fields);

		if (templates->templates[i].tags_num == 0) {
			templates->templates[i].fields = NULL;
			continue;
		}

		if (templates->templates[i].tags_num > templates->max_num_tags)
			templates->max_num_tags = templates->templates[i].tags_num;

		templates->templates[i].fields = (struct msg_field *) malloc ((sizeof *templates->templates[i].fields) * templates->templates[i].tags_num);

		for (int j = 0; j < templates->templates[i].tags_num; j++) {
			templates->templates[i].fields[j].pos = j;
			templates->templates[i].fields[j].tag = templates->templates[i].tags[j];
		}

		qsort (templates->templates[i].fields, templates->templates[i].tags_num, sizeof *templates->templates[i].fields, compare_tags);
	}
}

static void templates_update (struct templates * const t, const uint16_t mref, const char *p)
{
	if (mref >= t->num) {
		t->templates = realloc (t->templates, sizeof *t->templates * (mref + 1));
		memset (t->templates + t->num, 0, sizeof *t->templates * ((mref - t->num) + 1));
		t->num = mref + 1;
	}

	p = strchr (p, '|');

	if (p == NULL)
		return;

	p = strchr (p + 1, '|');

	if (p == NULL)
		return;

	p++;

	uint64_t mclass = strtoull (p, NULL, 10);

	p = strchr (p, '|');

	if (p == NULL)
		return;

	p++;

	int tags_num = 0;
	for (const char *e = p; *e != '\0' && *e != '|'; e++) {
		if (*e == ' ') {
			tags_num++;

			while (e[1] == ' ')
				e++;
		}
	}

	if (*p != '\0' && *p != '|')
		tags_num++;

	if (t->templates[mref].tags != NULL) {
		free (t->templates[mref].tags);
		t->templates[mref].tags = NULL;
	}

	t->templates[mref].mclass = mclass;
	t->templates[mref].tags_num = tags_num;

	if (tags_num != 0)
		t->templates[mref].tags = (uint32_t *) malloc (sizeof (uint32_t) * tags_num);

	uint32_t *tags = t->templates[mref].tags;
	for (int i = 0; i < tags_num; i++) {
		while (*p == ' ')
			p++;

		tags[i] = strtoul (p, NULL, 10);

		p = strchr (p, ' ');
	}
}

void templates_decode (uint8_t *data, const uint8_t * const end, struct templates ** const templates, struct templates ** const templates_old)
{
	templates_free (*templates_old);
	*templates_old = *templates;

	struct templates *t = malloc (sizeof *t);
	t->templates = NULL;
	t->num = 0;
	t->max_num_tags = 0;

	/* the templates message will always be a fully tagged message consting of only
	 * plain strings and compressed strings. */

	data++; /* skip initial pmap */

	do {
		uint8_t field_code = *data++;
		uint32_t tag;

		if (field_code == 251) { /* Tagged field uint8 */
			tag = *data++;
		} else if (field_code == 252) { /* Tagged field uint16 */
			uint16_t val;
			memcpy (&val, data, sizeof val);
			tag = le16_to_cpu (val);
			data += 2;
		} else if (field_code == 253) { /* Tagged field uint32 */
			uint32_t val;
			memcpy (&val, data, sizeof val);
			tag = le32_to_cpu (val);
			data += 4;
		}

		field_code = *data++;

		if (field_code == 241) { /* string */
			const size_t slen = strlen ((const char *) (data)) + 1;

			templates_update (t, tag, (char *)data);

			data += slen;
		} else if (field_code == 242) { /* compressed string */
			unsigned long long slen, dlen = 0;

			do {
				if (data >= end)
					break;

				dlen *= 128;
				dlen += (*data & 127) + 1;
			} while ((*data++ & 128) == 0);

			uint8_t *p;

			if (dlen == 1) { /* zstd */
				dlen = ZSTD_getFrameContentSize (data, end - data);
				slen = ZSTD_findFrameCompressedSize (data, end - data);

				p = malloc (dlen + 1);

				dlen = ZSTD_decompress (p, dlen, data, slen);

				if (ZSTD_isError (dlen)) {
					free (p);
					break;
				}
			} else {
				slen = end - data;
				p = malloc (dlen + 1);

				if (zlib_uncompress (data, p, &slen, &dlen) != Z_OK) {
					free (p);
					break;
				}
			}

			p[dlen] = '\0';

			templates_update (t, tag, (char *)p);

			free (p);

			data += slen;
		}
	} while (data < end);

	templates_rebalance (t);

#ifdef __clang__
	atomic_store_explicit ((struct templates * _Atomic *)templates, t, memory_order_release);
#else
	atomic_store_explicit (templates, t, memory_order_release);
#endif
}

void templates_initialize (struct templates ** const templates)
{
	struct templates *t = malloc (sizeof *t);
	t->templates = calloc (3, sizeof *t->templates);
	t->num = 3;
	t->max_num_tags = 0;

	t->templates[MDF_M_MESSAGESREFERENCE].mclass = MDF_MC_UNDEF;
	t->templates[MDF_M_MESSAGESREFERENCE].tags_num = 1;
	t->templates[MDF_M_MESSAGESREFERENCE].tags = (uint32_t *) malloc (sizeof (uint32_t) * t->templates[MDF_M_MESSAGESREFERENCE].tags_num);
	t->templates[MDF_M_MESSAGESREFERENCE].tags[0] = 0xffffffff;

	t->templates[MDF_M_LOGON].mclass = MDF_MC_UNDEF;
	t->templates[MDF_M_LOGON].tags_num = 3;
	t->templates[MDF_M_LOGON].tags = (uint32_t *) malloc (sizeof (uint32_t) * t->templates[MDF_M_LOGON].tags_num);
	t->templates[MDF_M_LOGON].tags[0] = MDF_F_USERNAME;
	t->templates[MDF_M_LOGON].tags[1] = MDF_F_PASSWORD;
	t->templates[MDF_M_LOGON].tags[2] = MDF_F_EXTRACREDENTIAL;

	t->templates[MDF_M_LOGOFF].mclass = MDF_MC_UNDEF;
	t->templates[MDF_M_LOGOFF].tags_num = 1;
	t->templates[MDF_M_LOGOFF].tags = (uint32_t *) malloc (sizeof (uint32_t) * t->templates[MDF_M_LOGOFF].tags_num);
	t->templates[MDF_M_LOGOFF].tags[0] = MDF_F_LOGOFFREASON;

	templates_rebalance (t);
	*templates = t;
}

void templates_free (struct templates *templates)
{
	if (templates == NULL)
		return;

	if (templates->templates != NULL) {
		for (int i = 0; i < templates->num; i++) {
			free (templates->templates[i].tags);
			free (templates->templates[i].fields);
		}

		free (templates->templates);
	}

	free (templates);
}
