#include <linux/slab.h>
#include <linux/mutex.h>
#include <linux/crc32.h>
#include <linux/kthread.h>
#include <linux/completion.h>

#include "nvidia-vgpu-vfio.h"
#include "migration.h"

static int nv_vgpu_migration_save_thread(void *data)
{
	struct nv_vgpu_device *vgpu = data;
	struct nv_vgpu_migration_ctx *ctx = vgpu->migration.ctx;
	u8 *buffer;
	size_t size;
	int ret;

	ret = nv_vgpu_migration_save_state(vgpu, &buffer, &size);
	if (ret) {
		ctx->state.result = ret;
	} else {
		ctx->state.data = buffer;
		ctx->state.data_size = size;
		ctx->state.result = 0;
	}

	complete(&ctx->state.done);
	return 0;
}

static int nv_vgpu_migration_restore_thread(void *data)
{
	struct nv_vgpu_device *vgpu = data;
	struct nv_vgpu_migration_ctx *ctx = vgpu->migration.ctx;
	int ret;

	ret = nv_vgpu_migration_restore_state(vgpu, ctx->state.data, ctx->state.data_size);
	ctx->state.result = ret;
	complete(&ctx->state.done);
	return 0;
}

int nv_vgpu_migration_init(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_migration_ctx *ctx;

	ctx = kzalloc(sizeof(*ctx), GFP_KERNEL);
	if (!ctx)
		return -ENOMEM;

	ctx->vgpu = vgpu;
	mutex_init(&ctx->lock);
	init_completion(&ctx->state.done);
	ctx->in_progress = false;

	vgpu->migration.ctx = ctx;
	return 0;
}

void nv_vgpu_migration_cleanup(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_migration_ctx *ctx = vgpu->migration.ctx;

	if (!ctx)
		return;

	if (ctx->save_thread) {
		kthread_stop(ctx->save_thread);
		ctx->save_thread = NULL;
	}
	if (ctx->restore_thread) {
		kthread_stop(ctx->restore_thread);
		ctx->restore_thread = NULL;
	}

	if (ctx->state.data) {
		kfree(ctx->state.data);
		ctx->state.data = NULL;
	}

	kfree(ctx);
	vgpu->migration.ctx = NULL;
}

int nv_vgpu_migration_save_state(struct nv_vgpu_device *vgpu, u8 **data, size_t *size)
{
	struct nv_vgpu_migration_header *hdr;
	struct nv_vgpu_migration_section *sections;
	struct nv_vgpu_migration_device_state *dev_state;
	struct nv_vgpu_migration_bar_state *bar_state;
	u8 *buffer;
	size_t total_size = 0;
	int i, num_sections = 0;

	mutex_lock(&vgpu->lock);

	total_size = sizeof(struct nv_vgpu_migration_header) +
		     (8 * sizeof(struct nv_vgpu_migration_section));

	for (i = 0; i < NV_VGPU_MAX_REGIONS; i++) {
		if (vgpu->bars[i].mapped) {
			total_size += sizeof(struct nv_vgpu_migration_bar_state) + vgpu->bars[i].size;
			num_sections++;
		}
	}

	buffer = kmalloc(total_size, GFP_KERNEL);
	if (!buffer) {
		mutex_unlock(&vgpu->lock);
		return -ENOMEM;
	}

	hdr = (struct nv_vgpu_migration_header *)buffer;
	hdr->magic = NV_VGPU_MIGRATION_MAGIC;
	hdr->version = NV_VGPU_MIGRATION_VERSION;
	hdr->flags = 0;
	hdr->num_sections = num_sections + 1;
	hdr->total_size = total_size;
	hdr->timestamp = ktime_get_ns();

	sections = (struct nv_vgpu_migration_section *)(buffer + sizeof(*hdr));

	dev_state = (struct nv_vgpu_migration_device_state *)(buffer + sizeof(*hdr) +
		(hdr->num_sections * sizeof(*sections)));

	dev_state->vgpu_type = vgpu->type;
	dev_state->vgpu_id = vgpu->vgpu_id;
	dev_state->gpu_instance_id = vgpu->gpu_instance_id;
	dev_state->compute_instance_id = vgpu->compute_instance_id;
	dev_state->state = vgpu->state;
	dev_state->fb_base = vgpu->fb_base;
	dev_state->fb_size = vgpu->fb_size;
	dev_state->bar0_config_base = vgpu->bar0_config_base;

	sections[0].type = NV_VGPU_MIG_SECTION_DEVICE_STATE;
	sections[0].flags = 0;
	sections[0].offset = (u8 *)dev_state - buffer;
	sections[0].size = sizeof(*dev_state);
	sections[0].checksum = crc32(0, dev_state, sizeof(*dev_state));

	u8 *data_ptr = (u8 *)dev_state + sizeof(*dev_state);
	for (i = 0; i < NV_VGPU_MAX_REGIONS; i++) {
		if (!vgpu->bars[i].mapped)
			continue;

		bar_state = (struct nv_vgpu_migration_bar_state *)data_ptr;
		bar_state->bar_index = i;
		bar_state->gpa = vgpu->bars[i].gpa;
		bar_state->size = vgpu->bars[i].size;
		bar_state->iova = vgpu->bars[i].iova;

		if (vgpu->bars[i].vaddr && vgpu->bars[i].size > 0) {
			memcpy(bar_state->data, vgpu->bars[i].vaddr, vgpu->bars[i].size);
		}

		sections[num_sections].type = NV_VGPU_MIG_SECTION_BAR_STATE;
		sections[num_sections].flags = 0;
		sections[num_sections].offset = data_ptr - buffer;
		sections[num_sections].size = sizeof(*bar_state) + vgpu->bars[i].size;
		sections[num_sections].checksum = crc32(0, bar_state, sections[num_sections].size);

		data_ptr += sections[num_sections].size;
		num_sections++;
	}

	hdr->num_sections = num_sections;
	hdr->total_size = data_ptr - buffer;

	*data = buffer;
	*size = hdr->total_size;

	mutex_unlock(&vgpu->lock);
	return 0;
}

