#pragma once
// Android 屏幕处理器：管理截屏子连接，接收 Java 侧 MediaCodec 输出的 H.264 NALU，
// 按协议封装后通过子连接发送给服务端。
#include "common/commands.h"
#include "client/IOCPClient.h"
#include <atomic>
#include <mutex>
#include <condition_variable>
#include <queue>
#include <vector>
#include <thread>
#include <cstring>
#include <cinttypes>
#include <android/log.h>

#define LOGI_SH(...) __android_log_print(ANDROID_LOG_INFO,  "YAMA_SCR", __VA_ARGS__)
#define LOGE_SH(...) __android_log_print(ANDROID_LOG_ERROR, "YAMA_SCR", __VA_ARGS__)

extern uint64_t g_myClientID;
// 定义在 main.cpp，供子连接的 OnReceive 调用
extern void DispatchControlEvent(uint32_t msgVal, uint64_t wParam, int32_t ptX, int32_t ptY);

// Linux/macOS 共用的 BITMAPINFOHEADER 布局（与 Windows 完全一致）
#pragma pack(push, 1)
struct BmpInfoHeader {
    uint32_t biSize;
    int32_t  biWidth;
    int32_t  biHeight;
    uint16_t biPlanes;
    uint16_t biBitCount;
    uint32_t biCompression;
    uint32_t biSizeImage;
    int32_t  biXPelsPerMeter;
    int32_t  biYPelsPerMeter;
    uint32_t biClrUsed;
    uint32_t biClrImportant;
};
#pragma pack(pop)

class AndroidScreenHandler : public IOCPManager {
public:
    AndroidScreenHandler(IOCPClient* client, int width, int height)
        : m_client(client), m_width(width), m_height(height),
          m_started(false), m_running(true), m_firstSent(false)
    {
        memset(&m_bmpHdr, 0, sizeof(m_bmpHdr));
        m_bmpHdr.biSize          = sizeof(BmpInfoHeader);
        m_bmpHdr.biWidth         = width;
        m_bmpHdr.biHeight        = height;
        m_bmpHdr.biPlanes        = 1;
        m_bmpHdr.biBitCount      = 32;
        m_bmpHdr.biCompression   = 0;
        m_bmpHdr.biSizeImage     = (uint32_t)(width * height * 4);
        // 1 = top-down H264（Android MediaCodec 标准编码顺序）；
        // 服务端据此用负步长写 DIB，抵消 GDI bottom-up 翻转
        m_bmpHdr.biClrImportant  = 1;

        m_sendThread = std::thread(&AndroidScreenHandler::SendLoop, this);
    }

    ~AndroidScreenHandler() {
        m_running = false;
        m_cond.notify_all();
        if (m_sendThread.joinable())
            m_sendThread.join();
    }

    // 子连接建立后立即调用，告知服务端屏幕尺寸和编码格式
    void SendBitmapInfo() {
        const uint32_t total = 1 + sizeof(BmpInfoHeader) + 2 * sizeof(uint64_t) + sizeof(ScreenSettings);
        std::vector<uint8_t> buf(total, 0);

        buf[0] = TOKEN_BITMAPINFO;
        memcpy(&buf[1], &m_bmpHdr, sizeof(BmpInfoHeader));

        uint64_t clientID = g_myClientID;
        uint64_t zero     = 0;
        size_t off = 1 + sizeof(BmpInfoHeader);
        memcpy(&buf[off],     &clientID, 8);
        memcpy(&buf[off + 8], &zero,     8);

        ScreenSettings ss = {};
        ss.MaxFPS       = 10;
        ss.ScreenWidth  = m_width;
        ss.ScreenHeight = m_height;
        ss.QualityLevel = QUALITY_GOOD;   // 告知服务端使用 H264 解码器
        ss.ScreenType   = USING_VIRTUAL;  // 虚拟显示
        memcpy(&buf[off + 16], &ss, sizeof(ss));

        m_client->Send2Server((char*)buf.data(), total);
        LOGI_SH("SendBitmapInfo %dx%d clientID=%" PRIu64, m_width, m_height, clientID);

        // H.264 是推送模式，不需要等 COMMAND_NEXT 才开始发帧。
        // 服务端在 auth 通过瞬间发 COMMAND_NEXT，此时 setManagerCallBack 尚未注册，
        // 消息被 WorkThread 丢弃，m_started 永远 false 导致帧全部卡在队列里。
        m_started = true;
        m_cond.notify_all();
    }

