#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        4096
#define N_BATCH      2048
#define N_THREADS       4
#define N_PREDICT     256

#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);
		return NULL;
	}
	
	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);
		return NULL;
	}

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

int main(int argc, char *argv[])
{
	int i, n_images, rc;
	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 n_bitmaps = 0;
	mtmd_input_chunks *chunks = NULL;
	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;

	llama_pos n_past;

	rc = 1;

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

	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 (n_bitmaps = 0; n_bitmaps < n_images; n_bitmaps++) {
		wrapper = mtmd_helper_bitmap_init_from_file(
		                 mctx, 
		                 argv[n_bitmaps + 1], 
		                 false, 
		                 init_opt);

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

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

	{
		mtmd_input_text text;

		memset(&text, 0, sizeof(text));
		text.text = prompt;
		text.text_len = strlen(prompt);
		text.add_special = true;   /* use model's BOS/EOS tokens */
		text.parse_special = true; /* treat BOS/EOS tokens as special (Qwen/Llama) */
		
		chunks = mtmd_input_chunks_init();
		if (mtmd_tokenize(mctx, chunks, &text, 
			(const mtmd_bitmap **)bitmaps, n_images) != 0) {
			fprintf(stderr, "mtmd_tokenize error\n");
			goto cleanup;
		}
	}

	/* Encode images + decode prompt; logits for the last one */
	if (mtmd_helper_eval_chunks(mctx, ctx, chunks, 0, 0, N_BATCH, 
		true, &n_past) != 0) {
		fprintf(stderr, "Failed to evaluate prompt/images\n");
		goto cleanup;
	}

	/* Generation loop */
	{
		int t, n;
		llama_token token;

		rc = 0;
		for (t = 0; t < N_PREDICT && n_past < N_CTX; t++) {
			token = llama_sampler_sample(smpl, ctx, -1);
			if (llama_vocab_is_eog(vocab, token))
				break;

			n = 0;
			char buf[256];
			n = llama_token_to_piece(vocab, token, buf, sizeof(buf), 0, true);
			if (n > 0)
				fwrite(buf, 1, n, stdout);
			fflush(stdout);

			if (llama_decode(ctx, llama_batch_get_one(&token, 1)) != 0) {
				rc = 1;
				fprintf(stderr, "llama_decode error\n");
				break;
			}
			
			n_past++;
		}
		printf("\n");
	}

cleanup:
	if (bitmaps) {
		for (i = 0; i < n_bitmaps && bitmaps[i]; i++)
			mtmd_bitmap_free(bitmaps[i]);
		free(bitmaps);
	}

	if (chunks)
		mtmd_input_chunks_free(chunks);
	if (smpl)
		llama_sampler_free(smpl);
	if (ctx)
		llama_free(ctx);
	if (model)
		llama_model_free(model);

	llama_backend_free();

	return rc;
}