ScreenCapture.cpp

#include "ScreenCapture.h"

#include "LogBuffer.h"
#include "XPLMUtilities.h"
#include "XPLMDisplay.h"

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <utility>
#include <cstring>

namespace vr_image_client
{

namespace
{
// Pi constant for converting between radians and degrees.
constexpr float kPi = 3.14159265358979323846f;
constexpr float kMinValidFov = 0.01f;
constexpr float kMaxValidFov = 360.0f;
}

// Construct the screen capture helper with a logging callback.
ScreenCapture::ScreenCapture(LogCallback log_callback)
    : log_callback_(std::move(log_callback))
{
}

// Look up field-of-view datarefs used when logging the first capture.
void ScreenCapture::InitializeDataRefs()
{
    // Prefer the modern horizontal FOV dataref.
    horizontal_fov_ref_ = XPLMFindDataRef("sim/graphics/view/horizontal_fov_deg");
    if (horizontal_fov_ref_ == nullptr)
    {
        // Fall back to the legacy horizontal FOV dataref.
        horizontal_fov_ref_ = XPLMFindDataRef("sim/graphics/view/field_of_view_deg");
    }
    // Prefer the modern vertical FOV dataref.
    vertical_fov_ref_ = XPLMFindDataRef("sim/graphics/view/vertical_fov_deg");
    if (vertical_fov_ref_ == nullptr)
    {
        // Fall back to the legacy vertical FOV dataref.
        vertical_fov_ref_ = XPLMFindDataRef("sim/graphics/view/field_of_view_vertical_deg");
    }
    if (horizontal_fov_ref_ == nullptr && vertical_fov_ref_ == nullptr)
    {
        // Warn once if neither FOV dataref is present.
        WriteToPluginLog("[VR_for_Mac] Warning: camera FOV datarefs unavailable; capture logs will omit FOV values.\n");
    }
}

// Reset transient state when starting a new capture session.
void ScreenCapture::ResetSession()
{
    logged_first_capture_info_ = false;
    pbo_index_ = 0;
    map_index_ = 0;
    ready_pbo_count_ = 0;
}

void ScreenCapture::SetScreenSize(int width, int height)
{
    if (width > 0 && height > 0)
    {
        cached_screen_width_ = width;
        cached_screen_height_ = height;
    }
    else
    {
        cached_screen_width_ = 0;
        cached_screen_height_ = 0;
    }
}

void ScreenCapture::SetCaptureWidthRatio(float ratio)
{
    if (ratio <= 0.0f)
    {
        capture_width_ratio_ = 1.0f;
        return;
    }
    capture_width_ratio_ = std::min(1.0f, ratio);
}

bool ScreenCapture::GetCachedScreenSize(int &width, int &height) const
{
    return TryGetCachedScreenSize(width, height);
}

// Release OpenGL resources when shutting down.
void ScreenCapture::ReleaseResources()
{
    if (!buffers_initialized_)
    {
        // Nothing to do if buffers were never allocated.
        return;
    }
    // Drop the pixel buffer objects if they were created.
    glBindBuffer(GL_PIXEL_PACK_BUFFER, 0);
    glDeleteBuffers(kPboCount, pbo_ids_);
    for (int i = 0; i < kPboCount; ++i)
    {
        pbo_ids_[i] = 0;
    }
    buffer_size_ = 0;
    buffer_width_ = 0;
    buffer_height_ = 0;
    pbo_index_ = 0;
    map_index_ = 0;
    ready_pbo_count_ = 0;
    buffers_initialized_ = false;
}

bool ScreenCapture::CaptureFrame(RawFrame &out_frame, bool &frame_ready)
{
    frame_ready = false;

    // Query the current framebuffer size.
    int screen_width = 0;
    int screen_height = 0;
    if (!ResolveScreenSize(screen_width, screen_height))
    {
        // Invalid dimensions mean the capture cannot proceed.
        return false;
    }

    // Narrow the readback to the central slice the VR headset can actually display.
    const bool crop_enabled = COROP_IMAGE_CAPTURE_TO_SPEED_UP && capture_width_ratio_ > 0.0f;
    const float ratio = crop_enabled ? capture_width_ratio_ : 1.0f;
    int capture_width = static_cast<int>(std::lround(static_cast<float>(screen_width) * ratio));
    const int max_valid_width = std::max(1, screen_width);
    if (capture_width < 1)
    {
        capture_width = 1;
    }
    else if (capture_width > max_valid_width)
    {
        capture_width = max_valid_width;
    }
    const int capture_x = std::max(0, (screen_width - capture_width) / 2);

    // Resize backing buffers if the screen size changed.
    if (!EnsureBuffers(capture_width, screen_height))
    {
        // Buffer preparation failed, so we cannot capture this frame.
        return false;
    }

    // Ensure tightly packed pixel alignment.
    glPixelStorei(GL_PACK_ALIGNMENT, 1);

    const int write_index = pbo_index_;

    // Issue an asynchronous read into the current pixel buffer.
    glBindBuffer(GL_PIXEL_PACK_BUFFER, pbo_ids_[write_index]);
    glBufferData(GL_PIXEL_PACK_BUFFER, static_cast<GLsizeiptr>(buffer_size_), nullptr, GL_STREAM_READ);
    glReadPixels(capture_x, 0, capture_width, screen_height, GL_RGB, GL_UNSIGNED_BYTE, nullptr);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, 0);

