#include <llama.h>
#include <mtmd.h>

#include "jjml_mtmd.h"

/*
 * FACTORIZE WITH llm ?
 */

void jjml_mtmd_batch_add(struct llama_batch &batch, llama_token id,
		llama_pos pos, const std::vector<llama_seq_id> &seq_ids, bool logits) {
	batch.token[batch.n_tokens] = id;
	batch.pos[batch.n_tokens] = pos;
	batch.n_seq_id[batch.n_tokens] = seq_ids.size();
	for (size_t i = 0; i < seq_ids.size(); ++i) {
		batch.seq_id[batch.n_tokens][i] = seq_ids[i];
	}
	batch.logits[batch.n_tokens] = logits;

	batch.n_tokens++;
}

void jjml_mtmd_batch_clear(struct llama_batch &batch) {
	batch.n_tokens = 0;
}

/*
 * MTMD SPECIFIC
 */

// helper struct to make working with embd batch easier
// note: this will be removed after llama_batch_ext refactoring
struct decode_embd_batch {
	int n_pos_per_embd;
	int n_mmproj_embd;
	std::vector<llama_pos> pos;
	std::vector<llama_pos> pos_view; // used by mrope
	std::vector<int32_t> n_seq_id;
	std::vector<llama_seq_id> seq_id_0;
	std::vector<llama_seq_id*> seq_ids;
	std::vector<int8_t> logits;
	llama_batch batch;
	decode_embd_batch(float *embd, int32_t n_tokens, int n_pos_per_embd,
			int n_mmproj_embd) :
			n_pos_per_embd(n_pos_per_embd), n_mmproj_embd(n_mmproj_embd) {
		pos.resize(n_tokens * n_pos_per_embd);
		n_seq_id.resize(n_tokens);
		seq_ids.resize(n_tokens + 1);
		logits.resize(n_tokens);
		seq_id_0.resize(1);
		seq_ids[n_tokens] = nullptr;
		batch = {
		/*n_tokens       =*/n_tokens,
		/*tokens         =*/nullptr,
		/*embd           =*/embd,
		/*pos            =*/pos.data(),
		/*n_seq_id       =*/n_seq_id.data(),
		/*seq_id         =*/seq_ids.data(),
		/*logits         =*/logits.data(), };
	}

	void set_position_normal(llama_pos pos_0, llama_seq_id seq_id) {
		seq_id_0[0] = seq_id;
		for (int i = 0; i < batch.n_tokens; i++) {
			batch.pos[i] = pos_0 + i;
			batch.n_seq_id[i] = 1;
			batch.seq_id[i] = seq_id_0.data();
			batch.logits[i] = false;
		}
	}

	// M-RoPE for image
	void set_position_mrope_2d(llama_pos pos_0, int nx, int ny,
			llama_seq_id seq_id) {
		GGML_ASSERT(n_pos_per_embd == 4);
		seq_id_0[0] = seq_id;
		for (int y = 0; y < ny; y++) {
			for (int x = 0; x < nx; x++) {
				int i = y * nx + x;
				pos[i] = pos_0;
				pos[i + batch.n_tokens] = pos_0 + y;
				pos[i + batch.n_tokens * 2] = pos_0 + x;
				pos[i + batch.n_tokens * 3] = 0; // last pos dim is unused
			}
		}
		for (int i = 0; i < batch.n_tokens; i++) {
			batch.n_seq_id[i] = 1;
			batch.seq_id[i] = seq_id_0.data();
			batch.logits[i] = false;
		}
	}

	// M-RoPE for audio
	void set_position_mrope_1d(llama_pos pos_0, llama_seq_id seq_id) {
		GGML_ASSERT(n_pos_per_embd == 4);
		seq_id_0[0] = seq_id;
		for (int i = 0; i < batch.n_tokens; i++) {
			pos[i] = pos_0 + i;
			pos[i + batch.n_tokens] = pos_0 + i;
			pos[i + batch.n_tokens * 2] = pos_0 + i;
			pos[i + batch.n_tokens * 3] = 0; // last pos dim is unused
		}
		for (int i = 0; i < batch.n_tokens; i++) {
			batch.n_seq_id[i] = 1;
			batch.seq_id[i] = seq_id_0.data();
			batch.logits[i] = false;
		}
	}

