#include "redox_glue.h"

#include <fcntl.h>
#include <errno.h>
#include <stdatomic.h>
#include <poll.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

unsigned long jiffies;

struct redox_mapped_region {
    void *addr;
    size_t size;
    int fd;
    struct redox_mapped_region *next;
};

struct redox_irq_entry {
    unsigned int irq;
    irq_handler_t handler;
    void *dev_id;
    int fd;
    pthread_t thread;
    bool active;
    struct redox_irq_entry *next;
};

struct redox_pci_region {
    struct pci_dev *pdev;
    unsigned int bar;
    void *addr;
    size_t size;
    bool claimed;
    struct redox_pci_region *next;
};

static pthread_mutex_t g_region_lock = PTHREAD_MUTEX_INITIALIZER;
static struct redox_mapped_region *g_regions;
static pthread_mutex_t g_irq_lock = PTHREAD_MUTEX_INITIALIZER;
static struct redox_irq_entry *g_irqs;
static pthread_mutex_t g_pci_region_lock = PTHREAD_MUTEX_INITIALIZER;
static struct redox_pci_region *g_pci_regions;

static void redox_jiffies_advance(unsigned long delta)
{
    __sync_add_and_fetch(&jiffies, delta);
}

struct redox_vmalloc_region {
    void *base;
    size_t total_size;
    struct redox_vmalloc_region *next;
};

static pthread_mutex_t g_vmalloc_lock = PTHREAD_MUTEX_INITIALIZER;
static struct redox_vmalloc_region *g_vmalloc_regions;

void *vmalloc(unsigned long size)
{
    size_t aligned = PAGE_ALIGN((size_t)size);
    size_t total = aligned + 2 * PAGE_SIZE;
    void *base = mmap(NULL, total, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);

    if (base == MAP_FAILED) {
        return NULL;
    }

    void *usable = (char *)base + PAGE_SIZE;
    if (mprotect(usable, aligned, PROT_READ | PROT_WRITE) != 0) {
        munmap(base, total);
        return NULL;
    }

    struct redox_vmalloc_region *reg = malloc(sizeof(*reg));
    if (!reg) {
        munmap(base, total);
        return NULL;
    }
    reg->base = base;
    reg->total_size = total;

    pthread_mutex_lock(&g_vmalloc_lock);
    reg->next = g_vmalloc_regions;
    g_vmalloc_regions = reg;
    pthread_mutex_unlock(&g_vmalloc_lock);

    return usable;
}

void vfree(const void *addr)
{
    if (!addr) {
        return;
    }

    void *usable_base = (char *)addr - 0;
    void *region_base = (char *)addr - PAGE_SIZE;

    struct redox_vmalloc_region **link;
    pthread_mutex_lock(&g_vmalloc_lock);
    link = &g_vmalloc_regions;
    while (*link) {
        if ((*link)->base == region_base) {
            struct redox_vmalloc_region *reg = *link;
            *link = reg->next;
            pthread_mutex_unlock(&g_vmalloc_lock);
            munmap(reg->base, reg->total_size);
            free(reg);
            return;
        }
        link = &(*link)->next;
    }
    pthread_mutex_unlock(&g_vmalloc_lock);

    free((void *)addr);
}

static void redox_track_region(void *addr, size_t size, int fd)
{
    struct redox_mapped_region *region = malloc(sizeof(*region));
    if (!region) {
        if (fd >= 0) {
            close(fd);
        }
        return;
    }

    region->addr = addr;
    region->size = size;
    region->fd = fd;

    pthread_mutex_lock(&g_region_lock);
    region->next = g_regions;
    g_regions = region;
    pthread_mutex_unlock(&g_region_lock);
}

static struct redox_mapped_region *redox_untrack_region(const void *addr)
{
    struct redox_mapped_region *prev = NULL;
    struct redox_mapped_region *cur;

    pthread_mutex_lock(&g_region_lock);
    cur = g_regions;
    while (cur) {
        if (cur->addr == addr) {
            if (prev) {
                prev->next = cur->next;
            } else {
                g_regions = cur->next;
            }
            pthread_mutex_unlock(&g_region_lock);
            return cur;
        }
        prev = cur;
        cur = cur->next;
    }
    pthread_mutex_unlock(&g_region_lock);
    return NULL;
}

