#pragma once

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

// 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;

    // 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;
};