scar-chat7/client/media/screen_capture.cpp

#include "screen_capture.h"
#include <iostream>
#include <cstring>
#include <random>
#include <sstream>
#include <iomanip>
#include <unistd.h>

#ifdef __linux__
#include <sdbus-c++/sdbus-c++.h>
#include <pipewire/pipewire.h>
#include <spa/param/video/format-utils.h>
#include <spa/param/video/type-info.h>
#include <spa/param/format-utils.h>
#include <spa/param/video/format.h>
#include <spa/utils/hook.h>
#include <spa/pod/builder.h>
#include <QImage>
#endif

namespace scar {

ScreenCapture::ScreenCapture()
    : capturing_(false),
      pw_loop_(nullptr),
      pw_stream_(nullptr),
      pw_context_(nullptr),
      frame_width_(0),
      frame_height_(0) {
    backend_ = detectBestBackend();
#ifdef __linux__
    memset(&stream_listener_, 0, sizeof(stream_listener_));
#endif
}

ScreenCapture::~ScreenCapture() {
    stop();
    cleanupPipewire();
    cleanupPortalConnection();
}

ScreenCaptureBackend ScreenCapture::detectBestBackend() {
#ifdef _WIN32
    return ScreenCaptureBackend::FFMPEG_WINDOWS;
#else
    // Check for Wayland/Hyprland
    const char* waylandDisplay = std::getenv("WAYLAND_DISPLAY");
    const char* hyprlandInstance = std::getenv("HYPRLAND_INSTANCE_SIGNATURE");
    
    if (waylandDisplay || hyprlandInstance) {
        std::cout << "Detected Wayland/Hyprland environment" << std::endl;
        return ScreenCaptureBackend::PORTAL_PIPEWIRE;
    }
    
    // Fallback to X11
    std::cout << "Detected X11 environment" << std::endl;
    return ScreenCaptureBackend::FFMPEG_X11;
#endif
}

bool ScreenCapture::start() {
    return start(backend_);
}

bool ScreenCapture::start(ScreenCaptureBackend backend) {
    if (capturing_) {
        return true;
    }
    
    backend_ = backend;
    
    switch (backend_) {
        case ScreenCaptureBackend::PORTAL_PIPEWIRE:
            return startPortalCapture();
        case ScreenCaptureBackend::FFMPEG_X11:
            return startX11Capture();
        case ScreenCaptureBackend::FFMPEG_WAYLAND:
            return startWaylandCapture();
        case ScreenCaptureBackend::FFMPEG_WINDOWS:
            return startWindowsCapture();
        default:
            std::cerr << "Unknown backend" << std::endl;
            return false;
    }
}

void ScreenCapture::stop() {
    if (!capturing_) {
        return;
    }
    
    std::cout << "Stopping screen capture..." << std::endl;
    
    cleanupPipewire();
    cleanupPortalConnection();
    
    capturing_ = false;
}

void ScreenCapture::setFrameCallback(FrameCallback callback) {
    frameCallback_ = std::move(callback);
}

bool ScreenCapture::startX11Capture() {
    std::cout << "X11 screen capture not yet implemented" << std::endl;
    // TODO: Implement FFmpeg x11grab
    return false;
}

bool ScreenCapture::startWaylandCapture() {
    std::cout << "FFmpeg Wayland capture not yet implemented" << std::endl;
    // TODO: Implement FFmpeg Wayland
    return false;
}

bool ScreenCapture::startWindowsCapture() {
    std::cout << "Windows screen capture not yet implemented" << std::endl;
    // TODO: Implement FFmpeg GDI
    return false;
}

bool ScreenCapture::startPortalCapture() {
#ifndef __linux__
    std::cerr << "Portal capture only supported on Linux" << std::endl;
    return false;
#else
    std::cout << "Starting Portal + Pipewire screen capture..." << std::endl;
    
    // Step 1: Initialize portal connection
    if (!initPortalConnection()) {
        std::cerr << "Failed to connect to xdg-desktop-portal" << std::endl;
        return false;
    }
    
    // Step 2: Create portal session
    std::string session_handle = createPortalSession();
    if (session_handle.empty()) {
        std::cerr << "Failed to create portal session" << std::endl;
        cleanupPortalConnection();
        return false;
    }
    
    // Step 3: Select sources (monitors/windows)
    if (!selectPortalSources(session_handle)) {
        std::cerr << "Failed to select portal sources" << std::endl;
        cleanupPortalConnection();
        return false;
    }
    
    // Step 4: Start session and get Pipewire node ID
    uint32_t node_id = 0;
    if (!startPortalSession(session_handle, node_id)) {
        std::cerr << "Failed to start portal session" << std::endl;
        cleanupPortalConnection();
        return false;
    }
    
    std::cout << "Got Pipewire node ID: " << node_id << std::endl;
    
    // Step 5: Open Pipewire remote
    int pw_fd = openPipeWireRemote(session_handle);
    if (pw_fd < 0) {
        std::cerr << "Failed to open Pipewire remote" << std::endl;
        cleanupPortalConnection();
        return false;
    }
    
    // Step 6: Initialize Pipewire and connect to stream
    if (!initPipewire(pw_fd, node_id)) {
        std::cerr << "Failed to initialize Pipewire" << std::endl;
        cleanupPortalConnection();
        return false;
    }
    
    capturing_ = true;
    std::cout << "Screen capture started successfully" << std::endl;
    return true;
#endif
}

#ifdef __linux__
bool ScreenCapture::initPortalConnection() {
    try {
        // Create session bus connection
        portal_connection_ = sdbus::createSessionBusConnection();
        
        // Create proxy for ScreenCast portal interface
        // Using standard org.freedesktop.portal.Desktop interface
        // xdg-desktop-portal-hyprland implements this standard interface
        screencast_proxy_ = sdbus::createProxy(
            *portal_connection_,
            sdbus::ServiceName{"org.freedesktop.portal.Desktop"},
            sdbus::ObjectPath{"/org/freedesktop/portal/desktop"}
        );
        
        std::cout << "Connected to xdg-desktop-portal (Hyprland implementation)" << std::endl;
        return true;
    }
    catch (const sdbus::Error& e) {
        std::cerr << "Failed to connect to portal: " << e.what() << std::endl;
        return false;
    }
}

void ScreenCapture::cleanupPortalConnection() {
    screencast_proxy_.reset();
    portal_connection_.reset();
}

std::string ScreenCapture::createPortalSession() {
    try {
        // Generate unique handle for this request
        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_int_distribution<> dis(0, 999999);
        
        std::stringstream ss;
        ss << "/org/freedesktop/portal/desktop/request/" 
           << getpid() << "_" << dis(gen);
        std::string request_handle = ss.str();
        
        ss.str("");
        ss << "/org/freedesktop/portal/desktop/session/" 
           << getpid() << "_" << dis(gen);
        std::string session_handle = ss.str();
        
        std::cout << "Creating portal session..." << std::endl;
        std::cout << "  Request handle: " << request_handle << std::endl;
        std::cout << "  Session handle: " << session_handle << std::endl;
        
        // Build options dictionary
        std::map<std::string, sdbus::Variant> options;
        options["handle_token"] = sdbus::Variant{request_handle.substr(request_handle.rfind('/') + 1)};
        options["session_handle_token"] = sdbus::Variant{session_handle.substr(session_handle.rfind('/') + 1)};
        
        // Call CreateSession
        sdbus::ObjectPath response_path;
        screencast_proxy_->callMethod("CreateSession")
            .onInterface("org.freedesktop.portal.ScreenCast")
            .withArguments(options)
            .storeResultsTo(response_path);
        
        std::cout << "Session created: " << session_handle << std::endl;
        session_path_ = session_handle;
        return session_handle;
    }
    catch (const sdbus::Error& e) {
        std::cerr << "Failed to create portal session: " << e.what() << std::endl;
        return "";
    }
}

bool ScreenCapture::selectPortalSources(const std::string& session_handle) {
    try {
        // Generate unique request handle
        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_int_distribution<> dis(0, 999999);
        
        std::stringstream ss;
        ss << "/org/freedesktop/portal/desktop/request/" 
           << getpid() << "_" << dis(gen);
        std::string request_handle = ss.str();
        
        std::cout << "Selecting portal sources..." << std::endl;
        
        // Build options dictionary
        std::map<std::string, sdbus::Variant> options;
        options["handle_token"] = sdbus::Variant{request_handle.substr(request_handle.rfind('/') + 1)};
        options["types"] = sdbus::Variant{uint32_t(1 | 2)}; // MONITOR (1) | WINDOW (2)
        options["multiple"] = sdbus::Variant{false};
        options["cursor_mode"] = sdbus::Variant{uint32_t(2)}; // Embedded cursor
        
        // Call SelectSources
        sdbus::ObjectPath response_path;
        screencast_proxy_->callMethod("SelectSources")
            .onInterface("org.freedesktop.portal.ScreenCast")
            .withArguments(sdbus::ObjectPath{session_handle}, options)
            .storeResultsTo(response_path);
        
        std::cout << "Sources selected (user will see hyprland-share-picker)" << std::endl;
        return true;
    }
    catch (const sdbus::Error& e) {
        std::cerr << "Failed to select portal sources: " << e.what() << std::endl;
        return false;
    }
}

bool ScreenCapture::startPortalSession(const std::string& session_handle, uint32_t& node_id) {
    try {
        // Generate unique request handle
        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_int_distribution<> dis(0, 999999);
        
        std::stringstream ss;
        ss << "/org/freedesktop/portal/desktop/request/" 
           << getpid() << "_" << dis(gen);
        std::string request_handle = ss.str();
        
        std::cout << "Starting portal session..." << std::endl;
        
        // Build options dictionary
        std::map<std::string, sdbus::Variant> options;
        options["handle_token"] = sdbus::Variant{request_handle.substr(request_handle.rfind('/') + 1)};
        
        // Call Start
        sdbus::ObjectPath response_path;
        screencast_proxy_->callMethod("Start")
            .onInterface("org.freedesktop.portal.ScreenCast")
            .withArguments(sdbus::ObjectPath{session_handle}, std::string(""), options)
            .storeResultsTo(response_path);
        
        std::cout << "Portal session started" << std::endl;
        
        // Get the actual node_id from the session's Streams property
        node_id = getStreamsNodeId(session_handle);
        
        if (node_id == 0) {
            std::cerr << "Failed to get valid node_id from streams" << std::endl;
            return false;
        }
        
        return true;
    }
    catch (const sdbus::Error& e) {
        std::cerr << "Failed to start portal session: " << e.what() << std::endl;
        return false;
    }
}

int ScreenCapture::openPipeWireRemote(const std::string& session_handle) {
    try {
        std::cout << "Opening PipeWire remote..." << std::endl;
        
        // Build options dictionary (empty for now)
        std::map<std::string, sdbus::Variant> options;
        
        // Call OpenPipeWireRemote
        sdbus::UnixFd pipewire_fd;
        screencast_proxy_->callMethod("OpenPipeWireRemote")
            .onInterface("org.freedesktop.portal.ScreenCast")
            .withArguments(sdbus::ObjectPath{session_handle}, options)
            .storeResultsTo(pipewire_fd);
        
        // Extract the actual file descriptor from UnixFd
        int fd = dup(pipewire_fd.get());
        
        if (fd < 0) {
            std::cerr << "Failed to duplicate PipeWire file descriptor" << std::endl;
            return -1;
        }
        
        std::cout << "PipeWire remote opened, fd: " << fd << std::endl;
        return fd;
    }
    catch (const sdbus::Error& e) {
        std::cerr << "Failed to open PipeWire remote: " << e.what() << std::endl;
        return -1;
    }
}

uint32_t ScreenCapture::getStreamsNodeId(const std::string& session_handle) {
    try {
        std::cout << "Getting streams from session..." << std::endl;
        
        // Create a proxy for the session object
        auto session_proxy = sdbus::createProxy(
            *portal_connection_,
            sdbus::ServiceName{"org.freedesktop.portal.Desktop"},
            sdbus::ObjectPath{session_handle}
        );
        
        // Get the Streams property from org.freedesktop.portal.Session interface
        auto streams_variant = session_proxy->getProperty("Streams")
            .onInterface("org.freedesktop.portal.Session");
        
        // Streams is a(ua{sv}) - array of (node_id, properties)
        auto streams = streams_variant.get<std::vector<sdbus::Struct<uint32_t, std::map<std::string, sdbus::Variant>>>>();
        
        if (streams.empty()) {
            std::cerr << "No streams available in session" << std::endl;
            return 0;
        }
        
        // Get the first stream's node_id
        uint32_t node_id = std::get<0>(streams[0]);
        auto& properties = std::get<1>(streams[0]);
        
        std::cout << "Found stream node_id: " << node_id << std::endl;
        
        // Print stream properties if available
        if (properties.count("size")) {
            auto size = properties["size"].get<sdbus::Struct<int32_t, int32_t>>();
            std::cout << "  Stream size: " << std::get<0>(size) << "x" << std::get<1>(size) << std::endl;
        }
        
        return node_id;
    }
    catch (const sdbus::Error& e) {
        std::cerr << "Failed to get streams: " << e.what() << std::endl;
        return 0;
    }
}

bool ScreenCapture::initPipewire(int fd, uint32_t node_id) {
    pw_init(nullptr, nullptr);
    
    std::cout << "Initializing PipeWire with fd: " << fd << ", node_id: " << node_id << std::endl;
    
    pw_loop_ = pw_thread_loop_new("screen-capture", nullptr);
    if (!pw_loop_) {
        std::cerr << "Failed to create Pipewire loop" << std::endl;
        return false;
    }
    
    pw_context_ = pw_context_new(
        pw_thread_loop_get_loop(pw_loop_),
        nullptr, 0
    );
    if (!pw_context_) {
        std::cerr << "Failed to create Pipewire context" << std::endl;
        pw_thread_loop_destroy(pw_loop_);
        pw_loop_ = nullptr;
        return false;
    }
    
    // Connect to the PipeWire remote using the file descriptor
    pw_core* core = pw_context_connect_fd(
        pw_context_,
        fd,
        nullptr, 0
    );
    if (!core) {
        std::cerr << "Failed to connect to PipeWire remote" << std::endl;
        cleanupPipewire();
        return false;
    }
    
    // Create stream properties
    pw_properties* props = pw_properties_new(
        PW_KEY_MEDIA_TYPE, "Video",
        PW_KEY_MEDIA_CATEGORY, "Capture",
        PW_KEY_MEDIA_ROLE, "Screen",
        nullptr
    );
    
    // Create the stream
    pw_stream_ = pw_stream_new(
        core,
        "screen-capture-stream",
        props
    );
    
    if (!pw_stream_) {
        std::cerr << "Failed to create PipeWire stream" << std::endl;
        pw_core_disconnect(core);
        cleanupPipewire();
        return false;
    }
    
    // Set up stream events
    static const pw_stream_events stream_events = {
        PW_VERSION_STREAM_EVENTS,
        nullptr,  // destroy
        nullptr,  // state_changed
        nullptr,  // control_info
        nullptr,  // io_changed
        onStreamParamChanged,  // param_changed
        nullptr,  // add_buffer
        nullptr,  // remove_buffer
        onStreamProcess,  // process
        nullptr,  // drained
        nullptr,  // command
        nullptr,  // trigger_done
    };
    
    pw_stream_add_listener(
        pw_stream_,
        &stream_listener_,
        &stream_events,
        this  // userdata pointer
    );
    
    // Connect to the node
    uint8_t buffer[1024];
    spa_pod_builder pod_builder = SPA_POD_BUILDER_INIT(buffer, sizeof(buffer));
    
    // Request any video format (we'll adapt to what we get)
    const struct spa_pod* params[1];
    params[0] = reinterpret_cast<const struct spa_pod*>(spa_pod_builder_add_object(
        &pod_builder,
        SPA_TYPE_OBJECT_Format, SPA_PARAM_EnumFormat,
        SPA_FORMAT_mediaType, SPA_POD_Id(SPA_MEDIA_TYPE_video),
        SPA_FORMAT_mediaSubtype, SPA_POD_Id(SPA_MEDIA_SUBTYPE_raw)
    ));
    
    int result = pw_stream_connect(
        pw_stream_,
        PW_DIRECTION_INPUT,
        node_id,
        (pw_stream_flags)(PW_STREAM_FLAG_AUTOCONNECT | PW_STREAM_FLAG_MAP_BUFFERS),
        params, 1
    );
    
    if (result < 0) {
        std::cerr << "Failed to connect stream to node " << node_id << ": " << strerror(-result) << std::endl;
        pw_core_disconnect(core);
        cleanupPipewire();
        return false;
    }
    
    // Start the thread loop
    if (pw_thread_loop_start(pw_loop_) < 0) {
        std::cerr << "Failed to start PipeWire thread loop" << std::endl;
        pw_core_disconnect(core);
        cleanupPipewire();
        return false;
    }
    
    std::cout << "PipeWire stream connected and running" << std::endl;
    return true;
}

void ScreenCapture::cleanupPipewire() {
    if (pw_loop_) {
        pw_thread_loop_stop(pw_loop_);
    }
    
    if (pw_stream_) {
        pw_stream_destroy(pw_stream_);
        pw_stream_ = nullptr;
    }
    
    if (pw_context_) {
        pw_context_destroy(pw_context_);
        pw_context_ = nullptr;
    }
    
    if (pw_loop_) {
        pw_thread_loop_destroy(pw_loop_);
        pw_loop_ = nullptr;
    }
    
    pw_deinit();
}

void ScreenCapture::onStreamProcess(void* userdata) {
    auto* self = static_cast<ScreenCapture*>(userdata);
    
    pw_buffer* buffer = pw_stream_dequeue_buffer(self->pw_stream_);
    if (!buffer) {
        std::cerr << "No buffer available" << std::endl;
        return;
    }
    
    spa_buffer* spa_buf = buffer->buffer;
    if (!spa_buf || !spa_buf->datas || !spa_buf->datas[0].data) {
        pw_stream_queue_buffer(self->pw_stream_, buffer);
        return;
    }
    
    // Get frame data
    spa_data& data = spa_buf->datas[0];
    uint8_t* frame_data = static_cast<uint8_t*>(data.data);
    uint32_t size = data.chunk->size;
    
    // Invoke the callback if set
    if (self->frameCallback_) {
        // Convert to vector for callback
        std::vector<uint8_t> frame_vec(frame_data, frame_data + size);
        self->frameCallback_(frame_vec, self->frame_width_, self->frame_height_);
    }
    
    // Return buffer to stream
    pw_stream_queue_buffer(self->pw_stream_, buffer);
}

void ScreenCapture::onStreamParamChanged(void* userdata, uint32_t id, const struct spa_pod* param) {
    auto* self = static_cast<ScreenCapture*>(userdata);
    
    if (!param || id != SPA_PARAM_Format) {
        return;
    }
    
    // Parse video format
    spa_video_info_raw video_info;
    if (spa_format_video_raw_parse(param, &video_info) < 0) {
        std::cerr << "Failed to parse video format" << std::endl;
        return;
    }
    
    // Update our stored dimensions
    self->frame_width_ = video_info.size.width;
    self->frame_height_ = video_info.size.height;
    
    std::cout << "Stream format changed:" << std::endl;
    std::cout << "  Resolution: " << self->frame_width_ << "x" << self->frame_height_ << std::endl;
    std::cout << "  Format: " << video_info.format << std::endl;
    std::cout << "  Framerate: " << video_info.framerate.num << "/" << video_info.framerate.denom << std::endl;
}

#else
// Stub implementations for non-Linux platforms
bool ScreenCapture::initPortalConnection() { return false; }
void ScreenCapture::cleanupPortalConnection() {}
std::string ScreenCapture::createPortalSession() { return ""; }
bool ScreenCapture::selectPortalSources(const std::string&) { return false; }
bool ScreenCapture::startPortalSession(const std::string&, uint32_t&) { return false; }
int ScreenCapture::openPipeWireRemote(const std::string&) { return -1; }
uint32_t ScreenCapture::getStreamsNodeId(const std::string&) { return 0; }
bool ScreenCapture::initPipewire(int, uint32_t) { return false; }
void ScreenCapture::cleanupPipewire() {}
void ScreenCapture::onStreamProcess(void*) {}
void ScreenCapture::onStreamParamChanged(void*, uint32_t, const struct spa_pod*) {}
#endif

} // namespace scar