    pbo_index_ = (pbo_index_ + 1) % kPboCount;
    if (ready_pbo_count_ < kPboCount)
    {
        ++ready_pbo_count_;
    }

    if (ready_pbo_count_ <= 1)
    {
        // Need at least one finished frame before mapping.
        return true;
    }

    const int map_target = map_index_;

    // Map the previously completed frame.
    glBindBuffer(GL_PIXEL_PACK_BUFFER, pbo_ids_[map_target]);
    unsigned char *mapped = static_cast<unsigned char *>(glMapBuffer(GL_PIXEL_PACK_BUFFER, GL_READ_ONLY));
    if (mapped == nullptr)
    {
        // Bail out if the GPU buffer failed to map to CPU memory.
        glBindBuffer(GL_PIXEL_PACK_BUFFER, 0);
        ready_pbo_count_ = 0;
        map_index_ = pbo_index_;
        return false;
    }

    // Copy the mapped pixel buffer as-is; receivers are responsible for any flipping.
    out_frame.width = capture_width;
    out_frame.height = screen_height;
    const std::size_t row_bytes =
        static_cast<std::size_t>(capture_width) * static_cast<std::size_t>(ScreenCapture::kBytesPerPixel);
    const std::size_t total_bytes = row_bytes * static_cast<std::size_t>(screen_height);
    out_frame.rgb.resize(total_bytes);
    std::memcpy(out_frame.rgb.data(), mapped, total_bytes);

    if (!logged_first_capture_info_)
    {
        // Log metadata only once per session.
        LogFirstCaptureInfo(screen_width, screen_height, capture_width);
        logged_first_capture_info_ = true;
    }

    // Unmap and release the pixel buffer to reuse it on the next frame.
    glUnmapBuffer(GL_PIXEL_PACK_BUFFER);
    glBindBuffer(GL_PIXEL_PACK_BUFFER, 0);

    map_index_ = (map_index_ + 1) % kPboCount;
    if (ready_pbo_count_ > 0)
    {
        --ready_pbo_count_;
    }

    frame_ready = true;
    return true;
}