void __iomem *redox_ioremap(phys_addr_t offset, size_t size)
{
    int fd = open("/scheme/memory/physical", O_RDWR);
    void *addr;

    if (fd >= 0) {
        addr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, (off_t)offset);
        if (addr != MAP_FAILED) {
            redox_track_region(addr, size, fd);
            return addr;
        }
        close(fd);
    }

    pr_err("ioremap failed for %#llx (%zu bytes): %s\n",
           (unsigned long long)offset, size, strerror(errno));
    return NULL;
}

void redox_iounmap(void __iomem *addr)
{
    struct redox_mapped_region *region;

    if (!addr) {
        return;
    }

    region = redox_untrack_region(addr);
    if (!region) {
        return;
    }

    munmap(region->addr, region->size);
    if (region->fd >= 0) {
        close(region->fd);
    }
    free(region);
}

void redox_iowrite32(u32 val, void __iomem *addr)
{
    *(volatile u32 *)addr = val;
}

u32 redox_ioread32(const void __iomem *addr)
{
    return *(volatile const u32 *)addr;
}

void redox_iowrite16(u16 val, void __iomem *addr)
{
    *(volatile u16 *)addr = val;
}

u16 redox_ioread16(const void __iomem *addr)
{
    return *(volatile const u16 *)addr;
}

void redox_iowrite8(u8 val, void __iomem *addr)
{
    *(volatile u8 *)addr = val;
}

u8 redox_ioread8(const void __iomem *addr)
{
    return *(volatile const u8 *)addr;
}

void redox_mmio_write32(void *base, u32 offset, u32 val)
{
    if (!base) {
        return;
    }
    *(volatile u32 *)((u8 *)base + offset) = val;
}

u32 redox_mmio_read32(void *base, u32 offset)
{
    if (!base) {
        return 0;
    }
    return *(volatile u32 *)((u8 *)base + offset);
}

void *redox_dma_alloc_coherent(size_t size, dma_addr_t *dma_handle)
{
    size_t aligned_size = PAGE_ALIGN(size);
    int mem_fd, region_fd;
    void *ptr;
    void *phys_buf;
    uint64_t virt_addr;

    mem_fd = open("/scheme/memory/scheme-root", O_RDWR);
    if (mem_fd < 0) {
        pr_err("dma_alloc_coherent: cannot open scheme:memory — falling back to heap\n");
        ptr = NULL;
        if (posix_memalign(&ptr, PAGE_SIZE, aligned_size) != 0) {
            return NULL;
        }
        memset(ptr, 0, aligned_size);
        if (dma_handle) {
            *dma_handle = (dma_addr_t)(uintptr_t)ptr;
        }
        return ptr;
    }

    region_fd = openat(mem_fd, "zeroed@wb?phys_contiguous", O_RDWR, 0);
    close(mem_fd);
    if (region_fd < 0) {
        pr_err("dma_alloc_coherent: cannot open zeroed@wb phys_contiguous region\n");
        ptr = NULL;
        if (posix_memalign(&ptr, PAGE_SIZE, aligned_size) != 0) {
            return NULL;
        }
        memset(ptr, 0, aligned_size);
        if (dma_handle) {
            *dma_handle = (dma_addr_t)(uintptr_t)ptr;
        }
        return ptr;
    }

    ptr = mmap(NULL, aligned_size, PROT_READ | PROT_WRITE, MAP_SHARED, region_fd, 0);
    if (ptr == MAP_FAILED) {
        pr_err("dma_alloc_coherent: mmap failed: %s\n", strerror(errno));
        close(region_fd);
        return NULL;
    }

    virt_addr = (uint64_t)(uintptr_t)ptr;

    int trans_fd = open("/scheme/memory/translation", O_RDWR);
    if (trans_fd >= 0) {
        uint8_t buf[8];
        memcpy(buf, &virt_addr, 8);
        ssize_t n = write(trans_fd, buf, 8);
        if (n == 8) {
            n = read(trans_fd, buf, 8);
            if (n == 8) {
                uint64_t phys;
                memcpy(&phys, buf, 8);
                if (dma_handle) {
                    *dma_handle = (dma_addr_t)phys;
                }
                close(trans_fd);
                redox_track_region(ptr, aligned_size, region_fd);
                return ptr;
            }
        }
        close(trans_fd);
        pr_err("dma_alloc_coherent: translation failed — using virtual address as DMA handle\n");
    }

    if (dma_handle) {
        *dma_handle = (dma_addr_t)virt_addr;
    }
    redox_track_region(ptr, aligned_size, region_fd);
    return ptr;
}

