#include <stdio.h>
#include <err.h>
#include <string.h>
#include <syslog.h>

#include "llm.h"
#include "mem.h"

#define MAX_TOKENS    300
#define N_CTX        1024
#define N_THREADS       4

#define SYSTEM_PROMPT \
    "You are a dictionary assistant. Define the target word in the context provided.\n" \
    "DO NOT repeat target word in definition or synonyms. Choose one or two of Formality levels.\n\n" \
    "Output format:\n" \
    "<WORD> (<part of speech>) - <Formality: Conversational|Formal|Literary|Figurative|Archaic>\n\n" \
    "DEFINITION:\n" \
    "<terse, highly accurate definition in context>\n\n" \
    "GENERAL TONE:\n" \
    "<One terse sentence describing general tone>\n\n" \
    "SYNONYMS:\n" \
    "<Three to four standard conversational synonyms>\n\n" \
    "EXAMPLES:\n" \
    "1. <example 1>\n" \
    "2. <example 2>\n\n"

struct llm_ctx {
	struct llama_model *model;
	struct llama_context *ctx;
	const struct llama_vocab *vocab;
	struct llama_sampler *smpl;

	/* Token history of what is currently resident in ctx's KV
	 * cache for sequence 0, in order. Used to compute a
	 * longest-common-prefix against each new request so we only
	 * decode the tail that actually changed. */
	llama_token *cached_tokens;
	int n_cached_tokens;

	/* Token count of the system prompt.Used as a sanity-check floor: 
	 * if a later request's LCP against the cache falls below this, the
	 * system-prompt prefix itself failed to match, which would
	 * indicate template drift. */
	int n_system_tokens;
};

static void llm_log_cb(enum ggml_log_level level, const char *s, void *ud)
{
	(void)ud;

	switch (level) {
	case GGML_LOG_LEVEL_WARN:
		syslog(LOG_WARNING, "%s", s);
		return;
	case GGML_LOG_LEVEL_ERROR:
		syslog(LOG_ERR, "%s", s);
		return;
	default:
		return;
	}
}

/* Tokenize `text` (already a fully rendered chat-template prompt) into 
 * a freshly malloc'd buffer. Returns token count via *out_n, or -1 on 
 * failure (*out is left untouched). Caller owns *out. */
static int tokenize_prompt(const struct llama_vocab *vocab,
	const char *text, int text_len, llama_token **out) 
{
	int n;

	n = -llama_tokenize(vocab, text, text_len, NULL, 0, false, true);
	if (n <= 0)
		return -1;

	*out = MALLOC((size_t)n * sizeof(llama_token));

	if (llama_tokenize(vocab, text, text_len, *out, n, false, true) < 0) {
		free(*out);
		return -1;
	}

	return n;
}

struct llm_ctx *llm_init(const char *model_path)
{
	char *sys_prompt;
	int sys_prompt_len;
	llama_token *sys_tokens;

	struct llm_ctx *llm;
	struct llama_model_params mparams;
	struct llama_context_params cparams;
	struct llama_sampler_chain_params sparams;
	struct llama_chat_message sys_msg[1];

	llama_log_set(llm_log_cb, NULL);
	llama_backend_init();

	mparams = llama_model_default_params();
	mparams.n_gpu_layers = 0;   /* force all layers onto CPU */

	/* LLAMA_LOAD_MODE_MLOCK (force model to RAM, no swap/compression)
	 * would have been preferrable. But my Raspberry Pi5 doesn't have
	 * enough memory. */
	mparams.load_mode = LLAMA_LOAD_MODE_NONE;

	llm = MALLOC(sizeof(*llm));
	llm->model = llama_model_load_from_file(model_path, mparams);
	if (!llm->model) {
		free(llm);
		return NULL;
	}

	if (!llama_model_chat_template(llm->model, NULL))
		syslog(LOG_WARNING, "model has no embedded chat template");

	/* Creation of context allocates KV cache and compute buffers.
	 * Do this once in init, instead of per request. */
	cparams = llama_context_default_params();
	cparams.n_ctx = N_CTX;
	cparams.n_threads = N_THREADS;
	cparams.n_threads_batch = N_THREADS;

	llm->ctx = llama_init_from_model(llm->model, cparams);
	if (!llm->ctx) {
		llama_model_free(llm->model);
		free(llm);
		return NULL;
	}

	llm->vocab = llama_model_get_vocab(llm->model);
	if (!llm->vocab) {
		llama_free(llm->ctx);
		llama_model_free(llm->model);
		free(llm);
		return NULL;
	}

	/* Sampler chain, reuse across requests as well. */
	sparams = llama_sampler_chain_default_params();
	llm->smpl = llama_sampler_chain_init(sparams);

	llama_sampler_chain_add(llm->smpl, 
		llama_sampler_init_penalties(
			64,      /* last_n: lookback window */
			1.1f,    /* repeat_penalty */
			0.0f,    /* frequency_penalty */
			0.0f     /* presence_penalty */
		)
	);

	llama_sampler_chain_add(llm->smpl, llama_sampler_init_greedy());

	llm->cached_tokens = MALLOC((size_t)N_CTX * sizeof(llama_token));
	llm->n_cached_tokens = 0;

	/* Measure system prompt's token count. Used to check template 
	 * drift in llm_run(), not seed the KV cache. */
	sys_msg[0].role = "system";
	sys_msg[0].content = SYSTEM_PROMPT;

