#ifndef _NV_VGPU_MIGRATION_H_
#define _NV_VGPU_MIGRATION_H_

#include "nvidia-vgpu-vfio.h"

#define NV_VGPU_MIGRATION_MAGIC 0x4E564750
#define NV_VGPU_MIGRATION_VERSION 1

struct nv_vgpu_migration_ctx {
	struct nv_vgpu_device *vgpu;
	struct mutex lock;
	struct nv_vgpu_migration_state state;
	struct task_struct *save_thread;
	struct task_struct *restore_thread;
	bool in_progress;
};

struct nv_vgpu_migration_header {
	u32 magic;
	u32 version;
	u32 flags;
	u32 num_sections;
	u64 total_size;
	u64 timestamp;
};

enum nv_vgpu_migration_section_type {
	NV_VGPU_MIG_SECTION_DEVICE_STATE = 0x01,
	NV_VGPU_MIG_SECTION_BAR_STATE = 0x02,
	NV_VGPU_MIG_SECTION_MEMORY_STATE = 0x03,
	NV_VGPU_MIG_SECTION_IRQ_STATE = 0x04,
	NV_VGPU_MIG_SECTION_GSP_STATE = 0x05,
	NV_VGPU_MIG_SECTION_ENGINE_STATE = 0x06,
	NV_VGPU_MIG_SECTION_CHANNEL_STATE = 0x07,
	NV_VGPU_MIG_SECTION_CONTEXT_STATE = 0x08,
};

struct nv_vgpu_migration_section {
	u32 type;
	u32 flags;
	u64 offset;
	u64 size;
	u32 checksum;
};

struct nv_vgpu_migration_device_state {
	u32 vgpu_type;
	u32 vgpu_id;
	u32 gpu_instance_id;
	u32 compute_instance_id;
	u32 state;
	u64 fb_base;
	u64 fb_size;
	u64 bar0_config_base;
};

struct nv_vgpu_migration_bar_state {
	u32 bar_index;
	u64 gpa;
	u64 size;
	u64 iova;
	u8 data[];
};

int nv_vgpu_migration_init(struct nv_vgpu_device *vgpu);
void nv_vgpu_migration_cleanup(struct nv_vgpu_device *vgpu);

int nv_vgpu_migration_save_state(struct nv_vgpu_device *vgpu, u8 **data, size_t *size);
int nv_vgpu_migration_restore_state(struct nv_vgpu_device *vgpu, u8 *data, size_t size);

int nv_vgpu_migration_prepare(struct nv_vgpu_device *vgpu);
int nv_vgpu_migration_complete(struct nv_vgpu_device *vgpu);
int nv_vgpu_migration_cancel(struct nv_vgpu_device *vgpu);

#endif