#include "zed/capture.h"

#include <opencv2/opencv.hpp>

#include <algorithm>
#include <chrono>
#include <cmath>
#include <iostream>

namespace {

using Clock = std::chrono::steady_clock;

// Going this long without a usable frame means the camera is wedged rather than
// dropping the odd frame. Measured in elapsed time rather than in failed grabs,
// because a stalling camera can spend seconds inside a single grab.
constexpr double STALL_TIMEOUT_S = 5.0;

// How long the SDK is left to reconnect the camera on its own before the loop
// stops waiting on it and rebuilds the camera anyway. A driver that says it is
// recovering is usually right, but it has no way to tell us it has given up.
constexpr double DRIVER_RECOVERY_TIMEOUT_S = 30.0;

// How long to wait between grabs while the camera is not delivering.
constexpr std::chrono::milliseconds GRAB_RETRY{500};

// How long to wait before each reopen: one entry per attempt, the last repeated
// for every attempt beyond it.
constexpr int REOPEN_BACKOFF_S[] = {1, 2, 5, 10};

std::chrono::seconds reopen_backoff(int attempt) {
    const size_t index = std::min(static_cast<size_t>(attempt), std::size(REOPEN_BACKOFF_S)) - 1;
    return std::chrono::seconds(REOPEN_BACKOFF_S[index]);
}

double seconds_since(Clock::time_point since) {
    return std::chrono::duration<double>(Clock::now() - since).count();
}

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

// Copy the depth map into a tightly packed [H, W] buffer of millimetres.
//
// The rows of an sl::Mat can be padded, so the buffer is not necessarily H * W
// floats. inf and nan mark depth the SDK could not measure; they are cleared to
// 0 here, as the reference driver does, because a hole is meaningful and a
// consumer cannot back-project a non-finite value.
void pack_depth(sl::Mat& depth_mat, std::vector<float>& out) {
    const int height = depth_mat.getHeight();
    const int width = depth_mat.getWidth();
    const size_t step = depth_mat.getStep(sl::MEM::CPU);
    const float* src = depth_mat.getPtr<sl::float1>();

    out.resize(static_cast<size_t>(height) * static_cast<size_t>(width));
    for (int y = 0; y < height; y++) {
        const float* row = src + static_cast<size_t>(y) * step;
        float* dst = out.data() + static_cast<size_t>(y) * static_cast<size_t>(width);
        for (int x = 0; x < width; x++) {
            dst[x] = std::isfinite(row[x]) ? row[x] : 0.0f;
        }
    }
}

} // namespace