	sys_prompt_len = llama_chat_apply_template(NULL, sys_msg, 1, false, NULL, 0);
	if (sys_prompt_len > 0) {	
		sys_prompt = MALLOC((size_t)sys_prompt_len + 1);
		if (llama_chat_apply_template(NULL, sys_msg, 1, false, 
			sys_prompt, sys_prompt_len + 1) >= 0) {
			llm->n_system_tokens = tokenize_prompt(llm->vocab, 
				sys_prompt, sys_prompt_len, &sys_tokens);
			if (llm->n_system_tokens > 0)
				free(sys_tokens);
			else
				llm->n_system_tokens = 0;
		} else {
			llm->n_system_tokens = 0;
		}
		free(sys_prompt);
	} else {
		llm->n_system_tokens = 0;
	}

	return llm;
}

void llm_free(struct llm_ctx *llm)
{
	if (!llm)
		return;
	llama_sampler_free(llm->smpl);
	llama_free(llm->ctx);
	llama_model_free(llm->model);
	llama_backend_free();
	free(llm->cached_tokens);
	free(llm);
}

void llm_run(struct llm_ctx *llm, const char *user_prompt, FILE *out)
{
	int i, n;
	char *prompt;
	int prompt_len;
	llama_token *new_tokens;
	int n_new_tokens;
	int n_common;
	llama_token new_token_id;
	struct llama_batch batch;

	struct llama_chat_message messages[] = {
		{ "system", SYSTEM_PROMPT },
		{ "user",   user_prompt }
	};

	/* Retokenize full prompt. Let's expensive than decode; no
	 * assumptions about the model's internal chat structure. */
	prompt_len = llama_chat_apply_template(NULL, messages, 2, true, NULL, 0);
	if (prompt_len <= 0) {
		syslog(LOG_ERR, "chat template size calculation failed");
		return;
	}

	prompt = MALLOC((size_t)prompt_len + 1);
	if (llama_chat_apply_template(NULL, messages, 2, true, prompt, 
		prompt_len + 1) < 0) {
		syslog(LOG_ERR, "failed to apply chat template");
		free(prompt);
		return;
	}

	n_new_tokens = tokenize_prompt(llm->vocab, prompt, prompt_len, &new_tokens);
	free(prompt);
	if (n_new_tokens <= 0) {
		syslog(LOG_ERR, "tokenization failed");
		return;
	}

	if (n_new_tokens + MAX_TOKENS > N_CTX) {
		syslog(LOG_ERR, "token count exceeds context window");
		free(new_tokens);
		return;
	}

	/* llm->cached_tokens contain tokens from the previous request 
	 * (system prompt + any tokens inserted by the model + previous 
	 * user prompt + closing tokens inserted by the model). new_tokens
	 * should have a similar structure, except the new user prompt. We
	 * walk the two token sets to locate the first new token: presumably
	 * the start of the new user prompt. */
	n_common = 0;
	while (n_common < llm->n_cached_tokens && n_common < n_new_tokens && 
		llm->cached_tokens[n_common] == new_tokens[n_common])
		n_common++;

	/* Sanity check */
	if (llm->n_cached_tokens > 0 && n_common < llm->n_system_tokens) {
		/* Template may be rendering inconsistently between requests */
		syslog(LOG_WARNING, "cached prefix is shorter than the system prompt");
	}

	/* Force at least one real decode every request. If n_common hit
	 * n_new_tokens, decode gets skipped entirely, so ctx keeps the
	 * stale logits from the previous request's EOG -- sampling them
	 * again just reproduces EOG immediately (empty response on an
	 * exact repeat). */
	if (n_common == n_new_tokens)
		n_common--;

	/* Drop everything in the cache past the common prefix. */
	llama_memory_seq_rm(llama_get_memory(llm->ctx), 
		0,          /* sequence id: 0 - we only use one sequence */
		n_common,   /* start at first differing token */
		-1          /* -1: clear to the end of sequence */
	);
	
	if (n_common < n_new_tokens) {
		/* Only decode the diff */
		batch = llama_batch_get_one(new_tokens + n_common, 
			n_new_tokens - n_common);
		if (llama_decode(llm->ctx, batch) != 0) {
			syslog(LOG_ERR, "prompt evaluation failed");
			free(new_tokens);
			llm->n_cached_tokens = n_common; 
			return;
		}

		/* Copy the newly tokens to the buffer */
		memcpy(llm->cached_tokens + n_common, new_tokens + n_common, 
			(size_t)(n_new_tokens - n_common) * sizeof(llama_token));
		llm->n_cached_tokens = n_new_tokens;
	}

	free(new_tokens);

	/* Sampler holds repetition-window state (from the penalties
	 * stage) across calls since we no longer rebuild it each
	 * request. Reset that state so one request's repeated words
	 * don't suppress an unrelated word in the next one. */
	llama_sampler_reset(llm->smpl);

	/* Generation loop */
	for (i = 0; i < MAX_TOKENS; i++) {
		new_token_id = llama_sampler_sample(llm->smpl, llm->ctx, -1);
		llama_sampler_accept(llm->smpl, new_token_id);

		if (llama_vocab_is_eog(llm->vocab, new_token_id))
			break;

		char buf[128];
		n = llama_token_to_piece(llm->vocab, new_token_id, buf, 
				sizeof(buf), 0, false);
		if (n > 0) {
			fwrite(buf, 1, (size_t)n, out);
			fflush(out);
		}
		
		batch = llama_batch_get_one(&new_token_id, 1);
		if (llama_decode(llm->ctx, batch) != 0) {
			syslog(LOG_ERR, "decode() error in generation loop");
			break;
		}

		/* New token is now in KV cache. Record it for the next 
		 * request's LCP check. */
		llm->cached_tokens[llm->n_cached_tokens++] = new_token_id;
	}

	fprintf(out, "\n\n");
	fflush(out);
}