void redox_dma_free_coherent(size_t size, void *vaddr, dma_addr_t dma_handle)
{
    (void)dma_handle;

    struct redox_mapped_region *region = redox_untrack_region(vaddr);
    if (region) {
        munmap(region->addr, region->size);
        if (region->fd >= 0) {
            close(region->fd);
        }
        free(region);
    } else {
        free(vaddr);
    }
}

/*
 * PCI device state — populated by the Rust side via redox_pci_set_device_info()
 * before amdgpu_redox_init() is called.  redox_pci_find_amd_gpu() returns a
 * pointer to this struct, or NULL if the device info has not been set yet.
 */
static struct pci_dev g_pci_dev;
static int g_pci_dev_populated;

#define REDOX_MAX_FIRMWARE_BYTES (64U * 1024U * 1024U)

void redox_pci_set_device_info(u16 vendor, u16 device,
                                u8 bus_number, u8 dev_number,
                                u8 func_number, u8 revision, u32 irq,
                                u64 bar0_addr, u64 bar0_size,
                                u64 bar2_addr, u64 bar2_size)
{
    memset(&g_pci_dev, 0, sizeof(g_pci_dev));
    g_pci_dev.vendor = vendor;
    g_pci_dev.device_id = device;
    g_pci_dev.bus_number = bus_number;
    g_pci_dev.dev_number = dev_number;
    g_pci_dev.func_number = func_number;
    g_pci_dev.revision = revision;
    g_pci_dev.irq = irq;
    g_pci_dev.resource_start[0] = (phys_addr_t)bar0_addr;
    g_pci_dev.resource_len[0] = bar0_size;
    g_pci_dev.resource_flags[0] = IORESOURCE_MEM;
    g_pci_dev.resource_start[2] = (phys_addr_t)bar2_addr;
    g_pci_dev.resource_len[2] = bar2_size;
    g_pci_dev.resource_flags[2] = IORESOURCE_MEM;
    g_pci_dev.driver_data = NULL;
    memset(&g_pci_dev.device_obj, 0, sizeof(g_pci_dev.device_obj));
    g_pci_dev.enabled = false;
    g_pci_dev.refcount = 1;
    g_pci_dev.mmio_base = NULL;
    g_pci_dev.is_amdgpu = 1;
    g_pci_dev_populated = 1;

    printk("PCI device info set: %02x:%02x.%u vendor=%#06x device=%#06x rev=%#04x irq=%u "
           "bar0=%#llx+%#llx bar2=%#llx+%#llx\n",
           bus_number, dev_number, func_number,
           vendor, device, revision, irq,
           (unsigned long long)bar0_addr, (unsigned long long)bar0_size,
           (unsigned long long)bar2_addr, (unsigned long long)bar2_size);
}

struct pci_dev *redox_pci_find_amd_gpu(void)
{
    if (!g_pci_dev_populated) {
        pr_err("redox_pci_find_amd_gpu: device info not set — "
               "call redox_pci_set_device_info() first\n");
        return NULL;
    }

    return &g_pci_dev;
}

void redox_pci_dev_put(struct pci_dev *pdev)
{
    if (!pdev) {
        return;
    }

    if (__sync_sub_and_fetch(&pdev->refcount, 1) == 0) {
        redox_pci_release_regions(pdev);
        if (pdev != &g_pci_dev) {
            free(pdev);
        }
    }
}

int redox_pci_enable_device(struct pci_dev *pdev)
{
    if (!pdev) {
        return -ENODEV;
    }

    if (pdev->enabled) {
        return 0;
    }

    pdev->command |= (u16)(PCI_COMMAND_MEMORY | PCI_COMMAND_IO | PCI_COMMAND_MASTER);
    pdev->enabled = true;

    dev_info(&pdev->device_obj, "PCI device enabled (memory + I/O + bus master); "
             "hardware command register was pre-configured by pcid-spawner\n");
    return 0;
}

void redox_pci_set_master(struct pci_dev *pdev)
{
    if (!pdev) {
        return;
    }

    pdev->command |= (u16)PCI_COMMAND_MASTER;
    dev_info(&pdev->device_obj, "PCI bus master enabled\n");
}

