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

#include <iostream>
#include <memory>
#include <string_view>
#include <vector>
#include <limits>

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); }
};

struct MaskBuffers {
    std::vector<unsigned char*> data;
    explicit MaskBuffers(size_t count):data(count,nullptr) {}
    ~MaskBuffers() {for(auto* p:data) free(p);}
};

int main(int argc, char** argv) {
    int dilation_radius=0;
    if(parse_mask_dilation(argc,argv,&dilation_radius)!=0) {std::cerr<<"Invalid --mask-dilation: integer 0..100 required\n";return 2;}
    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|mean|fit] [--rgb-mask] [--mask-dilation 0..100] "
                     "[--roi-lt L T | --roi-whlt W H L T]\n";
        return 2;
    }
    if(geometry_name && std::string_view(geometry_name)!="rect" && std::string_view(geometry_name)!="circle" &&
       std::string_view(geometry_name)!="mean" && std::string_view(geometry_name)!="fit") 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;
    if(dilation_radius && info.task!=YOLO_TASK_SEG) return 2;

    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
            : (geometry_name && std::string_view(geometry_name)=="mean" ? YOLO_SEG_GEOMETRY_MEAN_DIAMETER_CIRCLE
            : (geometry_name && std::string_view(geometry_name)=="fit" ? YOLO_SEG_GEOMETRY_ROBUST_FIT_CIRCLE : YOLO_SEG_GEOMETRY_MIN_AREA_RECT)))
        : YOLO_SEG_GEOMETRY_NONE;
    // ROI 默认数据：W=800（宽）、H=736（高）、L=0（左）、T=0（上），单位为原图像素。
    // 默认初始化同时填写结构体大小、版本并清零保留字段，ABI 5 的字段顺序不变。
    yolo_roi_v1_t roi = yolo_roi_default();
    int roi_values[4] = {0};
    const int roi_mode = parse_example_roi_args(argc, argv, roi_values);
    int has_roi = roi_mode;
    if (roi_mode == 1) {
        // 只提供 L、T：保留上面的默认宽高 800×736，只修改左上角位置。
        roi.x = roi_values[0]; // L：例如 --roi-lt 217 738 中的 217。
        roi.y = roi_values[1]; // T：例如 --roi-lt 217 738 中的 738。
    } else if (roi_mode == 2) {
        // 完整指定 W、H、L、T；例如 --roi-whlt 771 609 217 738。
        roi = yolo_roi_from_whlt(roi_values[0], roi_values[1], roi_values[2], roi_values[3]);
    } else if (roi_mode == 0) {
        // 未传 ROI 参数时沿用模型 roi_info.txt；文件不存在则全图推理。
        // 默认初始化不代表无参数时自动裁剪为 800×736。
        has_roi = load_model_roi(config, width, height, &roi);
    }
    // ROI 必须完整位于原图内；越界、非正宽高由推理接口报错，不自动裁边。
    if (has_roi < 0) {
        fprintf(stderr, "Invalid ROI arguments, roi_info.txt or source image dimensions\n");
        return 7;
    }
    if (has_roi) std::cout << "ROI W=" << roi.width << " H=" << roi.height
                           << " L=" << roi.x << " T=" << roi.y << '\n';
    yolo_result_set_t* raw_results = nullptr;
    error = {};
    error.struct_size = sizeof(error);
    error.version = YOLO_ERROR_INFO_VERSION;
    const auto status =
        has_roi ? yolo_runtime_infer_image_roi(runtime.get(), &image, &roi, &options, &raw_results, &error)
                : 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;
    if(yolo_result_set_get_count(results.get(), &count)!=YOLO_STATUS_OK) return 8;
    std::cout << "task=" << info.task << " results=" << count << '\n';

    // 1. Read performance and typed results directly through the public C ABI.
    yolo_performance_v1_t performance{};
    performance.struct_size=sizeof(performance);performance.version=YOLO_PERFORMANCE_VERSION;
    if(yolo_result_set_get_performance(results.get(),&performance)!=YOLO_STATUS_OK) return 8;
    std::cout << "preprocess_ms=" << performance.preprocess_ms
              << " inference_ms=" << performance.inference_ms
              << " postprocess_ms=" << performance.postprocess_ms
              << " total_ms=" << performance.total_ms << '\n';
    std::vector<yolo_mask_view_v1_t> masks;
    MaskBuffers owned_masks(count);
    const auto region=has_roi?roi:yolo_roi_from_whlt(width,height,0,0);
    for(uint32_t index=0;index<count;++index) {
        yolo_result_common_v1_t common{};
        common.struct_size=sizeof(common);common.version=YOLO_RESULT_COMMON_VERSION;
        if(yolo_result_set_get_common(results.get(),index,&common)!=YOLO_STATUS_OK) return 8;
        // Instance ID is index+1; category/class ID can be shared by many instances.
        std::cout << "instance=" << index+1 << " class=" << common.class_id
                  << " label=" << common.label << " score=" << common.score << '\n';
        if(info.task!=YOLO_TASK_CLS) {
            yolo_bbox_v1_t box{};box.struct_size=sizeof(box);box.version=YOLO_BBOX_VERSION;
            if(yolo_result_set_get_bbox(results.get(),index,&box)!=YOLO_STATUS_OK) return 8;
            std::cout << "bbox=" << box.x << ',' << box.y << ',' << box.width << ',' << box.height << '\n';
        }
        if(info.task==YOLO_TASK_OBB) {
            yolo_rotated_rect_v1_t box{};box.struct_size=sizeof(box);box.version=YOLO_ROTATED_RECT_VERSION;
            if(yolo_result_set_get_obb(results.get(),index,&box)!=YOLO_STATUS_OK) return 8;
            std::cout << "obb center=" << box.center_x << ',' << box.center_y
                      << " size=" << box.width << ',' << box.height << " radians=" << box.angle_radians << '\n';
        }
        if(info.task==YOLO_TASK_POSE) {
            uint32_t keypoint_count=0;
            if(yolo_result_set_get_pose_keypoint_count(results.get(),index,&keypoint_count)!=YOLO_STATUS_OK) return 8;
            std::vector<yolo_keypoint_v1_t> keypoints(keypoint_count);
            if(yolo_result_set_copy_pose_keypoints(results.get(),index,keypoints.data(),keypoint_count,&keypoint_count)!=YOLO_STATUS_OK) return 8;
            for(const auto& point:keypoints)
                std::cout << "keypoint=" << point.x << ',' << point.y << " score=" << point.score << '\n';
        }
        if(info.task==YOLO_TASK_SEG) {
            yolo_seg_geometry_v1_t geometry{};
            geometry.struct_size=sizeof(geometry);geometry.version=YOLO_SEG_GEOMETRY_VERSION;
            if(yolo_result_set_get_seg_geometry(results.get(),index,&geometry)!=YOLO_STATUS_OK) return 8;
            std::cout << "seg type=" << geometry.type << " center=" << geometry.center_x << ',' << geometry.center_y
                      << " radius=" << geometry.radius << " size=" << geometry.width << ',' << geometry.height
                      << " radians=" << geometry.angle_radians << " fit=" << geometry.fit_status << '\n';

            // 2. Each borrowed binary mask is 0/255, with local dimensions/stride
            // and an origin in the inference input image. Holes are retained.
            yolo_mask_view_v1_t mask{};
            mask.struct_size=sizeof(mask);mask.version=YOLO_MASK_VIEW_VERSION;
            if(yolo_result_set_get_seg_mask(results.get(),index,&mask)!=YOLO_STATUS_OK) return 8;
            // 独立调用公共库；原始推理掩膜和拟合几何保持原值。
            if(dilate_example_mask(&mask,dilation_radius,&region,&owned_masks.data[index])!=0) return 8;
            masks.push_back(mask);
            std::cout << "binary mask origin=" << mask.origin_x << ',' << mask.origin_y
                      << " size=" << mask.width << ',' << mask.height << " stride=" << mask.row_stride_bytes << '\n';
        }
    }

    // 3. Every Seg run demonstrates RGB instance IDs in memory by default.
    // --rgb-mask is still accepted for existing scripts, but is no longer needed.
    if(info.task==YOLO_TASK_SEG) {
        const uint64_t stride=uint64_t(width)*3;
        const uint64_t bytes=stride*uint64_t(height);
        if(stride>uint64_t(std::numeric_limits<int32_t>::max()) ||
           bytes>uint64_t(std::numeric_limits<size_t>::max())) return 8;
        std::vector<unsigned char> rgb(static_cast<size_t>(bytes)); // caller owns this buffer
        error={};error.struct_size=sizeof(error);error.version=YOLO_ERROR_INFO_VERSION;
        if(yolo_masks_compose_rgb_ids(masks.data(),count,width,height,rgb.data(),
               static_cast<int32_t>(stride),bytes,&error)!=YOLO_STATUS_OK) {
            std::cerr << "RGB mask failed: " << error.message << '\n';return 8;
        }

        // 4. Read actual pixel membership: RGB channels encode up to 3 IDs.
        // Background=(0,0,0), single=(id,id,id), double=(id1,id2,id1),
        // triple=(id1,id2,id3). Remove zero/duplicate channels to decode IDs.
        uint64_t foreground=0,overlap_pixels=0;
        std::vector<uint64_t> instance_pixels(count,0);
        for(int y=0;y<height;++y) for(int x=0;x<width;++x) {
            const auto* pixel=rgb.data()+size_t(y)*static_cast<size_t>(stride)+size_t(x)*3;
            unsigned char ids[3]{};int id_count=0;
            for(int channel=0;channel<3;++channel) {
                const auto id=pixel[channel];
                if(!id) continue;
                bool duplicate=false;
                for(int j=0;j<id_count;++j) if(ids[j]==id) duplicate=true;
                if(!duplicate) {
                    if(id>count) return 8;
                    ids[id_count++]=id;++instance_pixels[id-1];
                }
            }
            if(id_count) {
                if(!foreground) {
                    std::cout << "RGB_PIXEL x=" << x << " y=" << y << " RGB="
                              << int(pixel[0]) << ',' << int(pixel[1]) << ',' << int(pixel[2]) << " IDs=";
                    for(int j=0;j<id_count;++j) std::cout << (j ? "," : "") << int(ids[j]);
                    std::cout << '\n';
                }
                ++foreground;if(id_count>1) ++overlap_pixels;
            }
        }
        for(uint32_t i=0;i<count;++i) std::cout << "ID_PIXELS id=" << i+1 << " pixels=" << instance_pixels[i] << '\n';
        std::cout << "RGB_ID_MASK_OK width=" << width << " height=" << height << " instances=" << count
                  << " foreground=" << foreground << " overlaps=" << overlap_pixels << " bytes=" << bytes << '\n';
        // Use rgb here for your processing. No BMP is needed. rgb is freed by
        // vector; borrowed masks stay valid until results is released below.
    }
    std::cout << "MASK_DILATION radius=" << dilation_radius << '\n';
    const int json_status=write_result_json_masks(json_path,config,results.get(),info.task,
        info.task==YOLO_TASK_SEG?masks.data():nullptr,dilation_radius);
    return json_status==0 ? 0 : 8;
}