#pragma once
/**
 * X264Encoder.h - Linux H264 Encoder using libx264
 *
 * Features:
 *   - Dynamic library loading (dlopen/dlsym)
 *   - Automatic fallback if libx264 not available
 *   - Manual BGRA→I420 conversion (no libyuv dependency)
 *   - API compatible with Windows X264Encoder
 *
 * Requirements:
 *   - libx264 installed (apt install libx264-dev)
 *   - If not installed, H264 encoding is disabled
 */

#include <stdint.h>
#include <cstring>
#include <cstdlib>
#include <dlfcn.h>
#include <vector>

// Include x264 header for struct definitions
// The library is dynamically loaded at runtime
extern "C" {
#include "../compress/x264/x264.h"
}

// ============== X264Encoder Class ==============

class X264Encoder {
public:
    // Check if libx264 is available on this system
    static bool IsAvailable() {
        static int available = -1;
        if (available < 0) {
            void* handle = TryLoadLibrary();
            available = (handle != nullptr) ? 1 : 0;
            if (handle) {
                dlclose(handle);
                fprintf(stderr, ">>> X264Encoder: libx264 available\n");
            } else {
                fprintf(stderr, "*** X264Encoder: libx264 not found (%s)\n", dlerror());
            }
        }
        return available == 1;
    }

    X264Encoder()
        : m_x264Handle(nullptr)
        , m_encoder(nullptr)
        , m_picIn(nullptr)
        , m_picOut(nullptr)
        , m_width(0)
        , m_height(0)
    {
        memset(&m_param, 0, sizeof(m_param));
        clearFunctionPointers();
    }

    ~X264Encoder() {
        close();
    }

    bool open(int width, int height, int fps, int crf) {
        close();

        // Load library
        if (!loadLibrary()) {
            fprintf(stderr, "*** X264Encoder::open: loadLibrary failed\n");
            return false;
        }

        // Round to even dimensions (H264 requirement)
        m_width = width & ~1;
        m_height = height & ~1;

        // Initialize parameters
        if (fn_x264_param_default_preset) {
            fn_x264_param_default_preset(&m_param, "ultrafast", "zerolatency");
        }

        // Set encoder parameters
        m_param.i_width = m_width;
        m_param.i_height = m_height;
        m_param.i_log_level = X264_LOG_NONE;
        m_param.i_threads = 1;
        m_param.i_frame_total = 0;
        m_param.i_keyint_max = fps * 15;  // Keyframe every 15 seconds
        m_param.i_bframe = 0;              // No B-frames for low latency
        m_param.b_open_gop = 0;
        m_param.i_fps_num = fps;
        m_param.i_fps_den = 1;
        m_param.i_csp = X264_CSP_I420;

        // Rate control: CRF mode
        m_param.rc.i_rc_method = X264_RC_CRF;
        m_param.rc.f_rf_constant = (float)crf;

        // Apply baseline profile for compatibility
        if (fn_x264_param_apply_profile) {
            fn_x264_param_apply_profile(&m_param, "baseline");
        }

        // Allocate pictures
        m_picIn = (x264_picture_t*)calloc(1, sizeof(x264_picture_t));
        m_picOut = (x264_picture_t*)calloc(1, sizeof(x264_picture_t));
        if (!m_picIn || !m_picOut) {
            close();
            return false;
        }

        // Initialize input picture
        if (fn_x264_picture_init) {
            fn_x264_picture_init(m_picIn);
        }

        // Allocate picture buffer
        if (fn_x264_picture_alloc) {
            if (fn_x264_picture_alloc(m_picIn, X264_CSP_I420, m_width, m_height) < 0) {
                close();
                return false;
            }
        }

        // Open encoder
        m_encoder = fn_x264_encoder_open(&m_param);
        if (!m_encoder) {
            close();
            return false;
        }

        return true;
    }

    void close() {
        if (m_encoder && fn_x264_encoder_close) {
            fn_x264_encoder_close(m_encoder);
            m_encoder = nullptr;
        }

        if (m_picIn) {
            if (fn_x264_picture_clean) {
                fn_x264_picture_clean(m_picIn);
            }
            free(m_picIn);
            m_picIn = nullptr;
        }

        if (m_picOut) {
            free(m_picOut);
            m_picOut = nullptr;
        }

        unloadLibrary();
        m_width = m_height = 0;
    }