int redox_pci_request_regions(struct pci_dev *pdev, const char *name)
{
    unsigned int bar;
    (void)name;

    if (!pdev) {
        return -ENODEV;
    }

    pthread_mutex_lock(&g_pci_region_lock);
    for (bar = 0; bar < 6; ++bar) {
        struct redox_pci_region *region;
        if (!(pdev->resource_flags[bar] & IORESOURCE_MEM) || pdev->resource_len[bar] == 0) {
            continue;
        }
        for (region = g_pci_regions; region != NULL; region = region->next) {
            if (region->claimed && region->pdev != pdev) {
                u64 a0 = pdev->resource_start[bar];
                u64 a1 = a0 + pdev->resource_len[bar] - 1;
                u64 b0 = (u64)(uintptr_t)region->addr;
                u64 b1 = b0 + region->size - 1;
                if (!(a1 < b0 || b1 < a0)) {
                    pthread_mutex_unlock(&g_pci_region_lock);
                    return -EBUSY;
                }
            }
        }
        region = calloc(1, sizeof(*region));
        if (!region) {
            pthread_mutex_unlock(&g_pci_region_lock);
            return -ENOMEM;
        }
        region->pdev = pdev;
        region->bar = bar;
        region->size = (size_t)pdev->resource_len[bar];
        region->addr = redox_ioremap((phys_addr_t)pdev->resource_start[bar], region->size);
        if (!region->addr) {
            free(region);
            pthread_mutex_unlock(&g_pci_region_lock);
            return -ENOMEM;
        }
        region->claimed = true;
        region->next = g_pci_regions;
        g_pci_regions = region;
    }
    pthread_mutex_unlock(&g_pci_region_lock);
    return 0;
}

void redox_pci_release_regions(struct pci_dev *pdev)
{
    struct redox_pci_region **link;
    if (!pdev) {
        return;
    }

    pthread_mutex_lock(&g_pci_region_lock);
    link = &g_pci_regions;
    while (*link) {
        struct redox_pci_region *region = *link;
        if (region->pdev == pdev) {
            *link = region->next;
            pthread_mutex_unlock(&g_pci_region_lock);
            if (region->addr) {
                redox_iounmap((void __iomem *)region->addr);
            }
            free(region);
            pthread_mutex_lock(&g_pci_region_lock);
            link = &g_pci_regions;
            continue;
        }
        link = &region->next;
    }
    pthread_mutex_unlock(&g_pci_region_lock);
}

int redox_request_firmware(const struct firmware **fw, const char *name, void *dev)
{
    char path[512];
    int fd;
    struct stat st;
    struct firmware *image;
    u8 *data;
    ssize_t nread;

    (void)dev;
    if (!fw || !name) {
        return -EINVAL;
    }

    snprintf(path, sizeof(path), "/scheme/firmware/amdgpu/%s", name);
    fd = open(path, O_RDONLY);
    if (fd < 0) {
        return -ENOENT;
    }

    if (fstat(fd, &st) != 0 || st.st_size < 0) {
        close(fd);
        return -EIO;
    }

    if ((unsigned long long)st.st_size > REDOX_MAX_FIRMWARE_BYTES) {
        close(fd);
        return -EFBIG;
    }

    image = calloc(1, sizeof(*image));
    data = malloc((size_t)st.st_size);
    if (!image || !data) {
        free(image);
        free(data);
        close(fd);
        return -ENOMEM;
    }

    nread = read(fd, data, (size_t)st.st_size);
    close(fd);
    if (nread != st.st_size) {
        free(image);
        free(data);
        return -EIO;
    }

    image->size = (size_t)st.st_size;
    image->data = data;
    *fw = image;
    return 0;
}

void redox_release_firmware(const struct firmware *fw)
{
    struct firmware *owned = (struct firmware *)fw;

    if (!owned) {
        return;
    }

    free((void *)owned->data);
    free(owned);
}

static void *redox_irq_thread_main(void *arg)
{
    struct redox_irq_entry *entry = arg;
    unsigned char buf[64];

    pthread_setcancelstate(PTHREAD_CANCEL_ENABLE, NULL);
    pthread_setcanceltype(PTHREAD_CANCEL_DEFERRED, NULL);

    while (entry->active) {
        ssize_t n = read(entry->fd, buf, sizeof(buf));
        if (n < 0) {
            if (errno == EINTR) {
                continue;
            }
            break;
        }
        if (n == 0) {
            break;
        }
        entry->handler((int)entry->irq, entry->dev_id);
    }

    return NULL;
}

