#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <fstream>
#include <stdexcept>
#include <string>
#include <vector>

#include "rknn_api.h"

static std::vector<unsigned char> read_file(const char* path) {
    std::ifstream file(path, std::ios::binary | std::ios::ate);
    if (!file) {
        throw std::runtime_error(std::string("failed to open model: ") + path);
    }

    std::streamsize size = file.tellg();
    file.seekg(0, std::ios::beg);

    std::vector<unsigned char> data(size);
    if (!file.read(reinterpret_cast<char*>(data.data()), size)) {
        throw std::runtime_error(std::string("failed to read model: ") + path);
    }

    return data;
}

static const char* tensor_type_to_str(rknn_tensor_type type) {
    switch (type) {
        case RKNN_TENSOR_FLOAT32: return "FLOAT32";
        case RKNN_TENSOR_FLOAT16: return "FLOAT16";
        case RKNN_TENSOR_INT8:    return "INT8";
        case RKNN_TENSOR_UINT8:   return "UINT8";
        case RKNN_TENSOR_INT16:   return "INT16";
        case RKNN_TENSOR_UINT16:  return "UINT16";
        case RKNN_TENSOR_INT32:   return "INT32";
        case RKNN_TENSOR_UINT32:  return "UINT32";
        default:                  return "UNKNOWN";
    }
}

static const char* tensor_format_to_str(rknn_tensor_format fmt) {
    switch (fmt) {
        case RKNN_TENSOR_NCHW: return "NCHW";
        case RKNN_TENSOR_NHWC: return "NHWC";
        default:               return "UNKNOWN";
    }
}

static void print_tensor_attr(const char* tag, const rknn_tensor_attr& attr) {
    printf("%s[%d]\n", tag, attr.index);
    printf("  name: %s\n", attr.name);

    printf("  dims: ");
    for (int i = 0; i < attr.n_dims; ++i) {
        printf("%d", attr.dims[i]);
        if (i + 1 < attr.n_dims) {
            printf(" x ");
        }
    }
    printf("\n");

    printf("  n_elems: %d\n", attr.n_elems);
    printf("  size: %d\n", attr.size);
    printf("  fmt: %s\n", tensor_format_to_str(attr.fmt));
    printf("  type: %s\n", tensor_type_to_str(attr.type));
    printf("  qnt_type: %d\n", attr.qnt_type);
    printf("  zp: %d\n", attr.zp);
    printf("  scale: %f\n", attr.scale);
}

int main(int argc, char** argv) {
    if (argc < 2) {
        printf("Usage:\n");
        printf("  %s model.rknn\n", argv[0]);
        return -1;
    }

    const char* model_path = argv[1];

    try {
        std::vector<unsigned char> model_data = read_file(model_path);

        rknn_context ctx = 0;

        int ret = rknn_init(&ctx, model_data.data(), model_data.size(), 0, NULL);
        if (ret != RKNN_SUCC) {
            printf("rknn_init failed, ret=%d\n", ret);
            return -1;
        }

        printf("rknn_init OK: %s\n", model_path);

        rknn_input_output_num io_num;
        memset(&io_num, 0, sizeof(io_num));

        ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
        if (ret != RKNN_SUCC) {
            printf("RKNN_QUERY_IN_OUT_NUM failed, ret=%d\n", ret);
            rknn_destroy(ctx);
            return -1;
        }

        printf("input_num=%d, output_num=%d\n", io_num.n_input, io_num.n_output);

        std::vector<rknn_tensor_attr> input_attrs(io_num.n_input);
        std::vector<rknn_tensor_attr> output_attrs(io_num.n_output);

        for (uint32_t i = 0; i < io_num.n_input; ++i) {
            memset(&input_attrs[i], 0, sizeof(rknn_tensor_attr));
            input_attrs[i].index = i;

            ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &input_attrs[i], sizeof(rknn_tensor_attr));
            if (ret != RKNN_SUCC) {
                printf("RKNN_QUERY_INPUT_ATTR failed, index=%d, ret=%d\n", i, ret);
                rknn_destroy(ctx);
                return -1;
            }

            print_tensor_attr("input", input_attrs[i]);
        }

        for (uint32_t i = 0; i < io_num.n_output; ++i) {
            memset(&output_attrs[i], 0, sizeof(rknn_tensor_attr));
            output_attrs[i].index = i;

            ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &output_attrs[i], sizeof(rknn_tensor_attr));
            if (ret != RKNN_SUCC) {
                printf("RKNN_QUERY_OUTPUT_ATTR failed, index=%d, ret=%d\n", i, ret);
                rknn_destroy(ctx);
                return -1;
            }

            print_tensor_attr("output", output_attrs[i]);
        }

        if (io_num.n_input != 1) {
            printf("this smoke test only supports one input\n");
            rknn_destroy(ctx);
            return -1;
        }

        const rknn_tensor_attr& in_attr = input_attrs[0];

        std::vector<unsigned char> input_buffer(in_attr.size);
        memset(input_buffer.data(), 0, input_buffer.size());

        rknn_input input;
        memset(&input, 0, sizeof(input));

        input.index = 0;
        input.buf = input_buffer.data();
        input.size = input_buffer.size();
        input.pass_through = 0;
        input.type = in_attr.type;
        input.fmt = in_attr.fmt;

        ret = rknn_inputs_set(ctx, 1, &input);
        if (ret != RKNN_SUCC) {
            printf("rknn_inputs_set failed, ret=%d\n", ret);
            rknn_destroy(ctx);
            return -1;
        }

        ret = rknn_run(ctx, NULL);
        if (ret != RKNN_SUCC) {
            printf("rknn_run failed, ret=%d\n", ret);
            rknn_destroy(ctx);
            return -1;
        }

        printf("rknn_run OK\n");

        std::vector<rknn_output> outputs(io_num.n_output);

        for (uint32_t i = 0; i < io_num.n_output; ++i) {
            memset(&outputs[i], 0, sizeof(rknn_output));
            outputs[i].index = i;
            outputs[i].want_float = 0;
        }

        ret = rknn_outputs_get(ctx, io_num.n_output, outputs.data(), NULL);
        if (ret != RKNN_SUCC) {
            printf("rknn_outputs_get failed, ret=%d\n", ret);
            rknn_destroy(ctx);
            return -1;
        }

        for (uint32_t i = 0; i < io_num.n_output; ++i) {
            printf("output[%d] buf=%p size=%d\n", i, outputs[i].buf, outputs[i].size);
        }

        rknn_outputs_release(ctx, io_num.n_output, outputs.data());
        rknn_destroy(ctx);

        printf("model smoke test OK\n");
        return 0;

    } catch (const std::exception& e) {
        printf("exception: %s\n", e.what());
        return -1;
    }
}
