/**
 * @file cross_vendor_compute.c
 * @brief Sample application demonstrating multi-vendor GPU compute aggregation.
 *
 * This sample uses libmvgal to:
 * 1. Initialize MVGAL and detect GPUs from different vendors.
 * 2. Allocate a unified memory buffer shared across all detected GPUs.
 * 3. Split a compute workload (simulated) across multiple GPUs.
 * 4. Synchronize results using cross-vendor fences.
 *
 * SPDX-License-Identifier: MIT
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <mvgal/mvgal.h>
#include <mvgal/mvgal_scheduler.h>

#define WORKLOAD_SIZE (1024 * 1024 * 64) // 64 MB
#define NUM_CHUNKS 4

int main(int argc, char *argv[]) {
    mvgal_error_t err;
    int32_t gpu_count = 0;
    mvgal_context_t ctx = NULL;
    mvgal_buffer_t buffer = NULL;
    mvgal_memory_alloc_info_t alloc_info = {0};
    mvgal_gpu_descriptor_t gpus[MVGAL_MAX_GPUS];
    mvgal_workload_t workloads[NUM_CHUNKS];

    printf("MVGAL Cross-Vendor Compute Sample\n");
    printf("=================================\n\n");

    // 1. Initialize MVGAL
    err = mvgal_init(0);
    if (err != MVGAL_SUCCESS) {
        fprintf(stderr, "Failed to initialize MVGAL: %d\n", err);
        return EXIT_FAILURE;
    }

    // 2. Query available GPUs
    gpu_count = mvgal_gpu_enumerate(gpus, MVGAL_MAX_GPUS);
    if (gpu_count <= 0) {
        fprintf(stderr, "No compatible GPUs detected via MVGAL.\n");
        mvgal_shutdown();
        return EXIT_FAILURE;
    }

    printf("Detected %d GPUs:\n", gpu_count);
    for (int32_t i = 0; i < gpu_count; i++) {
        printf("  [%d] %s (Vendor: 0x%04x, VRAM: %llu MB)\n", 
               i, gpus[i].name, gpus[i].vendor_id, (unsigned long long)(gpus[i].vram_total / (1024 * 1024)));
    }
    printf("\n");

    // 3. Create context and set strategy
    err = mvgal_context_create(&ctx);
    if (err != MVGAL_SUCCESS) {
        fprintf(stderr, "Failed to create context: %d\n", err);
        mvgal_shutdown();
        return EXIT_FAILURE;
    }

    // Use task-based scheduling for this compute workload
    mvgal_scheduler_set_strategy(ctx, MVGAL_STRATEGY_TASK);
    printf("Using strategy: MVGAL_STRATEGY_TASK\n");

    // 4. Allocate shared memory across all GPUs
    alloc_info.size = WORKLOAD_SIZE;
    alloc_info.flags = MVGAL_MEMORY_FLAG_HOST_VALID | MVGAL_MEMORY_FLAG_GPU_VALID | MVGAL_MEMORY_FLAG_SHARED;
    alloc_info.sharing_mode = MVGAL_MEMORY_SHARING_CONCURRENT;
    alloc_info.gpu_mask = (1 << gpu_count) - 1; // All GPUs

    err = mvgal_memory_allocate(ctx, &alloc_info, &buffer);
    if (err != MVGAL_SUCCESS) {
        fprintf(stderr, "Failed to allocate shared memory: %d\n", err);
        mvgal_context_destroy(ctx);
        mvgal_shutdown();
        return EXIT_FAILURE;
    }

    printf("Allocated %zu bytes of shared cross-vendor VRAM\n", WORKLOAD_SIZE);

    // 5. Map and fill buffer with initial data
    void *ptr = NULL;
    err = mvgal_memory_map(buffer, 0, WORKLOAD_SIZE, &ptr);
    if (err == MVGAL_SUCCESS) {
        printf("Mapped buffer to CPU address %p. Initializing data...\n", ptr);
        memset(ptr, 0xAA, WORKLOAD_SIZE);
        mvgal_memory_unmap(buffer);
    } else {
        printf("Warning: Failed to map buffer (err=%d). Continuing with simulated compute.\n", err);
    }

    // 6. Submit simulated compute tasks
    printf("Submitting %d compute chunks across %d GPUs...\n", NUM_CHUNKS, gpu_count);
    
    for (int i = 0; i < NUM_CHUNKS; i++) {
        mvgal_workload_submit_info_t info = {
            .type = MVGAL_WORKLOAD_COMPUTE,
            .priority = 50,
            .deadline = 0,
            .gpu_mask = (1 << gpu_count) - 1,
            .dependency_count = 0,
            .dependencies = NULL,
            .user_data = (void*)(uintptr_t)i
        };

        err = mvgal_workload_submit(ctx, &info, &workloads[i]);
        if (err != MVGAL_SUCCESS) {
            fprintf(stderr, "Failed to submit workload chunk %d: %d\n", i, err);
        }
    }

    // 7. Synchronize and verify assignments
    printf("Waiting for completion...\n");
    for (int i = 0; i < NUM_CHUNKS; i++) {
        if (workloads[i]) {
            err = mvgal_workload_wait(workloads[i], 5000); // 5s timeout
            if (err == MVGAL_SUCCESS) {
                mvgal_workload_descriptor_t desc;
                if (mvgal_workload_get_descriptor(workloads[i], &desc) == MVGAL_SUCCESS) {
                    if (desc.assigned_gpu < (uint32_t)gpu_count) {
                        printf("  Chunk %d completed on GPU %d (%s).\n", 
                               i, desc.assigned_gpu, gpus[desc.assigned_gpu].name);
                    } else {
                        printf("  Chunk %d completed (no GPU assigned).\n", i);
                    }
                } else {
                    printf("  Chunk %d completed.\n", i);
                }
            } else {
                printf("  Chunk %d timed out or failed: %d\n", i, err);
            }
            mvgal_workload_destroy(workloads[i]);
        }
    }

    // 8. Cleanup
    mvgal_memory_free(buffer);
    mvgal_context_destroy(ctx);
    mvgal_shutdown();

    printf("\nSample finished successfully.\n");
    return EXIT_SUCCESS;
}