	llama_batch get_view(int offset, int n_tokens) {
		llama_pos *pos_ptr;
		pos_view.clear();
		pos_view.reserve(n_tokens * n_pos_per_embd);
		if (n_pos_per_embd > 1) {
			// mrope
			// for example, with layout of src: 1234...1234...1234...1234...
			//       offset 2 will give us dst: 34...34...34...34...
			for (int i = 0; i < n_pos_per_embd; i++) {
				// assume n_tokens is less than or equal to batch.n_tokens
				// batch.n_tokens is number of **total** tokens
				// n_tokens is number of viewed token
				size_t src_idx = i * batch.n_tokens + offset;
				pos_view.insert(pos_view.end(), pos.data() + src_idx,
						pos.data() + src_idx + n_tokens);
			}
			pos_ptr = pos_view.data();
		} else {
			// normal
			pos_ptr = pos.data() + offset;
		}
		return {
			/*n_tokens       =*/n_tokens,
			/*tokens         =*/nullptr,
			/*embd           =*/batch.embd + offset * n_mmproj_embd,
			/*pos            =*/pos_ptr,
			/*n_seq_id       =*/batch.n_seq_id + offset,
			/*seq_id         =*/batch.seq_id + offset,
			/*logits         =*/batch.logits + offset,
		};
	}
};

