#include "zed/camera.h"
#include "util/env.h"

#include <algorithm>
#include <cstdlib>
#include <iostream>
#include <iterator>
#include <string>
#include <utility>

namespace {

sl::RESOLUTION parse_resolution(int resolution) {
    switch (resolution) {
        case 1080:
            return sl::RESOLUTION::HD1080;
        case 2000:
            return sl::RESOLUTION::HD2K;
        case 720:
        default:
            return sl::RESOLUTION::HD720;
    }
}

const std::pair<const char*, sl::DEPTH_MODE> DEPTH_MODES[] = {
    {"NEURAL", sl::DEPTH_MODE::NEURAL},
    {"NEURAL_PLUS", sl::DEPTH_MODE::NEURAL_PLUS},
    {"NEURAL_LIGHT", sl::DEPTH_MODE::NEURAL_LIGHT},
    {"PERFORMANCE", sl::DEPTH_MODE::PERFORMANCE},
    {"QUALITY", sl::DEPTH_MODE::QUALITY},
    {"NONE", sl::DEPTH_MODE::NONE},
};

// A typo here silently changes depth quality if it falls back to a default, and
// the resulting depth is wrong in a way that only shows up downstream, so an
// unrecognised name stops the node instead.
sl::DEPTH_MODE parse_depth_mode(const std::string& mode) {
    for (const auto& entry : DEPTH_MODES) {
        if (mode == entry.first) {
            return entry.second;
        }
    }
    std::cerr << "Unknown ZED_DEPTH_MODE '" << mode << "'. Accepted values:";
    for (const auto& entry : DEPTH_MODES) {
        std::cerr << " " << entry.first;
    }
    std::cerr << std::endl;
    std::exit(1);
}

sl::FLIP_MODE parse_flip(const std::string& mode) {
    if (mode == "OFF") {
        return sl::FLIP_MODE::OFF;
    }
    return sl::FLIP_MODE::ON;
}

// A blank frame is checked on a strided sample rather than the full image, which
// is just as conclusive and costs nothing at frame rate.
constexpr int BLANK_FRAME_STRIDE = 16;

// sl::Camera::reboot() reaches the ZED 2 and newer. On the first-generation ZED
// and the ZED Mini it returns without touching the device, so calling it there
// only costs time and fills the log.
bool model_supports_reboot(sl::MODEL model) {
    return model != sl::MODEL::ZED && model != sl::MODEL::ZED_M;
}

// The SDK reconnects a camera whose link hiccuped by itself. Dropping the handle
// over one of these takes the recovery away from the SDK and hands it to us,
// which is slower and, on a device the kernel is about to suspend, often fatal.
bool is_transient_grab_error(sl::ERROR_CODE err) {
    return err == sl::ERROR_CODE::CAMERA_REBOOTING || err == sl::ERROR_CODE::CORRUPTED_FRAME;
}

void list_available_cameras() {
    auto devices = sl::Camera::getDeviceList();
    std::cout << "Detected ZED cameras: " << devices.size() << std::endl;
    for (std::size_t i = 0; i < devices.size(); i++) {
        const auto& dev = devices[i];
        std::cout << " - [" << i << "] model=" << sl::toString(dev.camera_model)
                  << ", serial=" << dev.serial_number << std::endl;
    }
}

} // namespace