namespace zed {

Capture::Capture(const Config& config) : config_(config), camera_(config) {}

Capture::~Capture() {
    stop();
    if (thread_.joinable()) {
        // The thread may still be inside the SDK. Detaching leaks it for the few
        // moments the process has left, which is preferable to blocking here.
        if (wait_exit(std::chrono::seconds(1))) {
            thread_.join();
        } else {
            thread_.detach();
        }
    }
}

void Capture::start() {
    {
        std::lock_guard<std::mutex> lock(mutex_);
        running_ = true;
    }
    thread_ = std::thread(&Capture::run, this);
}

void Capture::stop() {
    std::lock_guard<std::mutex> lock(mutex_);
    running_ = false;
    wake_.notify_all();
}

bool Capture::wait_exit(std::chrono::milliseconds timeout) {
    std::unique_lock<std::mutex> lock(mutex_);
    return exited_.wait_for(lock, timeout, [this] { return finished_; });
}

bool Capture::take_frame(Frame& out) {
    std::lock_guard<std::mutex> lock(mutex_);
    if (!has_frame_) {
        return false;
    }
    std::swap(out, latest_);
    has_frame_ = false;
    return true;
}

bool Capture::take_state_change(std::string& out) {
    std::lock_guard<std::mutex> lock(mutex_);
    if (!state_changed_) {
        return false;
    }
    out = state_;
    state_changed_ = false;
    return true;
}

void Capture::set_state(const char* state) {
    std::lock_guard<std::mutex> lock(mutex_);
    if (state_ == state) {
        return;
    }
    state_ = state;
    state_changed_ = true;
    std::cout << "camera state: " << state_ << std::endl;
}

void Capture::sleep_interruptible(std::chrono::milliseconds duration) {
    std::unique_lock<std::mutex> lock(mutex_);
    wake_.wait_for(lock, duration, [this] { return !running_; });
}

void Capture::begin_recovery(int attempt) {
    // The first attempt only closes the handle, which is enough when the driver
    // rather than the device lost its way; past that the device itself is wedged
    // and nothing short of resetting it will bring the frames back.
    if (attempt <= 1) {
        camera_.close();
    } else {
        camera_.reset();
    }
    const auto wait = reopen_backoff(attempt);
    std::cerr << "recovering camera (attempt " << attempt << "); next open in " << wait.count()
              << "s" << std::endl;
    sleep_interruptible(std::chrono::duration_cast<std::chrono::milliseconds>(wait));
}

void Capture::pack_frame(Frame& frame, int64_t frame_id, bool camera_info_first) {
    sl::Mat& image_mat = camera_.image();
    const int width = image_mat.getWidth();
    const int height = image_mat.getHeight();

    frame.frame_id = frame_id;
    frame.device_ns = camera_.device_timestamp_ns();
    frame.host_ns = host_timestamp_ns();
    frame.width = width;
    frame.height = height;
    frame.calib = camera_.calibration();
    frame.camera_info_first = camera_info_first;

    cv::Mat image_rgba(height, width, CV_8UC4, image_mat.getPtr<sl::uchar1>(),
                       image_mat.getStep(sl::MEM::CPU) * sizeof(sl::uchar4));
    frame.image_bgr.resize(static_cast<size_t>(height) * width * 3);
    cv::Mat image_bgr(height, width, CV_8UC3, frame.image_bgr.data());
    cv::cvtColor(image_rgba, image_bgr, cv::COLOR_BGRA2BGR);

    pack_depth(camera_.depth(), frame.depth);

    frame.has_point_cloud = config_.point_cloud;
    if (!frame.has_point_cloud) {
        return;
    }

    // Unlike the image and depth, the cloud used to be published straight from
    // the SDK buffer. It has to be copied now that the reader is another thread.
    sl::Mat& pc_mat = camera_.point_cloud();
    const int pc_width = pc_mat.getWidth();
    const int pc_height = pc_mat.getHeight();
    const size_t pc_step = pc_mat.getStep(sl::MEM::CPU);
    const sl::float4* src = pc_mat.getPtr<sl::float4>();

    const int stride = config_.pc_mode == PcMode::HALF ? 2 : 1;
    frame.pc_width = pc_width / stride;
    frame.pc_height = pc_height / stride;
    frame.point_cloud.resize(static_cast<size_t>(frame.pc_height) * frame.pc_width * 4);

    for (int y = 0; y < frame.pc_height; y++) {
        const sl::float4* row = src + static_cast<size_t>(y * stride) * pc_step;
        float* dst = frame.point_cloud.data() + static_cast<size_t>(y) * frame.pc_width * 4;
        for (int x = 0; x < frame.pc_width; x++) {
            const sl::float4& point = row[x * stride];
            dst[x * 4 + 0] = point.x;
            dst[x * 4 + 1] = point.y;
            dst[x * 4 + 2] = point.z;
            dst[x * 4 + 3] = point.w;
        }
    }
}

void Capture::run() {
    int recovery_attempts = 0;
    bool streamed_once = false;
    bool stall_reported = false;
    bool camera_info_pending = false;
    auto last_good_frame_at = Clock::now();
    int64_t frame_id = 0;

    set_state("opening");

    while (true) {
        {
            std::lock_guard<std::mutex> lock(mutex_);
            if (!running_) {
                break;
            }
        }

        const char* down_state = streamed_once ? "recovering" : "opening";

        if (!camera_.is_open()) {
            set_state(down_state);
            if (!camera_.open()) {
                recovery_attempts++;
                begin_recovery(recovery_attempts);
                continue;
            }
            // Discard the first frames so auto-exposure settles. This blocks, and
            // on this thread that is fine.
            camera_.warmup();
            camera_info_pending = true;
            // Opening the ZED takes seconds. Start the stall clock from here so
            // that cost is never mistaken for a stalled stream.
            last_good_frame_at = Clock::now();
            // recovery_attempts is deliberately not cleared here. A wedged device
            // can open and still deliver nothing; clearing on open would restart
            // the escalation ladder every cycle and never reach the reset.
            continue;
        }

        // A blank frame counts as a failed grab: a camera whose link has dropped
        // does not necessarily fail its grabs, it can keep reporting success and
        // hand back a uniform frame with an empty depth map.
        if (!camera_.grab() || camera_.image_is_blank()) {
            set_state(down_state);
            const double stalled_for = seconds_since(last_good_frame_at);
            // The SDK holds the camera open while it reconnects a dropped link.
            // Taking the handle away mid-reconnect aborts that and leaves us to
            // redo it the slow way, so it gets longer than a plain stall.
            const double patience =
                camera_.recovering() ? DRIVER_RECOVERY_TIMEOUT_S : STALL_TIMEOUT_S;
            if (stalled_for < patience) {
                sleep_interruptible(GRAB_RETRY);
                continue;
            }
            if (!stall_reported) {
                stall_reported = true;
                std::cerr << "camera delivered no usable frame for " << stalled_for << "s"
                          << std::endl;
            }
            recovery_attempts++;
            begin_recovery(recovery_attempts);
            continue;
        }

        if (stall_reported) {
            std::cout << "camera recovered after " << recovery_attempts << " attempt(s)"
                      << std::endl;
            stall_reported = false;
        }
        last_good_frame_at = Clock::now();
        recovery_attempts = 0;
        streamed_once = true;

        pack_frame(scratch_, frame_id++, camera_info_pending);
        camera_info_pending = false;
        {
            std::lock_guard<std::mutex> lock(mutex_);
            // Swapping rather than assigning hands the previous frame's buffers
            // back, so steady-state capture does not reallocate.
            std::swap(latest_, scratch_);
            has_frame_ = true;
        }
        set_state("streaming");
    }

    camera_.close();
    {
        std::lock_guard<std::mutex> lock(mutex_);
        finished_ = true;
        exited_.notify_all();
    }
}

} // namespace zed