// ivshmem_common.h
#ifndef IVSHMEM_COMMON_H
#define IVSHMEM_COMMON_H

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <error.h>
#include <unistd.h>
#include <fcntl.h>
#include <poll.h>
#include <sys/mman.h>
#include <sys/signalfd.h>

struct ivshm_regs {
    uint32_t id;
    uint32_t max_peers;
    uint32_t int_control;
    uint32_t doorbell;
    uint32_t state;
};

typedef struct {
    int uio_fd;
    int uio_num;
    uint32_t id;
    volatile struct ivshm_regs *regs;
    volatile uint32_t *state;
    volatile uint32_t *in;
    volatile uint32_t *out;
    volatile uint32_t *rw;
    uint32_t peer_count;
} ivshmem_ctx;

static inline uint32_t mmio_read32(void *addr) {
    return *(volatile uint32_t *)addr;
}

static inline void mmio_write32(void *addr, uint32_t val) {
    *(volatile uint32_t *)addr = val;
}

void *map(int fd, int uio_num, int index) {
    char path[128];
    snprintf(path, sizeof(path), "/sys/class/uio/uio%d/maps/map%d/size", uio_num, index);
    FILE *f = fopen(path, "r");
    if (!f) error(1, errno, "Open %s failed", path);
    unsigned long size;
    fscanf(f, "%lx", &size);
    fclose(f);
    return mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, index * getpagesize());
}

void ivshmem_init(ivshmem_ctx *ctx, const char *dev) {
    ctx->uio_fd = open(dev, O_RDWR);
    if (ctx->uio_fd < 0) error(1, errno, "open %s", dev);
    
    // extract uio number from dev path like "/dev/uio0"
    if (sscanf(dev, "/dev/uio%d", &ctx->uio_num) != 1) {
        error(1, 0, "Invalid UIO device path: %s", dev);
    }

    ctx->regs = map(ctx->uio_fd, ctx->uio_num, 0);
    ctx->rw   = map(ctx->uio_fd, ctx->uio_num, 1);
    void *io  = map(ctx->uio_fd, ctx->uio_num, 2);

    ctx->id = mmio_read32((void *)&ctx->regs->id);
    ctx->peer_count = mmio_read32((void *)&ctx->regs->max_peers);

    ctx->state = (volatile uint32_t *)ctx->rw;
    ctx->in    = (volatile uint32_t *)io;
    ctx->out   = (volatile uint32_t *)((char *)io + 4096);  // assuming 4KB per peer

    ctx->state[ctx->id] = 1;
    mmio_write32((void *)&ctx->regs->int_control, 1 << ctx->id);
}

void ivshmem_wait_irq(ivshmem_ctx *ctx) {
    struct pollfd pfd = {
        .fd = ctx->uio_fd,
        .events = POLLIN
    };
    uint32_t irq_count;
    poll(&pfd, 1, -1);
    read(ctx->uio_fd, &irq_count, sizeof(irq_count));
}

void ivshmem_cleanup(ivshmem_ctx *ctx) {
    ctx->state[ctx->id] = 0;
    close(ctx->uio_fd);
}

#endif