/*
 * Minimal RK3588 NPU smoke test for RKNN Runtime 2.3.2.
 *
 * The program deliberately uses deterministic synthetic inputs. Its purpose is
 * to verify model loading, core selection, command submission, and finite
 * output, not the semantic accuracy of MobileNet.
 */

#include <errno.h>
#include <inttypes.h>
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <time.h>

#include "rknn_api.h"

#define DEFAULT_MODEL "/opt/rknn/current/share/models/rk3588/mobilenet_v1.rknn"
#define MAX_TENSORS 64U

static void usage(FILE *stream, const char *program)
{
    fprintf(stream,
            "Usage: %s [--model PATH] [--core auto|0|1|2|all] "
            "[--iterations N]\n",
            program);
}

static int parse_core(const char *value, rknn_core_mask *mask)
{
    if (strcmp(value, "auto") == 0) {
        *mask = RKNN_NPU_CORE_AUTO;
    } else if (strcmp(value, "0") == 0) {
        *mask = RKNN_NPU_CORE_0;
    } else if (strcmp(value, "1") == 0) {
        *mask = RKNN_NPU_CORE_1;
    } else if (strcmp(value, "2") == 0) {
        *mask = RKNN_NPU_CORE_2;
    } else if (strcmp(value, "all") == 0) {
        *mask = RKNN_NPU_CORE_0_1_2;
    } else {
        return -1;
    }
    return 0;
}

static int load_file(const char *path, void **buffer, uint32_t *size)
{
    struct stat info;
    FILE *file = NULL;
    void *data = NULL;

    if (stat(path, &info) != 0) {
        fprintf(stderr, "cannot stat model %s: %s\n", path, strerror(errno));
        return -1;
    }
    if (info.st_size <= 0 || (uint64_t)info.st_size > UINT32_MAX) {
        fprintf(stderr, "invalid model size: %jd\n", (intmax_t)info.st_size);
        return -1;
    }

    file = fopen(path, "rb");
    if (file == NULL) {
        fprintf(stderr, "cannot open model %s: %s\n", path, strerror(errno));
        return -1;
    }
    data = malloc((size_t)info.st_size);
    if (data == NULL) {
        fprintf(stderr, "cannot allocate %jd bytes for model\n",
                (intmax_t)info.st_size);
        fclose(file);
        return -1;
    }
    if (fread(data, 1, (size_t)info.st_size, file) != (size_t)info.st_size) {
        fprintf(stderr, "short read from model %s\n", path);
        free(data);
        fclose(file);
        return -1;
    }
    fclose(file);
    *buffer = data;
    *size = (uint32_t)info.st_size;
    return 0;
}

static double elapsed_ms(const struct timespec *start, const struct timespec *end)
{
    double seconds = (double)(end->tv_sec - start->tv_sec) * 1000.0;
    double nanos = (double)(end->tv_nsec - start->tv_nsec) / 1000000.0;
    return seconds + nanos;
}