int redox_request_irq(unsigned int irq, irq_handler_t handler, unsigned long flags, const char *name, void *dev)
{
    char path[128];
    int fd;
    struct redox_irq_entry *entry;

    (void)flags;
    (void)name;
    if (!handler) {
        return -EINVAL;
    }

    snprintf(path, sizeof(path), "/scheme/irq/%u", irq);
    fd = open(path, O_RDWR);
    if (fd < 0) {
        return -ENOENT;
    }

    entry = calloc(1, sizeof(*entry));
    if (!entry) {
        close(fd);
        return -ENOMEM;
    }

    entry->irq = irq;
    entry->handler = handler;
    entry->dev_id = dev;
    entry->fd = fd;
    entry->active = true;

    pthread_mutex_lock(&g_irq_lock);
    entry->next = g_irqs;
    g_irqs = entry;
    pthread_mutex_unlock(&g_irq_lock);

    if (pthread_create(&entry->thread, NULL, redox_irq_thread_main, entry) != 0) {
        pthread_mutex_lock(&g_irq_lock);
        if (g_irqs == entry) {
            g_irqs = entry->next;
        } else {
            struct redox_irq_entry *cur = g_irqs;
            while (cur && cur->next != entry) {
                cur = cur->next;
            }
            if (cur) {
                cur->next = entry->next;
            }
        }
        pthread_mutex_unlock(&g_irq_lock);
        close(fd);
        free(entry);
        return -EFAULT;
    }

    return 0;
}

void redox_free_irq(unsigned int irq, void *dev_id)
{
    struct redox_irq_entry **link;
    struct redox_irq_entry *entry = NULL;

    pthread_mutex_lock(&g_irq_lock);
    link = &g_irqs;
    while (*link) {
        if ((*link)->irq == irq && (*link)->dev_id == dev_id) {
            entry = *link;
            *link = entry->next;
            entry->active = false;
            break;
        }
        link = &(*link)->next;
    }
    pthread_mutex_unlock(&g_irq_lock);

    if (!entry) {
        return;
    }

    pthread_cancel(entry->thread);
    pthread_join(entry->thread, NULL);
    close(entry->fd);
    free(entry);
}

void msleep(unsigned int msecs)
{
    struct timespec ts;

    ts.tv_sec = msecs / 1000U;
    ts.tv_nsec = (long)(msecs % 1000U) * 1000000L;
    nanosleep(&ts, NULL);
    redox_jiffies_advance(msecs_to_jiffies(msecs));
}

void udelay(unsigned long usecs)
{
    struct timespec ts;

    ts.tv_sec = usecs / 1000000UL;
    ts.tv_nsec = (long)(usecs % 1000000UL) * 1000L;
    nanosleep(&ts, NULL);
    redox_jiffies_advance(usecs_to_jiffies((unsigned int)usecs));
}

void mdelay(unsigned long msecs)
{
    msleep((unsigned int)msecs);
}

unsigned long msecs_to_jiffies(unsigned int msecs)
{
    return (unsigned long)msecs * HZ / 1000UL;
}

unsigned long usecs_to_jiffies(unsigned int usecs)
{
    return DIV_ROUND_UP_ULL((unsigned long long)usecs * HZ, 1000000ULL);
}

struct redox_pm_state {
    struct device *dev;
    atomic_int usage_count;
    bool enabled;
    bool allowed;
    bool active;
    bool ignore_children;
    bool no_pm;
    struct redox_pm_state *next;
};

static pthread_mutex_t g_pm_lock = PTHREAD_MUTEX_INITIALIZER;
static struct redox_pm_state *g_pm_states;

static struct redox_pm_state *redox_pm_find_state(struct device *dev)
{
    struct redox_pm_state *state;

    for (state = g_pm_states; state != NULL; state = state->next) {
        if (state->dev == dev) {
            return state;
        }
    }
    return NULL;
}

static struct redox_pm_state *redox_pm_get_state(struct device *dev)
{
    struct redox_pm_state *state = redox_pm_find_state(dev);
    struct pci_dev *pdev;

    if (state != NULL || dev == NULL) {
        return state;
    }

    state = kzalloc(sizeof(*state), 0);
    if (state == NULL) {
        return NULL;
    }

    state->dev = dev;
    atomic_init(&state->usage_count, 0);
    pdev = dev->pci_dev;
    if (pdev != NULL) {
        state->no_pm = pci_has_quirk(pdev, PCI_QUIRK_NO_PM);
    }

    pthread_mutex_lock(&g_pm_lock);
    state->next = g_pm_states;
    g_pm_states = state;
    pthread_mutex_unlock(&g_pm_lock);

    return state;
}

static bool redox_pm_blocked(struct redox_pm_state *state)
{
    return state == NULL || state->no_pm;
}

