#include <stdio.h>
#include <stdbool.h>
#include <err.h>
#include <unistd.h>

#include "llama.h"
#include "mem.h"

#define MODEL_PATH "Qwen2.5-7B-Instruct-Q4_K_M.gguf"

#define SYSTEM_PROMPT \
	"You are a dictionary assistant. " \
	"Given a target word, provide its definition and formality on a single line, " \
	"followed by a blank line, and then a single example sentence.\n\n" \
	"Output format:\n" \
	"<word>: <terse definition>. <Formality>.\n\n" \
	"Example: <example sentence>\n\n" \
	"Formality must be one of: Formal, Literary, Archaic, or Conversational. " \
	"Keep the definition terse. Do not include any intro, markdown, or extra chatter."

#define CHATML_FMT \
	"<|im_start|>system\n%s<|im_end|>\n" \
	"<|im_start|>user\nDefine the word: %s<|im_end|>\n" \
	"<|im_start|>assistant\n"

#define MAX_TOKENS 200

static int build_prompt(char **buf, const char *word)
{
	char *p;
	int len, written;

	if (!buf || !word)
		return -1;

	len = snprintf(NULL, 0, CHATML_FMT, SYSTEM_PROMPT, word);
	if (len <= 0)
		return -1;

	p = MALLOC((size_t)len + 1);

	written = snprintf(p, (size_t)len + 1, CHATML_FMT, SYSTEM_PROMPT, word);
	if (written <= 0) {
		free(p);
		return -1;
	}
	
	*buf = p;
	return written;
}

static void process_request(struct llama_model *model, const char *word)
{
	char *prompt;
	int prompt_len;
	int n_prompt_tokens, i;

	struct llama_context_params cparams;
	struct llama_context *ctx;
	llama_token *prompt_tokens;
	struct llama_batch batch;
	struct llama_sampler_chain_params sparams;
	struct llama_sampler *smpl;

	llama_token new_token_id;

	prompt = NULL;
	prompt_len = build_prompt(&prompt, word);

	if (prompt_len <= 0) {
		fprintf(stderr, "Error: failed to construct prompt for '%s'\n", word);
		return;
	}
	
	cparams = llama_context_default_params();
	cparams.n_ctx = 512; /* keep context small to keep RAM usage low */
	cparams.n_threads = 4; /* todo: tune */
	
	ctx = llama_init_from_model(model, cparams);
	if (!ctx) {
		fprintf(stderr, "Error: failed to create context\n");
		free(prompt);
		return;
	}

	const struct llama_vocab *vocab = llama_model_get_vocab(model);

	n_prompt_tokens = -llama_tokenize(vocab, prompt, prompt_len, NULL, 
		0, true, true);

	if (n_prompt_tokens <= 0) {
		fprintf(stderr, "Error: tokenization sizing failed\n");
		llama_free(ctx);
		free(prompt);
		return;
	}

	prompt_tokens = MALLOC((size_t)n_prompt_tokens * sizeof(llama_token));

	if (llama_tokenize(vocab, prompt, prompt_len, prompt_tokens, 
		n_prompt_tokens, true, true) < 0) {
		fprintf(stderr, "Error: Tokenization execution failed\n");
		free(prompt_tokens);
		llama_free(ctx);
		free(prompt);
		return;
	}

	free(prompt); /* no longer required */

	batch = llama_batch_get_one(prompt_tokens, n_prompt_tokens);
	if (llama_decode(ctx, batch) != 0) {
		fprintf(stderr, "Error: Prompt evaluation failed\n");
		free(prompt_tokens);
		llama_free(ctx);
		return;
	}

	free(prompt_tokens);

	/* Generation loop */
	sparams = llama_sampler_chain_default_params();
	smpl = llama_sampler_chain_init(sparams);
	llama_sampler_chain_add(smpl, llama_sampler_init_temp(0.2f)); /* todo: tune */
	llama_sampler_chain_add(smpl, llama_sampler_init_dist(LLAMA_DEFAULT_SEED));

	for (i = 0; i < MAX_TOKENS; i++) {	
		new_token_id = llama_sampler_sample(smpl, ctx, -1);
		if (llama_vocab_is_eog(vocab, new_token_id))
			break;

		char buf[128];
		int n = llama_token_to_piece(vocab, new_token_id, buf, 
			sizeof(buf), 0, true);

		if (n > 0) {
			fwrite(buf, 1, (size_t)n, stdout);
			fflush(stdout);
		}

		batch = llama_batch_get_one(&new_token_id, 1);
		if (llama_decode(ctx, batch) != 0)
			break;
	}

	llama_sampler_free(smpl);
	llama_free(ctx);
}

int main(int argc , char *argv[])
{
	struct llama_model *model;
	struct llama_model_params mparams;

	if (unveil(MODEL_PATH, "r") == -1)
		err(1, "unveil %s failed", MODEL_PATH);

	if (unveil(NULL, NULL) == -1)
		err(1, "unveil lock failed");

	if (pledge("stdio rpath", NULL) == -1)
		err(1, "initial pledge failed");

	if (argc < 2)
		errx(1, "usage: %s [prompt]", argv[0]);

	llama_backend_init();

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

	model = llama_model_load_from_file(MODEL_PATH, mparams);
	if (!model)
		errx(1, "failed to load model from file %s", MODEL_PATH);

	if (pledge("stdio", NULL) == -1)
		err(1, "secondary pledge failed");

	const char *word = argv[1];
	process_request(model, word);

	llama_model_free(model);
	llama_backend_free();

	return 0;
}