SimpleRemoter/macos/ScreenHandler.h

#pragma once

#import <CoreGraphics/CoreGraphics.h>
#import <dispatch/dispatch.h>
#import <IOKit/pwr_mgt/IOPMLib.h>
#import <IOSurface/IOSurface.h>
#import "../client/IOCPClient.h"
#import "../common/commands.h"  // QualityLevel, QualityProfile, ALGORITHM_*
#include <vector>
#include <atomic>
#include <mutex>
#include <thread>
#include <cstdint>
#include <memory>
#include <condition_variable>

// Forward declarations
class IOCPClient;
class H264Encoder;
class InputHandler;

// macOS BITMAPINFOHEADER (compatible with Windows)
#pragma pack(push, 1)
struct BITMAPINFOHEADER_MAC {
    uint32_t biSize;            // 40
    int32_t  biWidth;
    int32_t  biHeight;
    uint16_t biPlanes;          // 1
    uint16_t biBitCount;        // 32
    uint32_t biCompression;     // 0 (BI_RGB)
    uint32_t biSizeImage;
    int32_t  biXPelsPerMeter;   // 0
    int32_t  biYPelsPerMeter;   // 0
    uint32_t biClrUsed;         // 0
    uint32_t biClrImportant;    // 0
};
#pragma pack(pop)

// Algorithm constants from commands.h: ALGORITHM_GRAY, ALGORITHM_DIFF, ALGORITHM_H264, ALGORITHM_RGB565

class ScreenHandler : public IOCPManager {
public:
    ScreenHandler(IOCPClient* client);
    ~ScreenHandler();

    // Initialize screen capture (returns false if permission denied)
    bool init();

    // Start/stop capture loop
    void start(IOCPClient* client, uint64_t clientID);
    void stop();

    // Check if running
    bool isRunning() const { return m_running; }

    // Get screen dimensions
    int getWidth() const { return m_width; }
    int getHeight() const { return m_height; }

    // Send bitmap info to server (called after connection)
    void sendBitmapInfo();

    // Handle received commands
    void OnReceive(uint8_t* data, ULONG size);

    // Apply quality level
    void applyQualityLevel(int8_t level, bool persist = false);

private:
    // Capture the screen (returns BGRA data, bottom-up)
    bool captureScreen(std::vector<uint8_t>& buffer);

    // Send first full screen frame
    void sendFirstScreen();

    // Send differential frame
    void sendDiffFrame();

    // Send H264 encoded frame
    void sendH264Frame(bool keyframe);

    // Compare bitmaps and generate diff data
    uint32_t compareBitmap(const uint8_t* curr, const uint8_t* prev,
                           uint8_t* outBuf, uint32_t totalBytes, uint8_t algo);

    // Color conversion helpers
    void convertBGRAtoGray(const uint8_t* src, uint8_t* dst, uint32_t pixelCount);
    void convertBGRAtoRGB565(const uint8_t* src, uint16_t* dst, uint32_t pixelCount);

    // Capture loop thread function
    void captureLoop();

    // Get current time in milliseconds
    static uint64_t getTickMs();

    // Get current cursor position (in physical pixels)
    void getCursorPosition(int32_t& x, int32_t& y);

    // Get current cursor type index (matches Windows cursor indices)
    uint8_t getCursorTypeIndex();

private:
    IOCPClient* m_client;
    uint64_t m_clientID;

    std::atomic<bool> m_running;
    std::thread m_captureThread;
    std::mutex m_mutex;

    // Screen info
    int m_width;          // Physical pixel width (sent to server)
    int m_height;         // Physical pixel height (sent to server)
    int m_logicalWidth;   // Logical point width (for CGEvent)
    int m_logicalHeight;  // Logical point height (for CGEvent)
    double m_scaleFactor; // Retina scale factor (physical / logical)
    CGDirectDisplayID m_displayID;

    // Protocol
    BITMAPINFOHEADER_MAC m_bmpHeader;
    std::vector<uint8_t> m_prevFrame;
    std::vector<uint8_t> m_currFrame;
    std::vector<uint8_t> m_diffBuffer;
    std::vector<uint8_t> m_tempBuffer;  // 临时缓冲区，避免每帧分配

    // Quality settings
    std::atomic<uint8_t> m_algorithm;
    std::atomic<int> m_maxFPS;
    int8_t m_qualityLevel;

    // H264 encoder
    std::unique_ptr<H264Encoder> m_h264Encoder;
    int m_h264Bitrate;

    // Input handler for mouse/keyboard control
    std::unique_ptr<InputHandler> m_inputHandler;

    // Power management: prevent display sleep during remote desktop
    IOPMAssertionID m_displayAssertionID;

    // Cached color space (avoid per-frame creation)
    CGColorSpaceRef m_colorSpace;

    // CGDisplayStream (efficient continuous capture)
    CGDisplayStreamRef m_displayStream;
    dispatch_queue_t m_streamQueue;
    IOSurfaceRef m_latestSurface;
    std::mutex m_surfaceMutex;
    std::condition_variable m_surfaceCond;
    std::atomic<bool> m_hasNewFrame;

    // Initialize/cleanup display stream
    bool initDisplayStream();
    void cleanupDisplayStream();

    // Process frame from IOSurface (called by stream callback)
    void processIOSurface(IOSurfaceRef surface);

    // Capture from IOSurface to buffer (with vertical flip)
    bool captureFromIOSurface(IOSurfaceRef surface, std::vector<uint8_t>& buffer);
};