int nv_vgpu_migration_restore_state(struct nv_vgpu_device *vgpu, u8 *data, size_t size)
{
	struct nv_vgpu_migration_header *hdr = (struct nv_vgpu_migration_header *)data;
	struct nv_vgpu_migration_section *sections;
	struct nv_vgpu_migration_device_state *dev_state;
	struct nv_vgpu_migration_bar_state *bar_state;
	int i;

	if (!data || size < sizeof(*hdr))
		return -EINVAL;

	if (hdr->magic != NV_VGPU_MIGRATION_MAGIC)
		return -EINVAL;

	if (hdr->version > NV_VGPU_MIGRATION_VERSION)
		return -EINVAL;

	sections = (struct nv_vgpu_migration_section *)(data + sizeof(*hdr));

	mutex_lock(&vgpu->lock);

	for (i = 0; i < hdr->num_sections; i++) {
		u32 checksum;
		u8 *section_data = data + sections[i].offset;

		checksum = crc32(0, section_data, sections[i].size);
		if (checksum != sections[i].checksum) {
			nv_vgpu_err(vgpu, "Migration checksum mismatch section %d", i);
			mutex_unlock(&vgpu->lock);
			return -EIO;
		}

		switch (sections[i].type) {
		case NV_VGPU_MIG_SECTION_DEVICE_STATE:
			dev_state = (struct nv_vgpu_migration_device_state *)section_data;
			vgpu->type = dev_state->vgpu_type;
			vgpu->vgpu_id = dev_state->vgpu_id;
			vgpu->gpu_instance_id = dev_state->gpu_instance_id;
			vgpu->compute_instance_id = dev_state->compute_instance_id;
			vgpu->state = dev_state->state;
			vgpu->fb_base = dev_state->fb_base;
			vgpu->fb_size = dev_state->fb_size;
			vgpu->bar0_config_base = dev_state->bar0_config_base;
			break;

		case NV_VGPU_MIG_SECTION_BAR_STATE:
			bar_state = (struct nv_vgpu_migration_bar_state *)section_data;
			if (bar_state->bar_index < NV_VGPU_MAX_REGIONS) {
				vgpu->bars[bar_state->bar_index].gpa = bar_state->gpa;
				vgpu->bars[bar_state->bar_index].size = bar_state->size;
				vgpu->bars[bar_state->bar_index].iova = bar_state->iova;
				if (vgpu->bars[bar_state->bar_index].vaddr && bar_state->size > 0) {
					memcpy(vgpu->bars[bar_state->bar_index].vaddr, bar_state->data, bar_state->size);
				}
			}
			break;

		default:
			nv_vgpu_warn(vgpu, "Unknown migration section type: %d", sections[i].type);
			break;
		}
	}

	mutex_unlock(&vgpu->lock);
	return 0;
}

int nv_vgpu_migration_prepare(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_migration_ctx *ctx = vgpu->migration.ctx;

	if (!ctx)
		return -ENODEV;

	mutex_lock(&ctx->lock);
	if (ctx->in_progress) {
		mutex_unlock(&ctx->lock);
		return -EBUSY;
	}
	ctx->in_progress = true;
	ctx->state.version = NV_VGPU_MIGRATION_VERSION;
	ctx->state.flags = 0;
	ctx->state.timestamp = ktime_get_ns();
	reinit_completion(&ctx->state.done);
	mutex_unlock(&ctx->lock);

	ctx->save_thread = kthread_run(nv_vgpu_migration_save_thread, vgpu, "nv-vgpu-mig-save");
	if (IS_ERR(ctx->save_thread)) {
		ctx->in_progress = false;
		return PTR_ERR(ctx->save_thread);
	}

	return 0;
}

int nv_vgpu_migration_complete(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_migration_ctx *ctx = vgpu->migration.ctx;
	int ret;

	if (!ctx)
		return -ENODEV;

	wait_for_completion(&ctx->state.done);
	ret = ctx->state.result;

	mutex_lock(&ctx->lock);
	ctx->in_progress = false;
	ctx->save_thread = NULL;
	mutex_unlock(&ctx->lock);

	return ret;
}

int nv_vgpu_migration_cancel(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_migration_ctx *ctx = vgpu->migration.ctx;

	if (!ctx)
		return -ENODEV;

	if (ctx->save_thread) {
		kthread_stop(ctx->save_thread);
		ctx->save_thread = NULL;
	}
	if (ctx->restore_thread) {
		kthread_stop(ctx->restore_thread);
		ctx->restore_thread = NULL;
	}

	mutex_lock(&ctx->lock);
	ctx->in_progress = false;
	mutex_unlock(&ctx->lock);

	return 0;
}

u8 *nv_vgpu_migration_get_data(struct nv_vgpu_device *vgpu, size_t *size)
{
	struct nv_vgpu_migration_ctx *ctx = vgpu->migration.ctx;

	if (!ctx || !ctx->state.data)
		return NULL;

	if (size)
		*size = ctx->state.data_size;

	return ctx->state.data;
}