int pm_runtime_get_sync(struct device *dev)
{
    struct redox_pm_state *state = redox_pm_get_state(dev);

    if (redox_pm_blocked(state)) {
        return 1;
    }

    pthread_mutex_lock(&g_pm_lock);
    atomic_fetch_add_explicit(&state->usage_count, 1, memory_order_relaxed);
    state->active = true;
    int usage = atomic_load_explicit(&state->usage_count, memory_order_relaxed);
    pthread_mutex_unlock(&g_pm_lock);

    dev_info(dev, "runtime PM get_sync: usage=%d\n", usage);
    return usage;
}

int pm_runtime_get_noresume(struct device *dev)
{
    struct redox_pm_state *state = redox_pm_get_state(dev);

    if (redox_pm_blocked(state)) {
        return 1;
    }

    pthread_mutex_lock(&g_pm_lock);
    atomic_fetch_add_explicit(&state->usage_count, 1, memory_order_relaxed);
    int usage = atomic_load_explicit(&state->usage_count, memory_order_relaxed);
    pthread_mutex_unlock(&g_pm_lock);

    dev_info(dev, "runtime PM get_noresume: usage=%d\n", usage);
    return usage;
}

int pm_runtime_put_autosuspend(struct device *dev)
{
    struct redox_pm_state *state = redox_pm_get_state(dev);

    if (redox_pm_blocked(state)) {
        return 0;
    }

    pthread_mutex_lock(&g_pm_lock);
    if (atomic_load_explicit(&state->usage_count, memory_order_relaxed) > 0) {
        atomic_fetch_sub_explicit(&state->usage_count, 1, memory_order_relaxed);
    }
    int usage = atomic_load_explicit(&state->usage_count, memory_order_relaxed);
    if (usage == 0 && state->allowed && state->enabled) {
        state->active = false;
    }
    pthread_mutex_unlock(&g_pm_lock);

    dev_info(dev, "runtime PM put_autosuspend: usage=%d active=%d\n",
             usage, state->active ? 1 : 0);
    return usage;
}

int pm_runtime_put_noidle(struct device *dev)
{
    struct redox_pm_state *state = redox_pm_get_state(dev);

    if (redox_pm_blocked(state)) {
        return 0;
    }

    pthread_mutex_lock(&g_pm_lock);
    if (atomic_load_explicit(&state->usage_count, memory_order_relaxed) > 0) {
        atomic_fetch_sub_explicit(&state->usage_count, 1, memory_order_relaxed);
    }
    int usage = atomic_load_explicit(&state->usage_count, memory_order_relaxed);
    pthread_mutex_unlock(&g_pm_lock);

    dev_info(dev, "runtime PM put_noidle: usage=%d\n", usage);
    return usage;
}

int pm_runtime_idle(struct device *dev)
{
    struct redox_pm_state *state = redox_pm_get_state(dev);

    if (redox_pm_blocked(state)) {
        return 0;
    }

    pthread_mutex_lock(&g_pm_lock);
    if (atomic_load_explicit(&state->usage_count, memory_order_relaxed) == 0 && state->allowed && state->enabled) {
        state->active = false;
    }
    int usage = atomic_load_explicit(&state->usage_count, memory_order_relaxed);
    pthread_mutex_unlock(&g_pm_lock);

    dev_info(dev, "runtime PM idle: active=%d\n", state->active ? 1 : 0);
    return usage;
}

int pm_runtime_set_active(struct device *dev)
{
    struct redox_pm_state *state = redox_pm_get_state(dev);

    if (redox_pm_blocked(state)) {
        return 0;
    }

    pthread_mutex_lock(&g_pm_lock);
    state->active = true;
    if (atomic_load_explicit(&state->usage_count, memory_order_relaxed) < 1) {
        atomic_store_explicit(&state->usage_count, 1, memory_order_relaxed);
    }
    int usage = atomic_load_explicit(&state->usage_count, memory_order_relaxed);
    pthread_mutex_unlock(&g_pm_lock);

    dev_info(dev, "runtime PM set_active\n");
    return usage;
}

int pm_runtime_enable(struct device *dev)
{
    struct redox_pm_state *state = redox_pm_get_state(dev);

    if (redox_pm_blocked(state)) {
        return 0;
    }

    pthread_mutex_lock(&g_pm_lock);
    state->enabled = true;
    pthread_mutex_unlock(&g_pm_lock);

    dev_info(dev, "runtime PM enabled\n");
    return 0;
}

int pm_runtime_disable(struct device *dev)
{
    struct redox_pm_state *state = redox_pm_get_state(dev);

    if (state == NULL) {
        return 0;
    }

    pthread_mutex_lock(&g_pm_lock);
    state->enabled = false;
    pthread_mutex_unlock(&g_pm_lock);

    dev_info(dev, "runtime PM disabled\n");
    return 0;
}

