/*
 * GSP vGPU-IO Example Application
 * Demonstrates basic VF creation, memory allocation, and BAR configuration
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <libvgpu.h>

static volatile bool running = true;

static void signal_handler(int sig)
{
    (void)sig;
    running = false;
}

int main(int argc, char *argv[])
{
    const char *device_path = "/dev/gsp-vgpu";
    vgpu_device_t device = NULL;
    vgpu_vf_t vf = NULL;
    vgpu_handle_t handle = 0;
    uint64_t va = 0;
    uint64_t dma_addr = 0;
    int ret;

    if (argc > 1)
        device_path = argv[1];

    printf("GSP vGPU-IO Example Application\n");
    printf("Device: %s\n", device_path);

    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);

    /* Open device */
    ret = vgpu_open(device_path, &device);
    if (ret != VGPU_OK) {
        fprintf(stderr, "Failed to open device: %s\n", vgpu_strerror(ret));
        return 1;
    }
    printf("Device opened successfully\n");

    /* Get device info */
    char info[256];
    ret = vgpu_get_info(device, info, sizeof(info));
    if (ret == VGPU_OK)
        printf("Device info: %s\n", info);

    /* Create VF */
    vgpu_vf_config_t config = {
        .id = 1,
        .bar1_offset = 0x10000000,
        .bar1_size = 0x40000,
        .pt_base = 0,
        .pt_size = 0,
        .cmd_base = 0,
        .cmd_size = 0,
        .max_memory_mb = 1024,
        .num_engines = 4,
    };

    ret = vgpu_vf_create(device, &config, &vf);
    if (ret != VGPU_OK) {
        fprintf(stderr, "Failed to create VF: %s\n", vgpu_strerror(ret));
        vgpu_close(device);
        return 1;
    }
    printf("VF %u created\n", config.id);

    /* Configure BAR1 */
    ret = vgpu_bar_config(vf, 0x10000000, 0x40000);
    if (ret != VGPU_OK) {
        fprintf(stderr, "Failed to configure BAR: %s\n", vgpu_strerror(ret));
        vgpu_vf_destroy(vf);
        vgpu_close(device);
        return 1;
    }
    printf("BAR1 configured\n");

    /* Allocate VRAM */
    const uint64_t alloc_size = 1024 * 1024;  /* 1 MB */
    ret = vgpu_mem_alloc(vf, alloc_size, VGPU_MEM_VRAM, 0x1000, 0,
                         0, &handle, &va);
    if (ret != VGPU_OK) {
        fprintf(stderr, "Failed to allocate memory: %s\n", vgpu_strerror(ret));
        vgpu_vf_destroy(vf);
        vgpu_close(device);
        return 1;
    }
    printf("Allocated %llu bytes at VA 0x%llx (handle 0x%llx)\n",
           (unsigned long long)alloc_size, (unsigned long long)va,
           (unsigned long long)handle);

    /* Map memory */
    uint64_t mapped_va = 0;
    ret = vgpu_mem_map(vf, handle, va, alloc_size, 0x3, 0, 0, &mapped_va);
    if (ret != VGPU_OK) {
        fprintf(stderr, "Failed to map memory: %s\n", vgpu_strerror(ret));
        vgpu_mem_free(vf, handle);
        vgpu_vf_destroy(vf);
        vgpu_close(device);
        return 1;
    }
    printf("Mapped at 0x%llx\n", (unsigned long long)mapped_va);

    /* Pin for DMA */
    ret = vgpu_mem_pin(vf, handle, 0, alloc_size, VGPU_PIN_READ | VGPU_PIN_WRITE, &dma_addr);
    if (ret != VGPU_OK) {
        fprintf(stderr, "Failed to pin memory: %s\n", vgpu_strerror(ret));
    } else {
        printf("Pinned for DMA at 0x%llx\n", (unsigned long long)dma_addr);
    }

    /* Protect with WPR (write) */
    ret = vgpu_wpr_protect_write(vf, va, alloc_size);
    if (ret != VGPU_OK) {
        fprintf(stderr, "Failed to protect with WPR: %s\n", vgpu_strerror(ret));
    } else {
        printf("WPR write protection enabled\n");
    }

    /* Set WPR timeout */
    ret = vgpu_wpr_set_timeout(vf, 1000000000ULL);  /* 1 second */
    if (ret != VGPU_OK) {
        fprintf(stderr, "Failed to set WPR timeout: %s\n", vgpu_strerror(ret));
    }

    /* Schedule work */
    ret = vgpu_schedule(vf, handle, 0, 0, 0, 0);
    if (ret != VGPU_OK) {
        fprintf(stderr, "Failed to schedule: %s\n", vgpu_strerror(ret));
    } else {
        printf("Work scheduled\n");
    }

    /* Run loop */
    printf("\nRunning... Press Ctrl+C to exit\n");
    while (running) {
        ret = vgpu_process_events(device, 1000);
        if (ret != VGPU_OK && ret != VGPU_ERR_TIMEOUT) {
            fprintf(stderr, "Event processing error: %s\n", vgpu_strerror(ret));
        }
    }

    printf("\nShutting down...\n");

    /* Unprotect WPR */
    vgpu_wpr_unprotect(vf, va);

    /* Unpin */
    if (dma_addr)
        vgpu_mem_unpin(vf, handle, 0, alloc_size);

    /* Unmap */
    vgpu_mem_unmap(vf, handle, va, alloc_size);

    /* Free memory */
    vgpu_mem_free(vf, handle);

    /* Destroy VF */
    vgpu_vf_destroy(vf);

    /* Close device */
    vgpu_close(device);

    printf("Done\n");
    return 0;
}