#ifndef _NV_VGPU_GSP_COMM_H_
#define _NV_VGPU_GSP_COMM_H_

#include "nvidia-vgpu-vfio.h"

#define NV_VGPU_GSP_RPC_MAX_SIZE 4096
#define NV_VGPU_GSP_RPC_TIMEOUT_MS 5000
#define NV_VGPU_GSP_QUEUE_DEPTH 256

#define NV_VGPU_GSP_DOORBELL_OFFSET 0x100000
#define NV_VGPU_GSP_SHARED_MEM_OFFSET 0x200000
#define NV_VGPU_GSP_SHARED_MEM_SIZE (64 * 1024)

#define NV_VGPU_GSP_FIRMWARE_PATH "/lib/firmware/nvidia/tu102/gsp/"

enum nv_vgpu_gsp_rpc_type {
	NV_GSP_RPC_TYPE_ALLOC = 0x01,
	NV_GSP_RPC_TYPE_FREE = 0x02,
	NV_GSP_RPC_TYPE_MAP = 0x03,
	NV_GSP_RPC_TYPE_UNMAP = 0x04,
	NV_GSP_RPC_TYPE_PIN = 0x05,
	NV_GSP_RPC_TYPE_UNPIN = 0x06,
	NV_GSP_RPC_TYPE_SCHEDULE = 0x07,
	NV_GSP_RPC_TYPE_YIELD = 0x08,
	NV_GSP_RPC_TYPE_NOTIFY = 0x09,
	NV_GSP_RPC_TYPE_BAR_TRAP = 0x0A,
	NV_GSP_RPC_TYPE_MIGRATE = 0x0B,
	NV_GSP_RPC_TYPE_DEBUG = 0xFF,
};

struct nv_vgpu_gsp_rpc_header {
	u32 type;
	u32 sequence;
	u32 size;
	u32 flags;
	u64 timestamp;
};

struct nv_vgpu_gsp_rpc_msg {
	struct nv_vgpu_gsp_rpc_header hdr;
	u8 data[];
};

struct nv_vgpu_gsp_firmware {
	void *bootloader;
	size_t bootloader_size;
	void *booter_load;
	size_t booter_load_size;
	void *gsp;
	size_t gsp_size;
};

struct nv_vgpu_gsp_comm_ctx {
	struct nv_vgpu_device *vgpu;
	
	void __iomem *gsp_doorbell;
	void __iomem *gsp_shared_mem;
	size_t shared_mem_size;
	
	struct mutex queue_lock;
	struct list_head pending_queue;
	struct list_head complete_queue;
	atomic_t sequence_counter;
	
	wait_queue_head_t wait_queue;
	
	struct workqueue_struct *rpc_wq;
	struct work_struct rpc_work;
	
	struct delayed_work poll_work;
	bool poll_active;
	
	spinlock_t irq_lock;
	bool irq_enabled;
	
	struct nv_vgpu_gsp_firmware firmware;
	bool firmware_loaded;
};

struct nv_vgpu_gsp_rpc_entry {
	struct list_head list;
	struct nv_vgpu_gsp_rpc_msg *msg;
	struct completion done;
	int result;
};

int nv_vgpu_gsp_init(struct nv_vgpu_device *vgpu);
void nv_vgpu_gsp_cleanup(struct nv_vgpu_device *vgpu);

int nv_vgpu_gsp_send_rpc(struct nv_vgpu_device *vgpu, void *msg, size_t size);
int nv_vgpu_gsp_recv_rpc(struct nv_vgpu_device *vgpu, void *msg, size_t size, int timeout_ms);

int nv_vgpu_gsp_alloc(struct nv_vgpu_device *vgpu, u64 size, u32 align, u64 *handle);
int nv_vgpu_gsp_free(struct nv_vgpu_device *vgpu, u64 handle);
int nv_vgpu_gsp_map(struct nv_vgpu_device *vgpu, u64 handle, u64 iova, u64 size, u32 flags);
int nv_vgpu_gsp_unmap(struct nv_vgpu_device *vgpu, u64 handle, u64 iova, u64 size);
int nv_vgpu_gsp_schedule(struct nv_vgpu_device *vgpu, u32 runlist_id, u32 channel_id);
int nv_vgpu_gsp_notify_bar_trap(struct nv_vgpu_device *vgpu, u64 offset, u32 size, void *val);

int nv_vgpu_gsp_load_firmware(struct nv_vgpu_device *vgpu);
void nv_vgpu_gsp_release_firmware(struct nv_vgpu_device *vgpu);
int nv_vgpu_gsp_boot(struct nv_vgpu_device *vgpu);

/* ============================================================
 * New Firmware-Decompiled RPC Operations (enhanced APIs)
 * ============================================================ */

int nv_vgpu_gsp_rpc_alloc(struct nv_vgpu_device *vgpu,
                           u64 size, u64 alignment, u32 mem_type,
                           u32 flags, u64 addr_hint, u32 pid,
                           u64 *gpu_va, u64 *handle, u64 *actual_size);

int nv_vgpu_gsp_rpc_free(struct nv_vgpu_device *vgpu, u64 handle);

int nv_vgpu_gsp_rpc_map(struct nv_vgpu_device *vgpu,
                         u64 handle, u64 gpu_va, u64 size,
                         u32 prot, u32 flags, u64 offset,
                         u64 *mapped_va);

int nv_vgpu_gsp_rpc_unmap(struct nv_vgpu_device *vgpu,
                           u64 handle, u64 gpu_va, u64 size);

int nv_vgpu_gsp_rpc_pin(struct nv_vgpu_device *vgpu,
                         u64 handle, u64 offset, u64 size,
                         u32 flags, u64 *dma_addr, u32 *num_pages);

int nv_vgpu_gsp_rpc_unpin(struct nv_vgpu_device *vgpu,
                           u64 handle, u64 offset, u64 size);

int nv_vgpu_gsp_rpc_schedule(struct nv_vgpu_device *vgpu,
                              u64 work_handle, u64 callback, u64 arg,
                              u32 priority, u32 flags);

int nv_vgpu_gsp_rpc_yield(struct nv_vgpu_device *vgpu);

int nv_vgpu_gsp_rpc_bar_trap(struct nv_vgpu_device *vgpu,
                              u64 bar_addr, u32 access_type, u32 size,
                              u64 value, u32 bar_num, u64 *ret_value);

int nv_vgpu_gsp_rpc_migrate_save(struct nv_vgpu_device *vgpu,
                                  u64 src_gpu_id, u64 dst_gpu_id,
                                  u64 memory_bitmap, u64 page_table_ptr,
                                  u64 *checkpoint);

int nv_vgpu_gsp_rpc_migrate_restore(struct nv_vgpu_device *vgpu,
                                     u64 src_gpu_id, u64 dst_gpu_id,
                                     u64 memory_bitmap, u64 page_table_ptr,
                                     u64 checkpoint);

static inline u32 nv_vgpu_gsp_next_sequence(struct nv_vgpu_gsp_comm_ctx *ctx)
{
	return atomic_inc_return(&ctx->sequence_counter);
}

#endif