bool ScreenCapture::QueryFieldOfView(float &horizontal_fov, float &vertical_fov) const
{
    bool has_horizontal = false;
    bool has_vertical = false;

    if (horizontal_fov_ref_ != nullptr)
    {
        const float value = XPLMGetDataf(horizontal_fov_ref_);
        if (value > kMinValidFov && value < kMaxValidFov)
        {
            horizontal_fov = value;
            has_horizontal = true;
        }
    }
    if (vertical_fov_ref_ != nullptr)
    {
        const float value = XPLMGetDataf(vertical_fov_ref_);
        if (value > kMinValidFov && value < kMaxValidFov)
        {
            vertical_fov = value;
            has_vertical = true;
        }
    }

    if (!has_horizontal && !has_vertical)
    {
        return false;
    }

    int width = 0;
    int height = 0;
    if (!has_horizontal || !has_vertical)
    {
        if (!TryGetCachedScreenSize(width, height))
        {
            XPLMGetScreenSize(&width, &height);
        }
    }

    if (!has_vertical && has_horizontal && width > 0 && height > 0)
    {
        const float aspect = static_cast<float>(width) / static_cast<float>(height);
        if (aspect > 0.0f)
        {
            const float horizontal_rad = horizontal_fov * (kPi / 180.0f);
            const float vertical_rad = 2.0f * std::atan(std::tan(horizontal_rad / 2.0f) / aspect);
            const float vertical_deg = vertical_rad * (180.0f / kPi);
            if (vertical_deg > kMinValidFov && vertical_deg < kMaxValidFov)
            {
                vertical_fov = vertical_deg;
                has_vertical = true;
            }
        }
    }
    else if (!has_horizontal && has_vertical && width > 0 && height > 0)
    {
        const float aspect = static_cast<float>(width) / static_cast<float>(height);
        if (aspect > 0.0f)
        {
            const float vertical_rad = vertical_fov * (kPi / 180.0f);
            const float horizontal_rad = 2.0f * std::atan(std::tan(vertical_rad / 2.0f) * aspect);
            const float horizontal_deg = horizontal_rad * (180.0f / kPi);
            if (horizontal_deg > kMinValidFov && horizontal_deg < kMaxValidFov)
            {
                horizontal_fov = horizontal_deg;
                has_horizontal = true;
            }
        }
    }

    return has_horizontal && has_vertical;
}

void ScreenCapture::LogFirstCaptureInfo(int screen_width, int screen_height, int captured_width)
{
    if (!log_callback_)
    {
        // Without a callback there is nowhere to log.
        return;
    }

    // Attempt to fetch horizontal and vertical FOV values, if available.
    float horizontal_fov = 0.0f;
    float vertical_fov = 0.0f;
    bool has_horizontal = false;
    bool has_vertical = false;

    if (horizontal_fov_ref_ != nullptr)
    {
        // Retrieve the horizontal FOV when the dataref exists.
        const float value = XPLMGetDataf(horizontal_fov_ref_);
        if (value > kMinValidFov && value < kMaxValidFov)
        {
            horizontal_fov = value;
            has_horizontal = true;
        }
    }
    if (vertical_fov_ref_ != nullptr)
    {
        // Retrieve the vertical FOV when the dataref exists.
        const float value = XPLMGetDataf(vertical_fov_ref_);
        if (value > kMinValidFov && value < kMaxValidFov)
        {
            vertical_fov = value;
            has_vertical = true;
        }
    }

    if (!has_vertical && has_horizontal && screen_width > 0 && screen_height > 0)
    {
        // Use the aspect ratio to derive the missing vertical FOV.
        const float aspect = static_cast<float>(screen_width) / static_cast<float>(screen_height);
        if (aspect > 0.0f)
        {
            // Infer vertical FOV from horizontal FOV and aspect ratio.
            const float horizontal_rad = horizontal_fov * (kPi / 180.0f);
            const float vertical_rad = 2.0f * std::atan(std::tan(horizontal_rad / 2.0f) / aspect);
            vertical_fov = vertical_rad * (180.0f / kPi);
            has_vertical = true;
        }
    }
    else if (!has_horizontal && has_vertical && screen_width > 0 && screen_height > 0)
    {
        // Use the aspect ratio to derive the missing horizontal FOV.
        const float aspect = static_cast<float>(screen_width) / static_cast<float>(screen_height);
        if (aspect > 0.0f)
        {
            // Infer horizontal FOV from vertical FOV and aspect ratio.
            const float vertical_rad = vertical_fov * (kPi / 180.0f);
            const float horizontal_rad = 2.0f * std::atan(std::tan(vertical_rad / 2.0f) * aspect);
            horizontal_fov = horizontal_rad * (180.0f / kPi);
            has_horizontal = true;
        }
    }

    char buffer[256];
    if (has_horizontal && has_vertical)
    {
        // Report both FOV axes when both were retrieved.
        std::snprintf(
            buffer,
            sizeof(buffer),
            "[VR_for_Mac] First capture: screen_width=%d capture_width=%d height=%d | FOV horiz=%.2f deg vert=%.2f deg",
            screen_width,
            captured_width,
            screen_height,
            horizontal_fov,
            vertical_fov);
    }
    else if (has_horizontal)
    {
        // Report only the horizontal FOV when vertical is missing.
        std::snprintf(
            buffer,
            sizeof(buffer),
            "[VR_for_Mac] First capture: screen_width=%d capture_width=%d height=%d | Horizontal FOV=%.2f deg",
            screen_width,
            captured_width,
            screen_height,
            horizontal_fov);
    }
    else if (has_vertical)
    {
        // Report only the vertical FOV when horizontal is missing.
        std::snprintf(
            buffer,
            sizeof(buffer),
            "[VR_for_Mac] First capture: screen_width=%d capture_width=%d height=%d | Vertical FOV=%.2f deg",
            screen_width,
            captured_width,
            screen_height,
            vertical_fov);
    }
    else
    {
        // Indicate when no FOV data could be queried.
        std::snprintf(
            buffer,
            sizeof(buffer),
            "[VR_for_Mac] First capture: screen_width=%d capture_width=%d height=%d | FOV unavailable",
            screen_width,
            captured_width,
            screen_height);
    }

    // Send the composed message back through the plugin logging callback.
    log_callback_(buffer);
}

