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

#include <cstdlib>
#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* runtime) const { yolo_runtime_destroy(runtime); }
};

static bool valid_segmentation_mask(const yolo_result_t& result) {
    if (result.task != YOLO_TASK_SEG) return true;
    if (!result.mask_data && result.mask_width == 0 && result.mask_height == 0 &&
        result.mask_stride_bytes == 0 && result.mask_origin_x == 0 && result.mask_origin_y == 0 &&
        result.mask_bbox.x == 0.0f && result.mask_bbox.y == 0.0f &&
        result.mask_bbox.width == 0.0f && result.mask_bbox.height == 0.0f) return true;
    if (!result.mask_data || result.mask_width <= 0 || result.mask_height <= 0 ||
        result.mask_stride_bytes < result.mask_width) return false;
    if (static_cast<int>(result.mask_bbox.x) != result.mask_origin_x ||
        static_cast<int>(result.mask_bbox.y) != result.mask_origin_y ||
        static_cast<int>(result.mask_bbox.width) != result.mask_width ||
        static_cast<int>(result.mask_bbox.height) != result.mask_height) return false;
    bool foreground = false;
    for (int y = 0; y < result.mask_height; ++y) {
        const unsigned char* row = result.mask_data + static_cast<size_t>(y) * result.mask_stride_bytes;
        for (int x = 0; x < result.mask_width; ++x) {
            if (row[x] != 0 && row[x] != 255) return false;
            foreground = foreground || row[x] != 0;
        }
    }
    return foreground;
}

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");
    if (!config || !image_path || !json_path) {
        std::cerr << "Usage: yolo_cpp_example --config model.yaml --image sample.bmp --json result.json\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_runtime_t* raw_runtime = nullptr;
    if (yolo_runtime_create(config, &raw_runtime) != YOLO_STATUS_OK || !raw_runtime) return 5;
    std::unique_ptr<yolo_runtime_t, RuntimeDeleter> runtime(raw_runtime);
    if ((yolo_runtime_get_capabilities(runtime.get()) & YOLO_CAPABILITY_INFER_IMAGE) == 0) return 6;

    yolo_image_view_t image{};
    image.struct_size = sizeof(image);
    image.abi_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_result_list_t results{};
    const int status = yolo_runtime_infer_image(runtime.get(), &image, &results);
    if (status != YOLO_STATUS_OK) {
        std::cerr << "Inference failed: " << yolo_runtime_last_error(runtime.get()) << '\n';
        return 7;
    }

    std::cout << "results=" << results.count
              << " preprocess=" << results.performance.preprocess_ms
              << "ms inference=" << results.performance.inference_ms
              << "ms postprocess=" << results.performance.postprocess_ms
              << "ms total=" << results.performance.total_ms << "ms\n";
    for (int index = 0; index < results.count; ++index) {
        const auto& result = results.results[index];
        std::cout << '[' << index << "] class=" << result.class_id
                  << " label=" << result.label << " score=" << result.score
                  << " angle=" << result.bbox.angle << '\n';
        if (result.task == YOLO_TASK_SEG) {
            std::cout << "    mask_bbox=(" << result.mask_bbox.x << ',' << result.mask_bbox.y
                      << ',' << result.mask_bbox.width << ',' << result.mask_bbox.height << ')'
                      << " mask_roi=" << result.mask_width << 'x' << result.mask_height
                      << " origin=(" << result.mask_origin_x << ',' << result.mask_origin_y << ")\n";
        }
        if (!valid_segmentation_mask(result)) {
            std::cerr << "Invalid per-instance Seg mask contract\n";
            return 9;
        }
    }
    return write_result_json(json_path, config, &results) == 0 ? 0 : 8;
}