    /**
     * Encode a frame
     * @param bgra      Input BGRA image data
     * @param bpp       Bits per pixel (24 or 32)
     * @param stride    Bytes per row
     * @param width     Image width
     * @param height    Image height
     * @param outData   Output: pointer to encoded H264 data
     * @param outSize   Output: size of encoded data
     * @param direction 1 = normal, -1 = vertical flip
     * @return 0 on success, negative on error
     */
    int encode(uint8_t* bgra, uint8_t bpp, uint32_t stride,
               uint32_t width, uint32_t height,
               uint8_t** outData, uint32_t* outSize,
               int direction = 1)
    {
        if (!m_encoder || !m_picIn || !fn_x264_encoder_encode) {
            return -1;
        }

        // Check dimensions match
        if ((int)(width & ~1) != m_width || (int)(height & ~1) != m_height) {
            return -2;
        }

        // Convert BGRA to I420 directly into x264 picture planes
        if (bpp == 32) {
            convertBGRAtoI420(bgra, stride, direction,
                              m_picIn->img.plane[0], m_picIn->img.i_stride[0],
                              m_picIn->img.plane[1], m_picIn->img.i_stride[1],
                              m_picIn->img.plane[2], m_picIn->img.i_stride[2]);
        } else if (bpp == 24) {
            convertRGB24toI420(bgra, stride, direction,
                               m_picIn->img.plane[0], m_picIn->img.i_stride[0],
                               m_picIn->img.plane[1], m_picIn->img.i_stride[1],
                               m_picIn->img.plane[2], m_picIn->img.i_stride[2]);
        } else {
            return -3;
        }

        // Encode
        x264_nal_t* pNal = nullptr;
        int iNal = 0;
        int encodeSize = fn_x264_encoder_encode(m_encoder, &pNal, &iNal, m_picIn, m_picOut);

        if (encodeSize < 0) {
            return -4;
        }

        if (encodeSize == 0 || !pNal) {
            *outData = nullptr;
            *outSize = 0;
            return 0;
        }

        *outData = pNal->p_payload;
        *outSize = encodeSize;
        return 0;
    }

private:
    // Library handle
    void* m_x264Handle;

    // Encoder state
    x264_t* m_encoder;
    x264_param_t m_param;
    x264_picture_t* m_picIn;
    x264_picture_t* m_picOut;
    int m_width, m_height;

    // x264 function pointers
    void  (*fn_x264_param_default_preset)(x264_param_t*, const char*, const char*);
    int   (*fn_x264_param_apply_profile)(x264_param_t*, const char*);
    x264_t* (*fn_x264_encoder_open)(x264_param_t*);
    void  (*fn_x264_encoder_close)(x264_t*);
    int   (*fn_x264_encoder_encode)(x264_t*, x264_nal_t**, int*, x264_picture_t*, x264_picture_t*);
    void  (*fn_x264_picture_init)(x264_picture_t*);
    int   (*fn_x264_picture_alloc)(x264_picture_t*, int, int, int);
    void  (*fn_x264_picture_clean)(x264_picture_t*);

    void clearFunctionPointers() {
        fn_x264_param_default_preset = nullptr;
        fn_x264_param_apply_profile = nullptr;
        fn_x264_encoder_open = nullptr;
        fn_x264_encoder_close = nullptr;
        fn_x264_encoder_encode = nullptr;
        fn_x264_picture_init = nullptr;
        fn_x264_picture_alloc = nullptr;
        fn_x264_picture_clean = nullptr;
    }

    static void* TryLoadLibrary() {
        // Try multiple library versions (newest first)
        const char* libNames[] = {
            "libx264.so",        // symlink (if exists)
            "libx264.so.164",    // Ubuntu 24, Debian 12+
            "libx264.so.163",
            "libx264.so.162",
            "libx264.so.161",
            "libx264.so.160",
            "libx264.so.159",
            "libx264.so.157",
            "libx264.so.155",    // Ubuntu 20
            "libx264.so.152",
            "libx264.so.148",    // older distros
            nullptr
        };

        for (int i = 0; libNames[i]; i++) {
            void* handle = dlopen(libNames[i], RTLD_LAZY);
            if (handle) return handle;
        }
        return nullptr;
    }

