playback.hpp

#ifndef K4A_PLAYBACK_HPP
#define K4A_PLAYBACK_HPP

#include <k4a/k4a.hpp>
#include <k4arecord/playback.h>

#include <algorithm>
#include <chrono>
#include <string>
#include <vector>

namespace k4a
{

class data_block
{
public:
    data_block(k4a_playback_data_block_t handle = nullptr) noexcept : m_handle(handle) {}

    // No Copies allowed
    data_block(const data_block &) = delete;
    data_block &operator=(const data_block &) = delete;

    data_block(data_block &&other) noexcept
    {
        m_handle = other.m_handle;
        other.m_handle = nullptr;
    }

    ~data_block()
    {
        reset();
    }

    data_block &operator=(data_block &&other) noexcept
    {
        if (this != &other)
        {
            reset();
            m_handle = other.m_handle;
            other.m_handle = nullptr;
        }
        return *this;
    }

    explicit operator bool() const noexcept
    {
        return is_valid();
    }

    bool is_valid() const noexcept
    {
        return m_handle != nullptr;
    }

    void reset() noexcept
    {
        if (m_handle)
        {
            k4a_playback_data_block_release(m_handle);
            m_handle = nullptr;
        }
    }

    std::chrono::microseconds get_device_timestamp_usec() const noexcept
    {
        return std::chrono::microseconds(k4a_playback_data_block_get_device_timestamp_usec(m_handle));
    }

    size_t get_buffer_size() const noexcept
    {
        return k4a_playback_data_block_get_buffer_size(m_handle);
    }

    const uint8_t *get_buffer() const noexcept
    {
        return k4a_playback_data_block_get_buffer(m_handle);
    }

private:
    k4a_playback_data_block_t m_handle;
};

class playback
{
public:
    playback(k4a_playback_t handle = nullptr) noexcept : m_handle(handle) {}

    playback(playback &&other) noexcept : m_handle(other.m_handle)
    {
        other.m_handle = nullptr;
    }

    playback(const playback &) = delete;

    ~playback()
    {
        close();
    }

    playback &operator=(const playback &) = delete;

    playback &operator=(playback &&other) noexcept
    {
        if (this != &other)
        {
            close();
            m_handle = other.m_handle;
            other.m_handle = nullptr;
        }

        return *this;
    }

    explicit operator bool() const noexcept
    {
        return is_valid();
    }

    bool is_valid() const noexcept
    {
        return m_handle != nullptr;
    }

    void close() noexcept
    {
        if (m_handle != nullptr)
        {
            k4a_playback_close(m_handle);
            m_handle = nullptr;
        }
    }

    std::vector<uint8_t> get_raw_calibration() const
    {
        std::vector<uint8_t> calibration;
        size_t buffer = 0;
        k4a_buffer_result_t result = k4a_playback_get_raw_calibration(m_handle, nullptr, &buffer);

        if (result == K4A_BUFFER_RESULT_TOO_SMALL && buffer > 1)
        {
            calibration.resize(buffer);
            result = k4a_playback_get_raw_calibration(m_handle, &calibration[0], &buffer);
        }

        if (result != K4A_BUFFER_RESULT_SUCCEEDED)
        {
            throw error("Failed to read raw device calibration from recording!");
        }

        return calibration;
    }

    calibration get_calibration() const
    {
        calibration calib;
        k4a_result_t result = k4a_playback_get_calibration(m_handle, &calib);

        if (K4A_RESULT_SUCCEEDED != result)
        {
            throw error("Failed to read device calibration from recording!");
        }

        return calib;
    }

    k4a_record_configuration_t get_record_configuration() const
    {
        k4a_record_configuration_t config;
        k4a_result_t result = k4a_playback_get_record_configuration(m_handle, &config);

        if (K4A_RESULT_SUCCEEDED != result)
        {
            throw error("Failed to read record configuration!");
        }

        return config;
    }

    bool get_next_capture(capture *cap)
    {
        k4a_capture_t capture_handle;
        k4a_stream_result_t result = k4a_playback_get_next_capture(m_handle, &capture_handle);

        if (K4A_STREAM_RESULT_SUCCEEDED == result)
        {
            *cap = capture(capture_handle);
            return true;
        }
        else if (K4A_STREAM_RESULT_EOF == result)
        {
            return false;
        }

        throw error("Failed to get next capture!");
    }

    bool get_previous_capture(capture *cap)
    {
        k4a_capture_t capture_handle;
        k4a_stream_result_t result = k4a_playback_get_previous_capture(m_handle, &capture_handle);

        if (K4A_STREAM_RESULT_SUCCEEDED == result)
        {
            *cap = capture(capture_handle);
            return true;
        }
        else if (K4A_STREAM_RESULT_EOF == result)
        {
            return false;
        }

        throw error("Failed to get previous capture!");
    }

