#include "arrow/array.h"
#include "zed/capture.h"

#include <dora-node-api.h>

#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <initializer_list>
#include <iostream>
#include <string>
#include <utility>
#include <vector>

namespace {

rust::Vec<int64_t> make_vector_int(std::initializer_list<int64_t> values) {
    rust::Vec<int64_t> out;
    for (int64_t v : values) {
        out.push_back(v);
    }
    return out;
}

rust::Vec<double> make_distortion(const std::vector<double>& distortion) {
    rust::Vec<double> values;
    for (double v : distortion) {
        values.push_back(v);
    }
    return values;
}

void set_intrinsics(Metadata& metadata, const zed::Calibration& calib) {
    metadata.set_float("fx", calib.fx);
    metadata.set_float("fy", calib.fy);
    metadata.set_float("cx", calib.cx);
    metadata.set_float("cy", calib.cy);
    metadata.set_list_float("distortion", make_distortion(calib.distortion));
}

// Both clocks are needed. The device timestamp is the instant the frame was
// captured and is the only value two nodes reading the same frame can agree on.
// The host timestamp is what the rest of the dataflow and the session recorder
// share, and on a Jetson it can step when NTP synchronises after boot - which is
// exactly why the device timestamp is recorded alongside it rather than instead
// of it.
void set_frame_stamps(Metadata& metadata, const zed::Frame& frame) {
    metadata.set_int("frame_id", frame.frame_id);
    metadata.set_timestamp("device_timestamp_ns", frame.device_ns);
    metadata.set_timestamp("host_timestamp_ns", frame.host_ns);
}

// Publish one Arrow array. dora moves both structs out and releases them once
// the payload is in shared memory, so the buffers they borrow only have to
// outlive this call - which the frame the caller is holding does.
void send(DoraNode& dora_node, const char* id, arrow_c::Exported payload,
          rust::Box<Metadata> metadata) {
    auto result = send_arrow_output(dora_node.send_output, id,
                                    reinterpret_cast<uint8_t*>(&payload.array),
                                    reinterpret_cast<uint8_t*>(&payload.schema),
                                    std::move(metadata));

    if (!std::string(result.error).empty()) {
        std::cerr << "Error sending " << id << ": " << std::string(result.error) << std::endl;
    }
}

int64_t host_timestamp_ns() {
    return std::chrono::duration_cast<std::chrono::nanoseconds>(
               std::chrono::system_clock::now().time_since_epoch())
        .count();
}

} // namespace

int main() {
    // Read the environment before joining the dataflow, so a bad configuration
    // fails with its own message rather than after the node has registered.
    auto config = zed::Config::from_env();

    auto dora_node = init_dora_node();

    // The camera runs on its own thread. The SDK can block for an unbounded time
    // inside grab() and open() when the camera is disconnected, so nothing below
    // may call into it: this loop only reads what the capture thread has already
    // produced, which is what keeps Stop answerable while the camera is down.
    zed::Capture capture(config);
    capture.start();

    zed::Frame frame;
    while (true) {
        auto event = dora_node.events->next();
        auto type = event_type(event);

        if (type == DoraEventType::Stop || type == DoraEventType::AllInputsClosed) {
            break;
        }

        if (type != DoraEventType::Input) {
            continue;
        }

        std::string state;
        if (capture.take_state_change(state)) {
            auto state_metadata = new_metadata();
            state_metadata->set_timestamp("host_timestamp_ns", host_timestamp_ns());

            send(dora_node, "camera_state", arrow_c::string(state.data(), state.size()),
                 std::move(state_metadata));
        }

        // No new frame since the last tick: the camera is slower than the timer,
        // or it is down and the capture thread is recovering it.
        if (!capture.take_frame(frame)) {
            continue;
        }

        auto width = static_cast<uint32_t>(frame.width);
        auto height = static_cast<uint32_t>(frame.height);

        // --- image ---
        // The payload carries its own width, height and encoding, so none of
        // that travels as metadata: what is left is what describes the frame
        // rather than the picture.
        auto image_metadata = new_metadata();
        set_intrinsics(*image_metadata, frame.calib);
        set_frame_stamps(*image_metadata, frame);

        send(dora_node, "image",
             arrow_c::image(width, height, "bgr8", frame.image_bgr.data(), frame.image_bgr.size()),
             std::move(image_metadata));

        // --- depth ---
        // A depth map is an image whose pixels are distances, which is how ROS
        // spells it too: 32FC1 is one float32 per pixel, and `units` says what
        // those floats measure.
        auto depth_metadata = new_metadata();
        depth_metadata->set_string("units", "mm");
        set_intrinsics(*depth_metadata, frame.calib);
        set_frame_stamps(*depth_metadata, frame);

        send(dora_node, "depth",
             arrow_c::image(width, height, "32FC1",
                            reinterpret_cast<const uint8_t*>(frame.depth.data()),
                            frame.depth.size() * sizeof(float)),
             std::move(depth_metadata));

        // --- point_cloud ---
        if (frame.has_point_cloud) {
            // Unlike the two images, a cloud is a flat array of values: its
            // grid shape is only in the metadata.
            auto pc_metadata = new_metadata();
            pc_metadata->set_string("channels", "XYZRGBA");
            pc_metadata->set_string("units", "mm");
            pc_metadata->set_string("layout", "HWC");
            pc_metadata->set_list_int("shape",
                                      make_vector_int({frame.pc_height, frame.pc_width, 4}));
            set_intrinsics(*pc_metadata, frame.calib);
            set_frame_stamps(*pc_metadata, frame);

            send(dora_node, "point_cloud",
                 arrow_c::float32(frame.point_cloud.data(), frame.point_cloud.size()),
                 std::move(pc_metadata));
        }

        // --- camera_info ---
        if (frame.camera_info_first) {
            const float camera_data[4] = {frame.calib.fx, frame.calib.fy, frame.calib.cx,
                                          frame.calib.cy};

            auto cam_metadata = new_metadata();
            cam_metadata->set_list_float("distortion", make_distortion(frame.calib.distortion));
            set_frame_stamps(*cam_metadata, frame);

            send(dora_node, "camera_info", arrow_c::float32(camera_data, 4),
                 std::move(cam_metadata));
        }
    }

    capture.stop();
    if (!capture.wait_exit(std::chrono::seconds(2))) {
        // The capture thread is inside a blocking SDK call, where nothing can
        // reach it. Leave rather than let dora wait out its grace period and kill
        // the node: the frames are gone either way, and Stop has been answered.
        std::cerr << "capture thread still in the SDK; exiting anyway" << std::endl;
        std::cout.flush();
        std::_Exit(0);
    }
    return 0;
}