    bool loadLibrary() {
        m_x264Handle = TryLoadLibrary();
        if (!m_x264Handle) {
            return false;
        }

        // Load functions
        fn_x264_param_default_preset = (decltype(fn_x264_param_default_preset))
            dlsym(m_x264Handle, "x264_param_default_preset");

        fn_x264_param_apply_profile = (decltype(fn_x264_param_apply_profile))
            dlsym(m_x264Handle, "x264_param_apply_profile");

        fn_x264_picture_init = (decltype(fn_x264_picture_init))
            dlsym(m_x264Handle, "x264_picture_init");

        fn_x264_picture_alloc = (decltype(fn_x264_picture_alloc))
            dlsym(m_x264Handle, "x264_picture_alloc");

        fn_x264_picture_clean = (decltype(fn_x264_picture_clean))
            dlsym(m_x264Handle, "x264_picture_clean");

        fn_x264_encoder_close = (decltype(fn_x264_encoder_close))
            dlsym(m_x264Handle, "x264_encoder_close");

        // x264_encoder_open has version suffix based on X264_BUILD
        // Try common versions in order (newest first)
        const char* openNames[] = {
            "x264_encoder_open_164",
            "x264_encoder_open_163",
            "x264_encoder_open_162",
            "x264_encoder_open_161",
            "x264_encoder_open_160",
            "x264_encoder_open_159",
            "x264_encoder_open_157",
            "x264_encoder_open_155",
            "x264_encoder_open_152",
            "x264_encoder_open_148",
            nullptr
        };

        for (int i = 0; openNames[i]; i++) {
            fn_x264_encoder_open = (decltype(fn_x264_encoder_open))
                dlsym(m_x264Handle, openNames[i]);
            if (fn_x264_encoder_open) {
                fprintf(stderr, ">>> X264Encoder: found %s\n", openNames[i]);
                break;
            }
        }

        fn_x264_encoder_encode = (decltype(fn_x264_encoder_encode))
            dlsym(m_x264Handle, "x264_encoder_encode");

        // Check required functions
        if (!fn_x264_encoder_open || !fn_x264_encoder_encode || !fn_x264_encoder_close ||
            !fn_x264_param_default_preset || !fn_x264_picture_alloc) {
            fprintf(stderr, "*** X264Encoder: missing functions - open=%p encode=%p close=%p preset=%p alloc=%p\n",
                    (void*)fn_x264_encoder_open, (void*)fn_x264_encoder_encode,
                    (void*)fn_x264_encoder_close, (void*)fn_x264_param_default_preset,
                    (void*)fn_x264_picture_alloc);
            unloadLibrary();
            return false;
        }

        return true;
    }

    void unloadLibrary() {
        if (m_x264Handle) {
            dlclose(m_x264Handle);
            m_x264Handle = nullptr;
        }
        clearFunctionPointers();
    }

    /**
     * Convert BGRA to I420 (YUV 4:2:0 planar) directly into output planes
     * Using ITU-R BT.601 coefficients
     */
    void convertBGRAtoI420(const uint8_t* bgra, int stride, int direction,
                           uint8_t* yPlane, int yStride,
                           uint8_t* uPlane, int uStride,
                           uint8_t* vPlane, int vStride) {
        int srcStride = stride;
        int w = m_width;
        int h = m_height;

        // Direction: 1 = normal, -1 = flip vertically
        int startY = (direction > 0) ? 0 : (h - 1);
        int stepY = (direction > 0) ? 1 : -1;

        // Y plane: full resolution
        for (int j = 0; j < h; j++) {
            int srcY = startY + j * stepY;
            const uint8_t* srcRow = bgra + srcY * srcStride;
            uint8_t* dstRow = yPlane + j * yStride;

            for (int i = 0; i < w; i++) {
                uint8_t b = srcRow[i * 4 + 0];
                uint8_t g = srcRow[i * 4 + 1];
                uint8_t r = srcRow[i * 4 + 2];
                // Y = 0.257*R + 0.504*G + 0.098*B + 16
                int y = ((66 * r + 129 * g + 25 * b + 128) >> 8) + 16;
                dstRow[i] = (uint8_t)(y < 0 ? 0 : (y > 255 ? 255 : y));
            }
        }

        // U/V planes: half resolution (2x2 block averaging)
        int uvW = w / 2;
        int uvH = h / 2;

        for (int j = 0; j < uvH; j++) {
            int srcY0 = startY + (j * 2) * stepY;
            int srcY1 = startY + (j * 2 + 1) * stepY;
            const uint8_t* row0 = bgra + srcY0 * srcStride;
            const uint8_t* row1 = bgra + srcY1 * srcStride;

            for (int i = 0; i < uvW; i++) {
                // Average 4 pixels
                int r = 0, g = 0, b = 0;

                b += row0[(i*2+0)*4 + 0]; g += row0[(i*2+0)*4 + 1]; r += row0[(i*2+0)*4 + 2];
                b += row0[(i*2+1)*4 + 0]; g += row0[(i*2+1)*4 + 1]; r += row0[(i*2+1)*4 + 2];
                b += row1[(i*2+0)*4 + 0]; g += row1[(i*2+0)*4 + 1]; r += row1[(i*2+0)*4 + 2];
                b += row1[(i*2+1)*4 + 0]; g += row1[(i*2+1)*4 + 1]; r += row1[(i*2+1)*4 + 2];

                r >>= 2; g >>= 2; b >>= 2;

                // U = -0.148*R - 0.291*G + 0.439*B + 128
                int u = ((-38 * r - 74 * g + 112 * b + 128) >> 8) + 128;
                // V = 0.439*R - 0.368*G - 0.071*B + 128
                int v = ((112 * r - 94 * g - 18 * b + 128) >> 8) + 128;

                uPlane[j * uStride + i] = (uint8_t)(u < 0 ? 0 : (u > 255 ? 255 : u));
                vPlane[j * vStride + i] = (uint8_t)(v < 0 ? 0 : (v > 255 ? 255 : v));
            }
        }
    }