namespace zed {

Config Config::from_env() {
    Config c;
    c.resolution = util::get_env_int("ZED_RESOLUTION", c.resolution);
    c.fps = util::get_env_int("ZED_FPS", c.fps);
    c.depth_mode = parse_depth_mode(util::get_env_str("ZED_DEPTH_MODE", "NEURAL"));
    c.depth_min_mm = util::get_env_int("ZED_DEPTH_MIN_MM", c.depth_min_mm);
    c.depth_max_mm = util::get_env_int("ZED_DEPTH_MAX_MM", c.depth_max_mm);
    c.depth_fill = util::get_env_bool("ZED_DEPTH_FILL", c.depth_fill);
    c.point_cloud = util::get_env_bool("ZED_POINT_CLOUD", c.point_cloud);
    c.pc_mode = util::get_env_str("ZED_PC_MODE", "FULL") == "HALF" ? PcMode::HALF : PcMode::FULL;
    c.confidence_threshold = util::get_env_int("ZED_CONFIDENCE_THRESHOLD", c.confidence_threshold);
    c.texture_confidence_threshold =
        util::get_env_int("ZED_TEXTURE_CONFIDENCE_THRESHOLD", c.texture_confidence_threshold);
    c.flip = util::get_env_str("ZED_FLIP", c.flip);
    c.warmup_frames = util::get_env_int("ZED_WARMUP_FRAMES", c.warmup_frames);
    return c;
}

Camera::Camera(const Config& config) : config_(config) {
    // Process-wide, and must be set before any camera is opened. Without it the
    // epoch of the timestamps published alongside every frame depends on the SDK
    // default rather than on anything this node states.
    sl::setTimestampClock(sl::TIMESTAMP_CLOCK::SYSTEM_CLOCK);

    init_params_.camera_resolution = parse_resolution(config_.resolution);
    init_params_.camera_fps = config_.fps;
    init_params_.depth_mode = config_.depth_mode;
    init_params_.coordinate_units = sl::UNIT::MILLIMETER;
    init_params_.camera_image_flip = parse_flip(config_.flip);
    init_params_.depth_minimum_distance = config_.depth_min_mm;
    init_params_.depth_maximum_distance = config_.depth_max_mm;
    // Recover a hiccuping link inside the SDK rather than out here. Without this
    // a grab blocks for the roughly 17 s the SDK spends reconnecting, freezing the
    // dora event loop; with it the grab returns CAMERA_REBOOTING straight away and
    // the SDK reconnects in the background, holding the device open meanwhile.
    init_params_.async_grab_camera_recovery = true;
    // Let the SDK flag the frames a stalled camera hands back looking valid.
    init_params_.enable_image_validity_check = 1;

    runtime_params_.enable_depth = true;
    runtime_params_.enable_fill_mode = config_.depth_fill;
    // The two thresholds discard the depth the SDK is least certain of, which is
    // the depth at object edges: those pixels straddle the object and whatever is
    // behind it, and belong to neither.
    runtime_params_.confidence_threshold = config_.confidence_threshold;
    runtime_params_.texture_confidence_threshold = config_.texture_confidence_threshold;

    std::cout << "ZED depth: mode=" << sl::toString(init_params_.depth_mode)
              << ", range=" << config_.depth_min_mm << "-" << config_.depth_max_mm << " mm"
              << ", confidence=" << config_.confidence_threshold
              << ", texture_confidence=" << config_.texture_confidence_threshold
              << ", fill=" << (config_.depth_fill ? "on" : "off") << std::endl;

    list_available_cameras();
}

Camera::~Camera() {
    close();
}

bool Camera::open() {
    sl::ERROR_CODE err = zed_.open(init_params_);
    if (err != sl::ERROR_CODE::SUCCESS) {
        std::cerr << "Failed to open ZED camera: " << sl::toString(err) << std::endl;
        return false;
    }
    recovering_ = false;

    sl::CameraInformation info = zed_.getCameraInformation();
    serial_number_ = info.serial_number;
    model_ = info.camera_model;
    // After the first open, reopen the same physical camera. Recovery runs while
    // the device may be mid-re-enumeration, and picking whichever ZED answers
    // first is not the same thing as picking ours.
    init_params_.input.setFromSerialNumber(serial_number_);

    std::cout << "ZED camera opened (" << sl::toString(model_) << ", SN " << serial_number_ << ")"
              << std::endl;
    return true;
}

void Camera::close() {
    recovering_ = false;
    zed_.close();
}

void Camera::reset() {
    close();
    if (serial_number_ == 0 || !model_supports_reboot(model_)) {
        return;
    }
    // The ZED can stop delivering frames while staying enumerated: the video
    // endpoint halts, grabs fail, and reopening returns CAMERA FAILED TO SETUP
    // because the endpoint is still halted. A device reset clears that, and
    // reboot() is the SDK's way to ask for one without unplugging the camera.
    // The device re-enumerates afterwards, so the caller's backoff before the
    // next open is what lets it settle.
    sl::ERROR_CODE err = sl::Camera::reboot(static_cast<int>(serial_number_), true);
    std::cout << "reset ZED camera (SN " << serial_number_ << "): " << sl::toString(err)
              << std::endl;
}

bool Camera::is_open() const {
    return zed_.isOpened();
}

bool Camera::recovering() const {
    return recovering_;
}

void Camera::warmup() {
    for (int i = 0; i < config_.warmup_frames; i++) {
        zed_.grab(runtime_params_);
    }
}

bool Camera::grab() {
    sl::ERROR_CODE err = zed_.grab(runtime_params_);
    if (err != sl::ERROR_CODE::SUCCESS) {
        if (is_transient_grab_error(err)) {
            // Every grab repeats this for as long as the reconnect takes, so only
            // the first one is worth a line in the log.
            if (!recovering_) {
                recovering_ = true;
                std::cerr << "ZED camera link dropped (" << sl::toString(err)
                          << "); the SDK is reconnecting it" << std::endl;
            }
        } else {
            recovering_ = false;
            std::cerr << "Failed to grab frame: " << sl::toString(err) << std::endl;
        }
        return false;
    }
    if (recovering_) {
        recovering_ = false;
        std::cout << "ZED camera link restored" << std::endl;
    }

    zed_.retrieveImage(image_mat_, sl::VIEW::LEFT);
    zed_.retrieveMeasure(depth_mat_, sl::MEASURE::DEPTH);
    // At 1080p a full XYZRGBA cloud is ~14.7 MB per frame, and retrieving it
    // costs capture time whether or not anything consumes it.
    if (config_.point_cloud) {
        zed_.retrieveMeasure(point_cloud_mat_, sl::MEASURE::XYZRGBA);
    }
    return true;
}

Calibration Camera::calibration() const {
    sl::CameraInformation info = zed_.getCameraInformation();
    auto calib = info.camera_configuration.calibration_parameters;
    auto left = calib.left_cam;

    Calibration result;
    result.fx = left.fx;
    result.fy = left.fy;
    result.cx = left.cx;
    result.cy = left.cy;
    result.distortion.reserve(std::size(left.disto));
    for (double v : left.disto) {
        result.distortion.push_back(v);
    }
    return result;
}

bool Camera::image_is_blank() {
    const int height = image_mat_.getHeight();
    const int width = image_mat_.getWidth();
    if (height == 0 || width == 0) {
        return true;
    }

    const size_t step = image_mat_.getStep(sl::MEM::CPU);
    const sl::uchar4* src = image_mat_.getPtr<sl::uchar4>();

    unsigned char lowest[3] = {255, 255, 255};
    unsigned char highest[3] = {0, 0, 0};
    for (int y = 0; y < height; y += BLANK_FRAME_STRIDE) {
        const sl::uchar4* row = src + static_cast<size_t>(y) * step;
        for (int x = 0; x < width; x += BLANK_FRAME_STRIDE) {
            const unsigned char channels[3] = {row[x].x, row[x].y, row[x].z};
            for (int c = 0; c < 3; c++) {
                lowest[c] = std::min(lowest[c], channels[c]);
                highest[c] = std::max(highest[c], channels[c]);
            }
        }
    }
    // Sensor noise makes a single-colour frame impossible for a live camera, so
    // this tells a dead device apart from any real scene.
    return lowest[0] == highest[0] && lowest[1] == highest[1] && lowest[2] == highest[2];
}

int64_t Camera::device_timestamp_ns() {
    return static_cast<int64_t>(
        zed_.getTimestamp(sl::TIME_REFERENCE::IMAGE).getNanoseconds());
}

sl::Mat& Camera::image() {
    return image_mat_;
}

sl::Mat& Camera::depth() {
    return depth_mat_;
}

sl::Mat& Camera::point_cloud() {
    return point_cloud_mat_;
}

} // namespace zed