#include "yolo_runtime.h"
#include <array>
#include <cassert>
#include <cstdio>
#include <set>
#include <vector>

static yolo_mask_view_v1_t view(const unsigned char* data, int w, int h, int stride, int x=0, int y=0) {
    yolo_mask_view_v1_t m{};
    m.struct_size=sizeof(m); m.version=YOLO_MASK_VIEW_VERSION;
    m.data=data; m.width=w; m.height=h; m.row_stride_bytes=stride; m.origin_x=x; m.origin_y=y;
    return m;
}
int main() {
    // Input padding contains nonbinary sentinels; holes and disconnected pixels remain intact.
    const unsigned char a[]={255,0,255,99, 0,255,0,99, 255,0,255,99};
    const unsigned char b[]={255,255,0,99, 0,255,0,99, 255,0,0,99};
    const unsigned char c[]={255,0,0,99, 0,255,0,99, 0,0,255,99};
    std::vector<yolo_mask_view_v1_t> masks={view(a,3,3,4,1,1),view(b,3,3,4,1,1),view(c,3,3,4,1,1)};
    const int w=5,h=5,stride=17;
    std::vector<unsigned char> rgb(stride*h,77);
    yolo_error_info_v1_t error{};error.struct_size=sizeof(error);error.version=YOLO_ERROR_INFO_VERSION;
    const auto compose=[&](uint32_t n) {
        return yolo_masks_compose_rgb_ids(masks.data(),n,w,h,rgb.data(),stride,rgb.size(),&error);
    };
    // Independently reconstruct each membership from the channels and compare every input pixel.
    for(uint32_t n=1;n<=3;++n) {
        assert(compose(n)==YOLO_STATUS_OK);
        for(int y=0;y<h;++y) for(int x=0;x<w;++x) {
            std::set<int> expected,actual;
            for(uint32_t i=0;i<n;++i) {
                const auto& m=masks[i];
                if(x>=m.origin_x && x<m.origin_x+m.width && y>=m.origin_y && y<m.origin_y+m.height &&
                   m.data[(y-m.origin_y)*m.row_stride_bytes+x-m.origin_x]) expected.insert(i+1);
            }
            const auto* p=&rgb[y*stride+x*3];
            for(int k=0;k<3;++k) if(p[k]) actual.insert(p[k]);
            assert(actual==expected);
            if(expected.empty()) assert(p[0]==0 && p[1]==0 && p[2]==0);
            else if(expected.size()==1) assert(p[0]==*expected.begin() && p[0]==p[1] && p[1]==p[2]);
            else if(expected.size()==2) assert(p[0]==*expected.begin() && p[1]==*expected.rbegin() && p[2]==p[0]);
            else assert(p[0]==1 && p[1]==2 && p[2]==3);
        }
        for(int y=0;y<h;++y) assert(rgb[y*stride+15]==77 && rgb[y*stride+16]==77);
    }
    // Every error must preserve all caller bytes, including partial-overlap overflow.
    masks.push_back(masks[0]);
    const auto saved=rgb;
    assert(compose(4)==YOLO_STATUS_UNSUPPORTED && error.message[0]); assert(rgb==saved);
    masks[0].origin_x=4;
    assert(compose(1)==YOLO_STATUS_INVALID_ARGUMENT); assert(rgb==saved);
    masks[0].origin_x=1; masks[0].row_stride_bytes=2;
    assert(compose(1)==YOLO_STATUS_INVALID_ARGUMENT); assert(rgb==saved);
    masks[0].row_stride_bytes=4; masks[0].version=0;
    assert(compose(1)==YOLO_STATUS_ABI_MISMATCH); assert(rgb==saved);
    masks[0]=view(a,3,3,4,1,1);
    assert(yolo_masks_compose_rgb_ids(masks.data(),1,w,h,rgb.data(),stride,1,&error)==YOLO_STATUS_BUFFER_TOO_SMALL);
    assert(rgb==saved);
    const unsigned char bad[]={255,1};masks[0]=view(bad,2,1,2);
    assert(compose(1)==YOLO_STATUS_INVALID_ARGUMENT); assert(rgb==saved);
    assert(yolo_masks_compose_rgb_ids(nullptr,1,w,h,rgb.data(),stride,rgb.size(),nullptr)==YOLO_STATUS_INVALID_ARGUMENT);
    assert(yolo_masks_compose_rgb_ids(nullptr,256,w,h,rgb.data(),stride,rgb.size(),nullptr)==YOLO_STATUS_UNSUPPORTED);
    assert(rgb==saved);
    assert(yolo_masks_compose_rgb_ids(nullptr,0,w,h,rgb.data(),stride,rgb.size(),nullptr)==YOLO_STATUS_OK);
    for(int y=0;y<h;++y) for(int x=0;x<w*3;++x) assert(rgb[y*stride+x]==0);
    // Full byte ID range, without assigning IDs from semantic class IDs.
    const unsigned char foreground=255;
    masks.clear();
    for(int i=0;i<255;++i) masks.push_back(view(&foreground,1,1,1,i,0));
    rgb.assign(255*3,0);
    assert(yolo_masks_compose_rgb_ids(masks.data(),255,255,1,rgb.data(),765,rgb.size(),nullptr)==YOLO_STATUS_OK);
    for(int i=0;i<255;++i) for(int k=0;k<3;++k) assert(rgb[i*3+k]==i+1);
    std::puts("RGB_ID_MASK_OK: exact encoding, independent decode, offsets/strides/holes, 255 IDs, atomic errors");
}