    /**
     * Convert RGB24 to I420 (YUV 4:2:0 planar) directly into output planes
     */
    void convertRGB24toI420(const uint8_t* rgb, int stride, int direction,
                            uint8_t* yPlane, int yStride,
                            uint8_t* uPlane, int uStride,
                            uint8_t* vPlane, int vStride) {
        int srcStride = stride;
        int w = m_width;
        int h = m_height;

        int startY = (direction > 0) ? 0 : (h - 1);
        int stepY = (direction > 0) ? 1 : -1;

        // Y plane
        for (int j = 0; j < h; j++) {
            int srcY = startY + j * stepY;
            const uint8_t* srcRow = rgb + srcY * srcStride;
            uint8_t* dstRow = yPlane + j * yStride;

            for (int i = 0; i < w; i++) {
                uint8_t r = srcRow[i * 3 + 0];
                uint8_t g = srcRow[i * 3 + 1];
                uint8_t b = srcRow[i * 3 + 2];
                int y = ((66 * r + 129 * g + 25 * b + 128) >> 8) + 16;
                dstRow[i] = (uint8_t)(y < 0 ? 0 : (y > 255 ? 255 : y));
            }
        }

        // U/V planes
        int uvW = w / 2;
        int uvH = h / 2;

        for (int j = 0; j < uvH; j++) {
            int srcY0 = startY + (j * 2) * stepY;
            int srcY1 = startY + (j * 2 + 1) * stepY;
            const uint8_t* row0 = rgb + srcY0 * srcStride;
            const uint8_t* row1 = rgb + srcY1 * srcStride;

            for (int i = 0; i < uvW; i++) {
                int r = 0, g = 0, b = 0;

                r += row0[(i*2+0)*3 + 0]; g += row0[(i*2+0)*3 + 1]; b += row0[(i*2+0)*3 + 2];
                r += row0[(i*2+1)*3 + 0]; g += row0[(i*2+1)*3 + 1]; b += row0[(i*2+1)*3 + 2];
                r += row1[(i*2+0)*3 + 0]; g += row1[(i*2+0)*3 + 1]; b += row1[(i*2+0)*3 + 2];
                r += row1[(i*2+1)*3 + 0]; g += row1[(i*2+1)*3 + 1]; b += row1[(i*2+1)*3 + 2];

                r >>= 2; g >>= 2; b >>= 2;

                int u = ((-38 * r - 74 * g + 112 * b + 128) >> 8) + 128;
                int v = ((112 * r - 94 * g - 18 * b + 128) >> 8) + 128;

                uPlane[j * uStride + i] = (uint8_t)(u < 0 ? 0 : (u > 255 ? 255 : u));
                vPlane[j * vStride + i] = (uint8_t)(v < 0 ? 0 : (v > 255 ? 255 : v));
            }
        }
    }
};