int main(int argc, char **argv)
{
    const char *model_path = DEFAULT_MODEL;
    const char *core_name = "auto";
    rknn_core_mask core_mask = RKNN_NPU_CORE_AUTO;
    unsigned long iterations = 1;
    void *model = NULL;
    uint32_t model_size = 0;
    rknn_context context = 0;
    rknn_sdk_version sdk_version;
    rknn_input_output_num io_count;
    rknn_tensor_attr *input_attrs = NULL;
    rknn_input *inputs = NULL;
    rknn_output *outputs = NULL;
    int outputs_acquired = 0;
    int context_created = 0;
    int result = EXIT_FAILURE;
    double total_ms = 0.0;
    uint32_t last_top_index = 0;
    float last_top_value = -INFINITY;
    int ret;
    uint32_t i;

    for (i = 1; i < (uint32_t)argc; ++i) {
        if (strcmp(argv[i], "--model") == 0 && i + 1U < (uint32_t)argc) {
            model_path = argv[++i];
        } else if (strcmp(argv[i], "--core") == 0 && i + 1U < (uint32_t)argc) {
            core_name = argv[++i];
            if (parse_core(core_name, &core_mask) != 0) {
                fprintf(stderr, "invalid core selector: %s\n", core_name);
                usage(stderr, argv[0]);
                return 2;
            }
        } else if (strcmp(argv[i], "--iterations") == 0 &&
                   i + 1U < (uint32_t)argc) {
            char *end = NULL;
            errno = 0;
            iterations = strtoul(argv[++i], &end, 10);
            if (errno != 0 || end == argv[i] || *end != '\0' ||
                iterations == 0 || iterations > 1000) {
                fprintf(stderr, "iterations must be between 1 and 1000\n");
                return 2;
            }
        } else if (strcmp(argv[i], "--help") == 0 ||
                   strcmp(argv[i], "-h") == 0) {
            usage(stdout, argv[0]);
            return 0;
        } else {
            fprintf(stderr, "unknown or incomplete option: %s\n", argv[i]);
            usage(stderr, argv[0]);
            return 2;
        }
    }

    if (load_file(model_path, &model, &model_size) != 0) {
        goto cleanup;
    }

    ret = rknn_init(&context, model, model_size, 0, NULL);
    if (ret != RKNN_SUCC) {
        fprintf(stderr, "rknn_init failed: %d\n", ret);
        goto cleanup;
    }
    context_created = 1;

    ret = rknn_set_core_mask(context, core_mask);
    if (ret != RKNN_SUCC) {
        fprintf(stderr, "rknn_set_core_mask(%s) failed: %d\n", core_name, ret);
        goto cleanup;
    }

    memset(&sdk_version, 0, sizeof(sdk_version));
    ret = rknn_query(context, RKNN_QUERY_SDK_VERSION, &sdk_version,
                     sizeof(sdk_version));
    if (ret != RKNN_SUCC) {
        fprintf(stderr, "RKNN_QUERY_SDK_VERSION failed: %d\n", ret);
        goto cleanup;
    }

    memset(&io_count, 0, sizeof(io_count));
    ret = rknn_query(context, RKNN_QUERY_IN_OUT_NUM, &io_count,
                     sizeof(io_count));
    if (ret != RKNN_SUCC || io_count.n_input == 0 || io_count.n_output == 0 ||
        io_count.n_input > MAX_TENSORS || io_count.n_output > MAX_TENSORS) {
        fprintf(stderr, "invalid RKNN input/output count (%u/%u), ret=%d\n",
                io_count.n_input, io_count.n_output, ret);
        goto cleanup;
    }

    input_attrs = calloc(io_count.n_input, sizeof(*input_attrs));
    inputs = calloc(io_count.n_input, sizeof(*inputs));
    outputs = calloc(io_count.n_output, sizeof(*outputs));
    if (input_attrs == NULL || inputs == NULL || outputs == NULL) {
        fprintf(stderr, "cannot allocate tensor metadata\n");
        goto cleanup;
    }

    for (i = 0; i < io_count.n_input; ++i) {
        uint32_t byte;
        input_attrs[i].index = i;
        ret = rknn_query(context, RKNN_QUERY_INPUT_ATTR, &input_attrs[i],
                         sizeof(input_attrs[i]));
        if (ret != RKNN_SUCC || input_attrs[i].n_elems == 0) {
            fprintf(stderr, "query for input %u failed: %d\n", i, ret);
            goto cleanup;
        }

        inputs[i].index = i;
        inputs[i].size = input_attrs[i].n_elems;
        inputs[i].type = RKNN_TENSOR_UINT8;
        inputs[i].fmt = input_attrs[i].fmt == RKNN_TENSOR_UNDEFINED
                            ? RKNN_TENSOR_NHWC
                            : input_attrs[i].fmt;
        inputs[i].pass_through = 0;
        inputs[i].buf = malloc(inputs[i].size);
        if (inputs[i].buf == NULL) {
            fprintf(stderr, "cannot allocate input %u (%u bytes)\n", i,
                    inputs[i].size);
            goto cleanup;
        }
        for (byte = 0; byte < inputs[i].size; ++byte) {
            ((uint8_t *)inputs[i].buf)[byte] =
                (uint8_t)((byte * 17U + i * 23U) & 0xffU);
        }
    }

    ret = rknn_inputs_set(context, io_count.n_input, inputs);
    if (ret != RKNN_SUCC) {
        fprintf(stderr, "rknn_inputs_set failed: %d\n", ret);
        goto cleanup;
    }

    for (i = 0; i < io_count.n_output; ++i) {
        outputs[i].index = i;
        outputs[i].want_float = 1;
        outputs[i].is_prealloc = 0;
    }

    for (unsigned long iteration = 0; iteration < iterations; ++iteration) {
        struct timespec start;
        struct timespec end;

        if (clock_gettime(CLOCK_MONOTONIC, &start) != 0) {
            fprintf(stderr, "clock_gettime failed: %s\n", strerror(errno));
            goto cleanup;
        }
        ret = rknn_run(context, NULL);
        if (ret != RKNN_SUCC) {
            fprintf(stderr, "rknn_run failed at iteration %lu: %d\n",
                    iteration, ret);
            goto cleanup;
        }
        ret = rknn_outputs_get(context, io_count.n_output, outputs, NULL);
        if (ret != RKNN_SUCC) {
            fprintf(stderr, "rknn_outputs_get failed at iteration %lu: %d\n",
                    iteration, ret);
            goto cleanup;
        }
        outputs_acquired = 1;
        if (clock_gettime(CLOCK_MONOTONIC, &end) != 0) {
            fprintf(stderr, "clock_gettime failed: %s\n", strerror(errno));
            goto cleanup;
        }
        total_ms += elapsed_ms(&start, &end);

        last_top_value = -INFINITY;
        last_top_index = 0;
        for (i = 0; i < io_count.n_output; ++i) {
            const float *values = outputs[i].buf;
            uint32_t count = outputs[i].size / (uint32_t)sizeof(float);
            uint32_t value_index;
            if (values == NULL || count == 0) {
                fprintf(stderr, "output %u is empty\n", i);
                goto cleanup;
            }
            for (value_index = 0; value_index < count; ++value_index) {
                if (!isfinite(values[value_index])) {
                    fprintf(stderr, "output %u contains a non-finite value\n", i);
                    goto cleanup;
                }
                if (i == 0 && values[value_index] > last_top_value) {
                    last_top_value = values[value_index];
                    last_top_index = value_index;
                }
            }
        }

        ret = rknn_outputs_release(context, io_count.n_output, outputs);
        outputs_acquired = 0;
        if (ret != RKNN_SUCC) {
            fprintf(stderr, "rknn_outputs_release failed: %d\n", ret);
            goto cleanup;
        }
    }

    printf("RKNN_RESULT status=pass core=%s iterations=%lu avg_ms=%.3f "
           "top_index=%u top_value=%.7g api=%s driver=%s\n",
           core_name, iterations, total_ms / (double)iterations,
           last_top_index, last_top_value, sdk_version.api_version,
           sdk_version.drv_version);
    result = EXIT_SUCCESS;

cleanup:
    if (outputs_acquired) {
        (void)rknn_outputs_release(context, io_count.n_output, outputs);
    }
    if (inputs != NULL) {
        for (i = 0; i < io_count.n_input; ++i) {
            free(inputs[i].buf);
        }
    }
    free(outputs);
    free(inputs);
    free(input_attrs);
    if (context_created) {
        (void)rknn_destroy(context);
    }
    free(model);
    return result;
}