int pm_runtime_allow(struct device *dev)
{
    struct redox_pm_state *state = redox_pm_get_state(dev);

    if (state == NULL) {
        return 0;
    }

    pthread_mutex_lock(&g_pm_lock);
    state->allowed = true;
    pthread_mutex_unlock(&g_pm_lock);

    dev_info(dev, "runtime PM allowed\n");
    return 0;
}

int pm_runtime_forbid(struct device *dev)
{
    struct redox_pm_state *state = redox_pm_get_state(dev);

    if (state == NULL) {
        return 0;
    }

    pthread_mutex_lock(&g_pm_lock);
    state->allowed = false;
    pthread_mutex_unlock(&g_pm_lock);

    dev_info(dev, "runtime PM forbidden\n");
    return 0;
}

void pm_suspend_ignore_children(struct device *dev, int enable)
{
    struct redox_pm_state *state = redox_pm_get_state(dev);

    if (state == NULL) {
        return;
    }

    pthread_mutex_lock(&g_pm_lock);
    state->ignore_children = (bool)enable;
    pthread_mutex_unlock(&g_pm_lock);

    dev_info(dev, "runtime PM ignore_children=%d\n", enable ? 1 : 0);
}

/* ---- Workqueue implementation ---- */

enum {
    WORK_IDLE = 0,
    WORK_PENDING = 1,
    WORK_EXECUTING = 2,
    WORK_CANCELLED = 3,
};

static pthread_mutex_t g_work_queue_lock = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t g_work_queue_cond = PTHREAD_COND_INITIALIZER;
static struct work_struct *g_work_queue_head;
static struct work_struct *g_work_queue_tail;
static volatile int g_work_executing_count;
static volatile int g_work_worker_running;
static pthread_t g_work_worker_thread;

static void *redox_work_worker_main(void *arg)
{
    (void)arg;
    pthread_setcancelstate(PTHREAD_CANCEL_ENABLE, NULL);
    pthread_setcanceltype(PTHREAD_CANCEL_DEFERRED, NULL);

    for (;;) {
        struct work_struct *work = NULL;

        pthread_mutex_lock(&g_work_queue_lock);
        while (!g_work_queue_head && g_work_worker_running) {
            pthread_cond_wait(&g_work_queue_cond, &g_work_queue_lock);
        }
        if (!g_work_worker_running) {
            pthread_mutex_unlock(&g_work_queue_lock);
            break;
        }
        work = g_work_queue_head;
        if (work) {
            g_work_queue_head = work->next;
            if (!g_work_queue_head) {
                g_work_queue_tail = NULL;
            }
            work->next = NULL;
            work->state = WORK_EXECUTING;
            __sync_fetch_and_add(&g_work_executing_count, 1);
        }
        pthread_mutex_unlock(&g_work_queue_lock);

        if (work && work->func) {
            work->func(work);
        }

        if (work) {
            pthread_mutex_lock(&g_work_queue_lock);
            __sync_fetch_and_sub(&g_work_executing_count, 1);
            pthread_mutex_lock(&work->lock);
            if (work->state == WORK_EXECUTING) {
                work->state = WORK_IDLE;
            }
            pthread_cond_broadcast(&work->done);
            pthread_mutex_unlock(&work->lock);
            pthread_cond_broadcast(&g_work_queue_cond);
            pthread_mutex_unlock(&g_work_queue_lock);
        }
    }

    return NULL;
}

static void redox_ensure_worker(void)
{
    pthread_mutex_lock(&g_work_queue_lock);
    if (!g_work_worker_running) {
        g_work_worker_running = 1;
        pthread_create(&g_work_worker_thread, NULL, redox_work_worker_main, NULL);
    }
    pthread_mutex_unlock(&g_work_queue_lock);
}

void redox_schedule_work(struct work_struct *work)
{
    if (!work || !work->func) {
        return;
    }

    redox_ensure_worker();

    pthread_mutex_lock(&g_work_queue_lock);
    pthread_mutex_lock(&work->lock);
    if (work->state != WORK_IDLE) {
        pthread_mutex_unlock(&work->lock);
        pthread_mutex_unlock(&g_work_queue_lock);
        return;
    }
    work->state = WORK_PENDING;
    work->next = NULL;
    pthread_mutex_unlock(&work->lock);

    if (g_work_queue_tail) {
        g_work_queue_tail->next = work;
        g_work_queue_tail = work;
    } else {
        g_work_queue_head = g_work_queue_tail = work;
    }
    pthread_cond_signal(&g_work_queue_cond);
    pthread_mutex_unlock(&g_work_queue_lock);
}