int32_t jjml_mtmd_decode_image_chunk(mtmd_context *ctx,
		struct llama_context *lctx, const mtmd_input_chunk *chunk,
		float *encoded_embd, llama_pos n_past, llama_seq_id seq_id,
		int32_t n_batch, llama_pos *new_n_past) {
	auto chunk_type = mtmd_input_chunk_get_type(chunk);
	const char *name =
			chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE ? "image" : "audio";
	if (chunk_type == MTMD_INPUT_CHUNK_TYPE_TEXT) {
//        LOG_ERR("failed to decode chunk: input chunk not of image/audio type\n");
		return -1;
	}

	const llama_model *model = llama_get_model(lctx);
	int n_mmproj_embd = llama_model_n_embd(model);
	int n_pos_per_embd = mtmd_decode_use_mrope(ctx) ? 4 : 1;

	int32_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk);
	int32_t i_batch = 0;
	int32_t n_img_batches = GGML_PAD(n_tokens, n_batch) / n_batch;
	decode_embd_batch batch_embd(encoded_embd, n_tokens, n_pos_per_embd,
			n_mmproj_embd);

	if (mtmd_decode_use_mrope(ctx)) {
		if (chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE) {
			const auto image_tokens = mtmd_input_chunk_get_tokens_image(chunk);
			if (!image_tokens) {
//                LOG_ERR("failed to decode chunk: image tokens are null\n");
				return -1;
			}
			const int nx = mtmd_image_tokens_get_nx(image_tokens);
			const int ny = mtmd_image_tokens_get_ny(image_tokens);
			batch_embd.set_position_mrope_2d(n_past, nx, ny, seq_id);
		} else if (chunk_type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {
			batch_embd.set_position_mrope_1d(n_past, seq_id);
		} else {
			GGML_ABORT("invalid chunk type for M-RoPE");
		}
	} else {
		batch_embd.set_position_normal(n_past, seq_id);
	}
#ifdef JJML_PRE_LLAMA_0_0_8775
	if (mtmd_decode_use_non_causal(ctx)) {
#else
	if (mtmd_decode_use_non_causal(ctx, chunk)) {
#endif
		llama_set_causal_attn(lctx, false);
		// TODO @ngxson : need to make sure only one image is processed at a time, and n_ubatch must be enough to hold the image
	}

	while (i_batch < n_img_batches) { // split into batches
		int pos_offset = i_batch * n_batch;
		int n_tokens_batch = std::min(n_batch, n_tokens - pos_offset);
		llama_batch batch_embd_view = batch_embd.get_view(pos_offset,
				n_tokens_batch);

//        LOG_INF("decoding %s batch %d/%d, n_tokens_batch = %d\n", name, i_batch+1, n_img_batches, n_tokens_batch);

		int64_t t1 = ggml_time_ms();
		int32_t ret = llama_decode(lctx, batch_embd_view);
		if (ret != 0) {
//            LOG_ERR("failed to decode %s\n", name);
			llama_set_causal_attn(lctx, true); // restore causal attn
			return ret;
		}

//        LOG_INF("%s decoded (batch %d/%d) in %" PRId64 " ms\n", name, i_batch+1, n_img_batches, ggml_time_ms() - t1);

		i_batch++;
	}

	n_past += mtmd_input_chunk_get_n_pos(chunk);
	*new_n_past = n_past;

#ifdef JJML_PRE_LLAMA_0_0_8775
	if (mtmd_decode_use_non_causal(ctx)) {
#else
	if (mtmd_decode_use_non_causal(ctx, chunk)) {
#endif
		llama_set_causal_attn(lctx, true);
	}
	return 0;
}

int32_t jjml_mtmd_eval_chunk_single(mtmd_context *ctx,
		struct llama_context *lctx, const mtmd_input_chunk *chunk,
		llama_pos n_past, llama_seq_id seq_id, int32_t n_batch,
		bool logits_last, llama_pos *new_n_past) {
	int32_t ret;
	llama_batch text_batch = llama_batch_init(n_batch, 0, 1);
	auto chunk_type = mtmd_input_chunk_get_type(chunk);

	if (chunk_type == MTMD_INPUT_CHUNK_TYPE_TEXT) {
		size_t n_tokens;
		const auto tokens = mtmd_input_chunk_get_tokens_text(chunk, &n_tokens);
		// LOG_INF("decoding text chunk, n_tokens = %zu\n", n_tokens);
		size_t i = 0;
		while (i < n_tokens) { // split into batches
			text_batch.n_tokens = 0; // clear the batch
			for (; i < n_tokens && text_batch.n_tokens < n_batch; i++) {
				int32_t j = text_batch.n_tokens;
				text_batch.token[j] = tokens[i];
				text_batch.pos[j] = n_past++;
				text_batch.n_seq_id[j] = 1;
				text_batch.seq_id[j][0] = seq_id;
				text_batch.logits[j] = false;

				text_batch.n_tokens++;
			}
			bool is_last_token = (i == n_tokens);
			if (logits_last && is_last_token) {
				text_batch.logits[text_batch.n_tokens - 1] = true;
			}
			ret = llama_decode(lctx, text_batch);
			if (ret != 0) {
//                LOG_ERR("failed to decode text\n");
				llama_batch_free(text_batch);
				return ret;
			}
			*new_n_past += text_batch.n_tokens;
		}

	} else if (chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE
			|| chunk_type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {
		const char *name =
				chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE ? "image" : "audio";
		int64_t t0 = ggml_time_ms();

//        LOG_INF("encoding %s slice...\n", name);

		ret = mtmd_encode_chunk(ctx, chunk);
		if (ret != 0) {
//            LOG_ERR("failed to encode %s slice\n", name);
			llama_batch_free(text_batch);
			return ret;
		}

//        LOG_INF("%s slice encoded in %" PRId64 " ms\n", name, ggml_time_ms() - t0);

		float *embd = mtmd_get_output_embd(ctx);
		ret = jjml_mtmd_decode_image_chunk(ctx, lctx, chunk, embd, n_past,
				seq_id, n_batch, new_n_past);
		if (ret != 0) {
//            LOG_ERR("failed to decode %s\n", name);
			llama_batch_free(text_batch);
			return ret;
		}
	} else {
		GGML_ABORT("chunk type not supported");
	}

	llama_batch_free(text_batch);
	return 0;
}

int32_t jjml_mtmd_eval_chunks(mtmd_context *ctx, struct llama_context *lctx,
		const mtmd_input_chunks *chunks, llama_pos n_past, llama_seq_id seq_id,
		int32_t n_batch, bool logits_last, llama_pos *new_n_past) {
	size_t n_chunks = mtmd_input_chunks_size(chunks);
	if (n_chunks == 0) {
//        LOG_ERR("no chunks to eval\n");
		return 0;
	}

	for (size_t i = 0; i < n_chunks; i++) {
		bool chunk_logits_last = (i == n_chunks - 1) && logits_last;
		auto chunk = mtmd_input_chunks_get(chunks, i);

		int32_t res = jjml_mtmd_eval_chunk_single(ctx, lctx, chunk, n_past,
				seq_id, n_batch, chunk_logits_last, &n_past);
		if (res != 0) {
//            LOG_ERR("failed to eval chunk %zu\n", i);
			return res;
		}
		*new_n_past = n_past;
	}

	return 0;
}