bool ScreenCapture::ResolveScreenSize(int &width, int &height)
{
    if (cached_screen_width_ > 0 && cached_screen_height_ > 0)
    {
        width = cached_screen_width_;
        height = cached_screen_height_;
        return true;
    }
    XPLMGetScreenSize(&width, &height);
    if (width > 0 && height > 0)
    {
        cached_screen_width_ = width;
        cached_screen_height_ = height;
        return true;
    }
    return false;
}

bool ScreenCapture::TryGetCachedScreenSize(int &width, int &height) const
{
    if (cached_screen_width_ > 0 && cached_screen_height_ > 0)
    {
        width = cached_screen_width_;
        height = cached_screen_height_;
        return true;
    }
    return false;
}

// Allocate or resize the pixel buffer objects that back capture requests.
bool ScreenCapture::EnsureBuffers(int width, int height)
{
    const std::size_t required_size =
        static_cast<std::size_t>(width) * static_cast<std::size_t>(height) * ScreenCapture::kBytesPerPixel;
    if (required_size == 0u)
    {
        // Skip capture setup when the framebuffer size is zero.
        return false;
    }

    if (!buffers_initialized_)
    {
        // Generate multiple buffers so reads and uploads can overlap.
        glGenBuffers(ScreenCapture::kPboCount, pbo_ids_);
        bool allocation_failed = false;
        for (int i = 0; i < ScreenCapture::kPboCount; ++i)
        {
            if (pbo_ids_[i] == 0)
            {
                allocation_failed = true;
                break;
            }
        }
        if (allocation_failed)
        {
            // Clean up if PBO generation failed.
            glDeleteBuffers(ScreenCapture::kPboCount, pbo_ids_);
            for (int i = 0; i < ScreenCapture::kPboCount; ++i)
            {
                pbo_ids_[i] = 0;
            }
            return false;
        }
        buffer_width_ = width;
        buffer_height_ = height;
        buffer_size_ = required_size;
        for (int i = 0; i < ScreenCapture::kPboCount; ++i)
        {
            // Allocate storage for each PBO according to the current screen size.
            glBindBuffer(GL_PIXEL_PACK_BUFFER, pbo_ids_[i]);
            glBufferData(GL_PIXEL_PACK_BUFFER, static_cast<GLsizeiptr>(buffer_size_), nullptr, GL_STREAM_READ);
        }
        glBindBuffer(GL_PIXEL_PACK_BUFFER, 0);
        pbo_index_ = 0;
        map_index_ = 0;
        ready_pbo_count_ = 0;
        buffers_initialized_ = true;
    }
    else if (width != buffer_width_ || height != buffer_height_)
    {
        // Update cached dimensions before reallocating the buffers.
        buffer_width_ = width;
        buffer_height_ = height;
        buffer_size_ = required_size;
        for (int i = 0; i < ScreenCapture::kPboCount; ++i)
        {
            // Reallocate existing PBOs when the dimensions change.
            glBindBuffer(GL_PIXEL_PACK_BUFFER, pbo_ids_[i]);
            glBufferData(GL_PIXEL_PACK_BUFFER, static_cast<GLsizeiptr>(buffer_size_), nullptr, GL_STREAM_READ);
        }
        glBindBuffer(GL_PIXEL_PACK_BUFFER, 0);
        pbo_index_ = 0;
        map_index_ = 0;
        ready_pbo_count_ = 0;
    }
    return true;
}

} // namespace vr_image_client