static void *redox_delayed_work_timer(void *arg)
{
    struct delayed_work *dwork = arg;
    unsigned long delay_ms = dwork->delay;

    if (delay_ms > 0) {
        struct timespec ts;
        ts.tv_sec = delay_ms / 1000U;
        ts.tv_nsec = (long)(delay_ms % 1000U) * 1000000L;
        nanosleep(&ts, NULL);
    }

    pthread_mutex_lock(&dwork->work.lock);
    if (dwork->work.state == WORK_CANCELLED) {
        pthread_mutex_unlock(&dwork->work.lock);
        dwork->timer_active = 0;
        return NULL;
    }
    pthread_mutex_unlock(&dwork->work.lock);

    redox_schedule_work(&dwork->work);
    dwork->timer_active = 0;
    return NULL;
}

void redox_schedule_delayed_work(struct delayed_work *dwork, unsigned long delay)
{
    if (!dwork || !dwork->work.func) {
        return;
    }

    dwork->delay = delay;
    dwork->timer_active = 1;

    if (pthread_create(&dwork->timer_thread, NULL, redox_delayed_work_timer, dwork) != 0) {
        dwork->timer_active = 0;
        redox_schedule_work(&dwork->work);
    }
}

void redox_cancel_work_sync(struct work_struct *work)
{
    if (!work) {
        return;
    }

    pthread_mutex_lock(&g_work_queue_lock);
    if (work->state == WORK_PENDING) {
        struct work_struct **link = &g_work_queue_head;
        while (*link) {
            if (*link == work) {
                *link = work->next;
                if (g_work_queue_tail == work) {
                    g_work_queue_tail = NULL;
                    struct work_struct *t = g_work_queue_head;
                    while (t && t->next) t = t->next;
                    g_work_queue_tail = t;
                }
                break;
            }
            link = &(*link)->next;
        }
        work->next = NULL;
        work->state = WORK_CANCELLED;
        pthread_mutex_unlock(&g_work_queue_lock);
        return;
    }
    pthread_mutex_unlock(&g_work_queue_lock);

    pthread_mutex_lock(&work->lock);
    while (work->state == WORK_EXECUTING) {
        pthread_cond_wait(&work->done, &work->lock);
    }
    pthread_mutex_unlock(&work->lock);
}

void redox_cancel_delayed_work_sync(struct delayed_work *dwork)
{
    if (!dwork) {
        return;
    }

    if (dwork->timer_active) {
        pthread_mutex_lock(&dwork->work.lock);
        dwork->work.state = WORK_CANCELLED;
        pthread_mutex_unlock(&dwork->work.lock);
        pthread_join(dwork->timer_thread, NULL);
        dwork->timer_active = 0;
    }

    redox_cancel_work_sync(&dwork->work);
}

void redox_flush_workqueue(struct workqueue_struct *wq)
{
    (void)wq;
    redox_flush_scheduled_work();
}

void redox_flush_scheduled_work(void)
{
    pthread_mutex_lock(&g_work_queue_lock);
    while (g_work_queue_head || g_work_executing_count > 0) {
        pthread_cond_wait(&g_work_queue_cond, &g_work_queue_lock);
    }
    pthread_mutex_unlock(&g_work_queue_lock);
}

/* ---- Completion timeout ---- */

unsigned long redox_wait_for_completion_timeout(struct completion *c, unsigned long timeout)
{
    struct timespec abs_time;
    unsigned long timeout_ms = jiffies_to_msecs(timeout);

    clock_gettime(CLOCK_REALTIME, &abs_time);
    abs_time.tv_sec += timeout_ms / 1000U;
    abs_time.tv_nsec += (long)(timeout_ms % 1000U) * 1000000L;
    if (abs_time.tv_nsec >= 1000000000L) {
        abs_time.tv_sec++;
        abs_time.tv_nsec -= 1000000000L;
    }

    pthread_mutex_lock(&c->mutex);
    if (c->done) {
        pthread_mutex_unlock(&c->mutex);
        return timeout;
    }

    int rc = pthread_cond_timedwait(&c->cond, &c->mutex, &abs_time);
    unsigned long result = 0;
    if (rc == 0 && c->done) {
        result = 1UL;
    }
    pthread_mutex_unlock(&c->mutex);
    return result;
}