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

#include "mem.h"
#include "mtmd.h"
#include "mtmd-helper.h"
#include "llama.h"

#define N_CTX        1024
#define N_THREADS       4

#define MODEL_PATH   "model.gguf"
#define MMPROJ_PATH  "mmproj.gguf"

#define PROMPT  \
    "These are consecutive frames of a video." \
    "Describe what happens."

static char * build_prompt(struct llama_model *model, int n_images)
{
	int n, i;
	const char *tmpl, *marker;
	char *content, *buf;
	size_t content_len, buf_len;
	struct llama_chat_message msg;

	tmpl = llama_model_chat_template(model, NULL);
	if (!tmpl) {
		fprintf(stderr, "Failed to load chat template\n");
		return NULL;
	}

	marker = mtmd_default_marker();
	if (!marker) {
		fprintf(stderr, "Null default media marker\n");
		return NULL;
	}

	content_len = n_images * (strlen(marker) + 1) + strlen(PROMPT) + 1;
	content = MALLOC(content_len);
	content[0] = '\0';

	for (i = 0; i < n_images; i++) {
		strcat(content, marker); /* Media marker string <__media__>*/
		strcat(content, "\n");
	}
	strcat(content, PROMPT);

	msg.role = "user";
	msg.content = content;

	n = llama_chat_apply_template(tmpl, &msg, 1, true, NULL, 0);
	if (n < 0) {
		fprintf(stderr, "Chat template size check failed\n");
		free(content);
	}
	
	buf_len = n + 1;
	buf = MALLOC(buf_len);
	n = llama_chat_apply_template(tmpl, &msg, 1, true, buf, buf_len);
	if (n < 0) {
		fprintf(stderr, "Failed to apply chat template\n");
		free(content);
	}

	buf[n] = '\0';
	return buf;
}

int main(int argc, char *argv[])
{
	int n_images;
	const char *prompt;
	const struct llama_vocab *vocab;

	struct llama_model *model = NULL;
	struct llama_model_params mparams;
	struct llama_context *ctx = NULL;
	struct llama_context_params cparams;
	struct llama_sampler *smpl = NULL;
	struct llama_sampler_chain_params sparams;

	int i;
	mtmd_input_chunks *chunks;
	mtmd_bitmap **bitmaps = NULL;
	struct mtmd_context *mctx = NULL;
	struct mtmd_context_params mtparams;
	struct mtmd_helper_init_opt init_opt;
	struct mtmd_helper_bitmap_wrapper wrapper;

	n_images = argc - 1;
	if (n_images < 1) {
		fprintf(stderr, "Usage: %s image1 [image2 ...]\n", argv[0]);
		return 1;
	}

	llama_backend_init();

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

	model = llama_model_load_from_file(MODEL_PATH, mparams);
	if (!model) {
		fprintf(stderr, "Failed to load model: %s\n", MODEL_PATH);
		goto cleanup;
	}

	if (!llama_model_chat_template(model, NULL))
		warnx("No embedded chat template");

	cparams = llama_context_default_params();
	cparams.n_ctx = N_CTX;
	cparams.n_threads = N_THREADS;
	cparams.n_threads_batch = N_THREADS;

	ctx = llama_init_from_model(model, cparams);
	if (!ctx) {
		fprintf(stderr, "Failed to initialize context\n");
		goto cleanup;
	}

	vocab = llama_model_get_vocab(model);
	if (!vocab) {
		fprintf(stderr, "Failed to initialize context\n");
		goto cleanup;
	}

	sparams = llama_sampler_chain_default_params();
	smpl = llama_sampler_chain_init(sparams);
	if (!smpl) {
		fprintf(stderr, "Failed to initialize sampler\n");
		goto cleanup;
	}

	llama_sampler_chain_add(smpl, 
		llama_sampler_init_penalties(
		    llama_vocab_n_tokens(vocab),
			64,      /* last_n: lookback window */
			1.1f,    /* repeat_penalty */
			0.0f,    /* frequency_penalty */
			0.0f     /* presence_penalty */
		)
	);
	llama_sampler_chain_add(smpl, llama_sampler_init_greedy());

	mtparams = mtmd_context_params_default();
	mtparams.use_gpu = false;
	mtparams.n_threads = N_THREADS;
	mctx = mtmd_init_from_file(MMPROJ_PATH, model, mtparams);
	if (!mctx) {
		fprintf(stderr, "Failed to load mmproj: %s\n", MMPROJ_PATH);
		goto cleanup;
	}

	init_opt = mtmd_helper_init_opt_default();
	bitmaps = CALLOC(n_images, sizeof(*bitmaps));
	for (i = 0; i < n_images; i++) {
		wrapper = mtmd_helper_bitmap_init_from_file(
		                 mctx, 
		                 argv[i + 1], 
		                 false, 
		                 init_opt);

		bitmaps[i] = wrapper.bitmap;
		if (!bitmaps[i]) {
			fprintf(stderr, "Failed to load image: %s\n", argv[i + 1]);
			goto cleanup;
		}
	}

	prompt = build_prompt(model, n_images);
	if (!prompt) {
		fprintf(stderr, "Prompt error");
		goto cleanup;
	}

	printf("Init success\n");

cleanup:
	if (bitmaps) {
		for (i = 0; bitmaps[i]; i++)
			mtmd_bitmap_free(bitmaps[i]);
		free(bitmaps);
	}
	if (smpl)
		llama_sampler_free(smpl);
	if (ctx)
		llama_free(ctx);
	if (model)
		llama_model_free(model);
	llama_backend_free();

	return 0;
}