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

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

#define SYSTEM_PROMPT \
    "You are a lexicographer. Define the target word strictly in the context provided.\n\n" \
    "Output format strictly as follows:\n" \
    "<WORD> (<part of speech>) — <Formality: Conversational|Formal|Literary|Archaic>\n\n" \
    "DEFINITION:\n" \
    "<terse, highly accurate definition>\n\n" \
    "WHEN TO USE:\n" \
    "<1 sentence on when to use>\n\n" \
    "EXAMPLES:\n" \
    "1. <example 1>\n" \
    "2. <example 2>\n\n" \
    "EXAMPLE INPUT:\n" \
    "Define torrent as in 'a torrent of rain poured down'\n\n" \
    "EXAMPLE OUTPUT:\n" \
    "TORRENT (noun) — Literary\n\n" \
    "DEFINITION:\n" \
    "A strong and fast-moving stream of water or other liquid.\n\n" \
    "WHEN TO USE:\n" \
    "Use when describing a sudden, overwhelming rush of water or emotion.\n\n" \
    "EXAMPLES:\n" \
    "1. A torrent of rain poured down on the village.\n" \
    "2. She faced a torrent of angry emails after the announcement.\n\n" \
    "CRITICAL: Do NOT echo the user's input, context sentence, or 'DEFINE' command. Start immediately with <WORD>."

#define MAX_TOKENS  300

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

	llama_token *prompt_tokens;
	llama_token new_token_id;

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

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

	const char *tmpl = llama_model_chat_template(model, NULL);
	if (!tmpl)
		fprintf(stderr, "Warning: model has no embedded chat template\n");

	prompt_len = llama_chat_apply_template(tmpl, messages, 2, true, NULL, 0);
	if (prompt_len <= 0) {
		fprintf(stderr, "Error: failed to calculate chat template size\n");
		return;
	}

	prompt = MALLOC((size_t)prompt_len + 1);
	if (llama_chat_apply_template(NULL, messages, 2, true, prompt, prompt_len + 1) < 0) {
		fprintf(stderr, "Error: failed to apply chat template\n");
		free(prompt);
		return;
	}

	cparams = llama_context_default_params();
	cparams.n_ctx = 1024;    /* context size in tokens */
	cparams.n_threads = 4;
	cparams.n_threads_batch = 4;

	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);
	if (!vocab) {
		fprintf(stderr, "Error: failed to obtain model vocabulary\n");
		free(prompt);
		llama_free(ctx);
		return;
	}

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

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

	if (n_prompt_tokens + MAX_TOKENS > (int)cparams.n_ctx) {
		fprintf(stderr, "Error: token count exceeds context size\n");
		free(prompt);
		llama_free(ctx);
		return;
	}

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

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

	free(prompt); /* Prompt buffer is fully tokenized and no longer needed */

	/* Ingest prompt tokens in one batch (parallelizes matrix 
	 * multiplications across tokens in the batch) */
	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_penalties(
		64,      /* last_n: lookback window (64 is standard) */
		1.1f,    /* repeat_penalty */
		0.0f,    /* frequency_penalty */
		0.0f     /* presence_penalty */
	));

	/* Pick the top token */
	llama_sampler_chain_add(smpl, llama_sampler_init_greedy());

	for (i = 0; i < MAX_TOKENS; i++) {	
		/* Model outputs next tokens for every token in the prompt. 
		 * We need the one after the last token in the prompt */
		new_token_id = llama_sampler_sample(smpl, ctx, -1);

		/* Tell the sampler chain which token was chosen */
		llama_sampler_accept(smpl, new_token_id);

		/* Check for end-of-generation (EOG) tokens:
		 * EOS: end-of-sequence
		 * EOT: end-of-turn */
		if (llama_vocab_is_eog(vocab, new_token_id))
			break;

		/* Convert numeric token id to printable text */
		char buf[128];
		int n = llama_token_to_piece(vocab, new_token_id, buf, 
			sizeof(buf), 0, false);

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

		/* Create batch with 1 token for the next forward pass */
		batch = llama_batch_get_one(&new_token_id, 1);

		if (llama_decode(ctx, batch) != 0) {
			fprintf(stderr, "llama_decode failed!\n");
			break;
		}
	}

	printf("\n\n");
	fflush(stdout);

	//llama_perf_context_print(ctx);

	llama_sampler_free(smpl);
	llama_free(ctx);
}

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

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

	const char *model_path = argv[1];

	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");

	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 *prompt = argv[2];
	process_request(model, prompt);

	llama_model_free(model);
	llama_backend_free();

	return 0;
}