    // 由 JNI 线程调用，投递 MediaCodec 输出的 NALU
    void OnFrameData(const uint8_t* data, uint32_t size, bool isKeyframe) {
        if (!m_running || size == 0) return;
        {
            std::lock_guard<std::mutex> lk(m_mutex);
            while (m_queue.size() >= 6) m_queue.pop();
            m_queue.push({std::vector<uint8_t>(data, data + size), isKeyframe});
        }
        m_cond.notify_one();
    }

    virtual VOID OnReceive(PBYTE data, ULONG size) override {
        if (!size) return;
        switch (data[0]) {
        case COMMAND_NEXT:
            // 推送模式下已在 SendBitmapInfo 里开始推流；此处保留兼容，无副作用。
            LOGI_SH("COMMAND_NEXT received");
            m_started = true;
            m_cond.notify_all();
            break;
        case CMD_QUALITY_LEVEL:
            if (size >= 2) LOGI_SH("QualityLevel=%d", (int)(int8_t)data[1]);
            break;
        case COMMAND_SCREEN_CONTROL: {
            // 服务端通过子连接下发鼠标/键盘控制包（MSG64 固定 48 字节）
            // 坐标从 lParam（offset 24）读取，与 Windows/Linux/macOS 客户端一致：
            // lParam = MAKELPARAM(enc_x, enc_y)，低16位=x，高16位=y。
            // 不读 pt.x/pt.y (offset 40/44)：64-bit MSG 的 pt 在 offset 36，
            // 与 MSG64 offset 40 不同，直接读会得到错误坐标。
            if ((ULONG)size < 1 + 48u) break;
            const uint8_t* p = data + 1;
            uint64_t msgVal = 0, wParam = 0, lParam = 0;
            memcpy(&msgVal,  p + 8,  8);
            memcpy(&wParam,  p + 16, 8);
            memcpy(&lParam,  p + 24, 8);
            int32_t ptX = (int32_t)(int16_t)(lParam & 0xFFFF);
            int32_t ptY = (int32_t)(int16_t)((lParam >> 16) & 0xFFFF);
            DispatchControlEvent((uint32_t)msgVal, wParam, ptX, ptY);
            break;
        }
        default:
            break;
        }
    }

private:
    struct Frame { std::vector<uint8_t> data; bool isKeyframe; };

    IOCPClient* m_client;
    int         m_width, m_height;
    std::atomic<bool> m_started;
    std::atomic<bool> m_running;
    std::atomic<bool> m_firstSent;
    std::atomic<bool> m_skipLogged{false};

    BmpInfoHeader m_bmpHdr;

    std::queue<Frame>       m_queue;
    std::mutex              m_mutex;
    std::condition_variable m_cond;
    std::thread             m_sendThread;

    void SendLoop() {
        while (m_running) {
            std::unique_lock<std::mutex> lk(m_mutex);
            m_cond.wait(lk, [&]{ return (!m_queue.empty() && m_started) || !m_running; });
            if (!m_running) break;

            Frame f = std::move(m_queue.front());
            m_queue.pop();
            lk.unlock();

            // 第一帧必须是关键帧
            if (!m_firstSent && !f.isKeyframe) {
                if (!m_skipLogged.exchange(true))
                    LOGI_SH("SendLoop: waiting for first IDR, skipping P-frames");
                continue;
            }

            if (!m_firstSent)
                LOGI_SH("SendH264 first IDR size=%u", (uint32_t)f.data.size());
            SendH264(f.data.data(), (uint32_t)f.data.size());
            m_firstSent = true;
        }
    }

    // 格式: [TOKEN_NEXTSCREEN:1][ALGORITHM_H264:1][cursorX:4][cursorY:4][cursorType:1][NALU:N]
    void SendH264(const uint8_t* nalu, uint32_t naluSize) {
        const uint32_t hdrSize = 1 + 1 + 4 + 4 + 1;
        std::vector<uint8_t> pkt(hdrSize + naluSize);

        pkt[0] = TOKEN_NEXTSCREEN;
        pkt[1] = ALGORITHM_H264;
        // cursor: (0,0), type: IDC_ARROW=1
        memset(&pkt[2], 0, 8);
        pkt[10] = 1;
        memcpy(&pkt[hdrSize], nalu, naluSize);

        m_client->Send2Server((char*)pkt.data(), pkt.size());
    }
};