#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <fstream>
#include <iterator>
#include <string>
#include <vector>

#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

#include "stereo_bolt/c_api.h"

namespace {

bool copy_matrix(const cv::FileNode& node, const char* key,
                 double* dst, int expected_values) {
    cv::Mat value;
    node[key] >> value;
    if (value.empty() || static_cast<int>(value.total()) < expected_values) {
        std::fprintf(stderr, "calibration field %s is missing or too short\n", key);
        return false;
    }
    cv::Mat flat = value.reshape(1, 1);
    cv::Mat as_double;
    flat.convertTo(as_double, CV_64F);
    for (int i = 0; i < expected_values; ++i) {
        dst[i] = as_double.at<double>(0, i);
    }
    return true;
}

bool load_calibration(const char* path, StereoBoltCalibC* out) {
    cv::FileStorage fs(path, cv::FileStorage::READ);
    if (!fs.isOpened()) {
        std::fprintf(stderr, "cannot open calibration: %s\n", path);
        return false;
    }
    const cv::FileNode left = fs["LeftCamera"];
    const cv::FileNode right = fs["RightCamera"];
    const cv::FileNode stereo = fs["Stereo"];
    if (left.empty() || right.empty() || stereo.empty()) {
        std::fprintf(stderr, "calibration is missing camera/stereo nodes\n");
        return false;
    }

    *out = StereoBoltCalibC{};
    left["ImageWidth"] >> out->image_width;
    left["ImageHeight"] >> out->image_height;
    stereo["Baseline"] >> out->baseline_mm;
    return copy_matrix(left, "CameraMatrix", out->left_K, 9) &&
           copy_matrix(left, "DistCoeffs", out->left_D, 5) &&
           copy_matrix(right, "CameraMatrix", out->right_K, 9) &&
           copy_matrix(right, "DistCoeffs", out->right_D, 5) &&
           copy_matrix(stereo, "R", out->R, 9) &&
           copy_matrix(stereo, "T", out->T, 3);
}

std::vector<uint8_t> read_bytes(const char* path) {
    std::ifstream stream(path, std::ios::binary);
    return {std::istreambuf_iterator<char>(stream), std::istreambuf_iterator<char>()};
}

}  // namespace

int main(int argc, char** argv) {
    if (argc != 7) {
        std::fprintf(stderr,
                     "usage: %s <calib.xml> <model.rknn> <left.bmp> <right.bmp> "
                     "<frame_id> <expected_count>\n",
                     argv[0]);
        return 2;
    }

    const int frame_id = std::stoi(argv[5]);
    const int expected = std::stoi(argv[6]);
    StereoBoltCalibC calib{};
    if (!load_calibration(argv[1], &calib)) return 2;

    const std::vector<uint8_t> model_bytes = read_bytes(argv[2]);
    if (model_bytes.empty()) {
        std::fprintf(stderr, "cannot read model: %s\n", argv[2]);
        return 2;
    }
    StereoBoltModelC model{};
    model.data = model_bytes.data();
    model.size = model_bytes.size();
    model.backend = SB_YOLO_BACKEND_RKNN;
    model.imgsz = 1024;
    model.conf = 0.25f;
    model.iou = 0.45f;

    StereoBoltParamsC params = sb_default_params();
    params.wd_z_min_mm = 900.0;
    params.wd_z_max_mm = 4000.0;
    params.plane_use_sgbm = 1;
    params.height_from_plane = 1;
    params.plane_height_along_axis = 1;
    params.plane_anchor_max_foot_gap_std_mm = 20.0;
    params.plane_sgbm_scale = 0.25;
    params.plane_max_samples = 40000;

    StereoBoltCtx* ctx = sb_create_ex(&calib, &model, &params);
    if (ctx == nullptr) {
        std::fprintf(stderr, "sb_create_ex failed: %s\n", sb_last_error(nullptr));
        return 2;
    }

    const cv::Mat left = cv::imread(argv[3], cv::IMREAD_GRAYSCALE);
    const cv::Mat right = cv::imread(argv[4], cv::IMREAD_GRAYSCALE);
    if (left.empty() || right.empty() || left.size() != right.size()) {
        std::fprintf(stderr, "cannot read equal-size stereo inputs\n");
        sb_destroy(ctx);
        return 2;
    }

    StereoBoltModuleResultC precise{};
    const sb_status_t precise_status = sb_process_bolt_module_buffers(
        ctx,
        left.data, left.cols, left.rows, static_cast<int>(left.step),
        right.data, right.cols, right.rows, static_cast<int>(right.step),
        1, expected, &precise);
    std::printf(
        "FRAME %d PRECISE api_status=%d success=%d yolo_left=%d yolo_right=%d "
        "pairs=%d bolts=%d reason=%s plane_rms=%.3f\n",
        frame_id, static_cast<int>(precise_status), precise.success,
        precise.failure.n_left_rois, precise.failure.n_right_rois,
        precise.failure.n_pairs, precise.data.n_bolts, precise.failure.reason,
        precise.data.ground_plane.rms_mm);
    std::printf("FRAME %d HEIGHTS", frame_id);
    for (int i = 0; i < precise.data.n_bolts; ++i) {
        const StereoBoltModuleBoltC& bolt = precise.data.bolts[i];
        std::printf(" id%d=%.3f(%s)", bolt.bolt_id, bolt.height_mm,
                    bolt.confidence == 0 ? "trusted" : "suspect");
    }
    std::printf("\n");

    const bool precise_ok =
        precise_status == SB_OK && precise.success == 1 &&
        precise.failure.n_left_rois == expected &&
        precise.failure.n_right_rois == expected &&
        precise.failure.n_pairs == expected &&
        precise.data.n_bolts == expected;
    sb_free_bolt_module_result(&precise);

    cv::Mat left_binned;
    cv::Mat right_binned;
    cv::resize(left, left_binned, cv::Size(left.cols / 2, left.rows / 2),
               0.0, 0.0, cv::INTER_AREA);
    cv::resize(right, right_binned, cv::Size(right.cols / 2, right.rows / 2),
               0.0, 0.0, cv::INTER_AREA);
    StereoBoltRangeSummaryC range{};
    std::vector<StereoBoltRangeBoltC> range_bolts(
        static_cast<size_t>(std::max(expected, 16)));
    const sb_status_t range_status = sb_range_bolt_binned(
        ctx,
        left_binned.data, left_binned.cols, left_binned.rows,
        static_cast<int>(left_binned.step),
        right_binned.data, right_binned.cols, right_binned.rows,
        static_cast<int>(right_binned.step),
        1, &range, range_bolts.data(), static_cast<int>(range_bolts.size()));
    std::printf(
        "FRAME %d RANGE api_status=%d found=%d bolts=%d distance_mm=%.3f status=%s\n",
        frame_id, static_cast<int>(range_status), range.found, range.bolt_count,
        range.distance_mm, range.status);

    const bool range_ok =
        range_status == SB_OK && range.found == 1 && range.bolt_count == expected;
    sb_destroy(ctx);
    std::printf("FRAME %d VERDICT=%s\n", frame_id,
                precise_ok && range_ok ? "PASS" : "FAIL");
    return precise_ok && range_ok ? 0 : 3;
}