    bool get_tag(const char *name, std::string *out) const
    {
        std::string tag;
        size_t buffer = 0;
        k4a_buffer_result_t result = k4a_playback_get_tag(m_handle, name, nullptr, &buffer);

        if (result == K4A_BUFFER_RESULT_TOO_SMALL && buffer > 0)
        {
            tag.resize(buffer);
            result = k4a_playback_get_tag(m_handle, name, &tag[0], &buffer);
            if (result == K4A_BUFFER_RESULT_SUCCEEDED && tag[buffer - 1] == 0)
            {
                // std::string expects there to not be as null terminator at the end of its data but
                // k4a_playback_get_tag adds a null terminator, so we drop the last character of the string after we
                // get the result back.
                tag.resize(buffer - 1);
            }
        }

        if (result != K4A_BUFFER_RESULT_SUCCEEDED)
        {
            return false;
        }

        *out = std::move(tag);

        return true;
    }

    bool get_next_imu_sample(k4a_imu_sample_t *sample)
    {
        k4a_stream_result_t result = k4a_playback_get_next_imu_sample(m_handle, sample);

        if (K4A_STREAM_RESULT_SUCCEEDED == result)
        {
            return true;
        }
        else if (K4A_STREAM_RESULT_EOF == result)
        {
            return false;
        }

        throw error("Failed to get next IMU sample!");
    }

    bool get_previous_imu_sample(k4a_imu_sample_t *sample)
    {
        k4a_stream_result_t result = k4a_playback_get_previous_imu_sample(m_handle, sample);

        if (K4A_STREAM_RESULT_SUCCEEDED == result)
        {
            return true;
        }
        else if (K4A_STREAM_RESULT_EOF == result)
        {
            return false;
        }

        throw error("Failed to get previous IMU sample!");
    }

    void seek_timestamp(std::chrono::microseconds offset, k4a_playback_seek_origin_t origin)
    {
        k4a_result_t result = k4a_playback_seek_timestamp(m_handle, offset.count(), origin);

        if (K4A_RESULT_SUCCEEDED != result)
        {
            throw error("Failed to seek recording!");
        }
    }

    std::chrono::microseconds get_recording_length() const noexcept
    {
        return std::chrono::microseconds(k4a_playback_get_recording_length_usec(m_handle));
    }

    void set_color_conversion(k4a_image_format_t format)
    {
        k4a_result_t result = k4a_playback_set_color_conversion(m_handle, format);

        if (K4A_RESULT_SUCCEEDED != result)
        {
            throw error("Failed to set color conversion!");
        }
    }

    bool get_next_data_block(const char *track, data_block *block)
    {
        k4a_playback_data_block_t block_handle;
        k4a_stream_result_t result = k4a_playback_get_next_data_block(m_handle, track, &block_handle);

        if (K4A_STREAM_RESULT_SUCCEEDED == result)
        {
            *block = data_block(block_handle);
            return true;
        }
        else if (K4A_STREAM_RESULT_EOF == result)
        {
            return false;
        }

        throw error("Failed to get next data block!");
    }

    bool get_previous_data_block(const char *track, data_block *block)
    {
        k4a_playback_data_block_t block_handle;
        k4a_stream_result_t result = k4a_playback_get_previous_data_block(m_handle, track, &block_handle);

        if (K4A_STREAM_RESULT_SUCCEEDED == result)
        {
            *block = data_block(block_handle);
            return true;
        }
        else if (K4A_STREAM_RESULT_EOF == result)
        {
            return false;
        }

        throw error("Failed to get previous data block!");
    }

    bool get_attachment(const char *attachment, std::vector<uint8_t> *data)
    {
        size_t data_size = 0;
        k4a_buffer_result_t result = k4a_playback_get_attachment(m_handle, attachment, nullptr, &data_size);
        if (result == K4A_BUFFER_RESULT_TOO_SMALL)
        {
            data->resize(data_size);
            result = k4a_playback_get_attachment(m_handle, attachment, &(*data)[0], &data_size);
            if (result != K4A_BUFFER_RESULT_SUCCEEDED)
            {
                throw error("Failed to read attachment!");
            }
            return true;
        }
        return false;
    }

    static playback open(const char *path)
    {
        k4a_playback_t handle = nullptr;
        k4a_result_t result = k4a_playback_open(path, &handle);

        if (K4A_RESULT_SUCCEEDED != result)
        {
            throw error("Failed to open recording!");
        }

        return playback(handle);
    }

private:
    k4a_playback_t m_handle;
};

} // namespace k4a

#endif