#include "bmp_loader.h"
#include "result_json.h"
#include "yolo_runtime.h"

#include <iostream>
#include <memory>
#include <string_view>

static const char* value_after(int argc, char** argv, std::string_view option) {
    for (int index = 1; index + 1 < argc; ++index)
        if (argv[index] == option) return argv[index + 1];
    return nullptr;
}

struct RuntimeDeleter {
    void operator()(yolo_runtime_t* value) const { yolo_runtime_destroy(value); }
};
struct ResultDeleter {
    void operator()(yolo_result_set_t* value) const { yolo_result_set_release(value); }
};

int main(int argc, char** argv) {
    const char* config = value_after(argc, argv, "--config");
    const char* image_path = value_after(argc, argv, "--image");
    const char* json_path = value_after(argc, argv, "--json");
    const char* geometry_name = value_after(argc, argv, "--seg-geometry");
    if (!config || !image_path || !json_path) {
        std::cerr << "Usage: yolo_cpp_example --config model.yaml --image sample.bmp "
                     "--json result.json [--seg-geometry rect|circle]\n";
        return 2;
    }
    if (yolo_runtime_get_abi_version() != YOLO_RUNTIME_ABI_VERSION) return 3;
    unsigned char* raw_pixels = nullptr;
    int width = 0, height = 0;
    if (load_bmp_rgb(image_path, &raw_pixels, &width, &height) != 0) return 4;
    std::unique_ptr<unsigned char, decltype(&free_bmp_rgb)> pixels(raw_pixels, free_bmp_rgb);

    yolo_error_info_v1_t error{};
    error.struct_size = sizeof(error);
    error.version = YOLO_ERROR_INFO_VERSION;
    yolo_runtime_t* raw_runtime = nullptr;
    if (yolo_runtime_create(config, &raw_runtime, &error) != YOLO_STATUS_OK || !raw_runtime) {
        std::cerr << "Create failed: " << error.message << '\n';
        return 5;
    }
    std::unique_ptr<yolo_runtime_t, RuntimeDeleter> runtime(raw_runtime);
    yolo_runtime_info_v1_t info{};
    info.struct_size = sizeof(info);
    info.version = YOLO_RUNTIME_INFO_VERSION;
    if (yolo_runtime_get_info(runtime.get(), &info) != YOLO_STATUS_OK) return 6;

    yolo_image_view_v1_t image{};
    image.struct_size = sizeof(image);
    image.version = YOLO_IMAGE_VIEW_ABI_VERSION;
    image.data = pixels.get();
    image.width = width;
    image.height = height;
    image.row_stride_bytes = width * 3;
    image.pixel_format = YOLO_PIXEL_FORMAT_RGB888;
    yolo_infer_options_v1_t options{};
    options.struct_size = sizeof(options);
    options.version = YOLO_INFER_OPTIONS_VERSION;
    options.seg_geometry_mode = info.task == YOLO_TASK_SEG
        ? (geometry_name && std::string_view(geometry_name) == "circle"
            ? YOLO_SEG_GEOMETRY_MIN_ENCLOSING_CIRCLE : YOLO_SEG_GEOMETRY_MIN_AREA_RECT)
        : YOLO_SEG_GEOMETRY_NONE;
    yolo_result_set_t* raw_results = nullptr;
    error = {};
    error.struct_size = sizeof(error);
    error.version = YOLO_ERROR_INFO_VERSION;
    const auto status =
        yolo_runtime_infer_image(runtime.get(), &image, &options, &raw_results, &error);
    if (status != YOLO_STATUS_OK || !raw_results) {
        std::cerr << "Inference failed: " << error.message << '\n';
        return 7;
    }
    std::unique_ptr<yolo_result_set_t, ResultDeleter> results(raw_results);
    uint32_t count = 0;
    yolo_result_set_get_count(results.get(), &count);
    std::cout << "task=" << info.task << " results=" << count << '\n';
    return write_result_json(json_path, config, results.get(), info.task) == 0 ? 0 : 8;
}