#include "ShaderYUV420.h"


#include <QOpenGLWidget>
#include <QOpenGLFunctions_3_3_Core>

#include "spdlog/spdlog.h"

Image_YUV420::Image_YUV420(Image_YUV420&& other): width(other.width), height(other.height),
          yData(std::move(other.yData)), uData(std::move(other.uData)), vData(std::move(other.vData))
{}

Image_YUV420::Image_YUV420(const Image_YUV420& other)
        : width(other.width), height(other.height),
          yData(other.yData), uData(other.uData), vData(other.vData)
{}

Image_YUV420& Image_YUV420::operator=(Image_YUV420&& other) 
{
    if (this != &other) {
        width = other.width;
        height = other.height;
        yData = std::move(other.yData);
        uData = std::move(other.uData);
        vData = std::move(other.vData);
    }
    return *this;
}
Image_YUV420& Image_YUV420::operator=(const Image_YUV420& other) 
{
    if (this != &other) {
        width = other.width;
        height = other.height;
        yData = other.yData;
        uData = other.uData;
        vData = other.vData;
    }
    return *this;
}

bool Image_YUV420::isValid() const 
{
    return width > 0 && height > 0 && !yData.isEmpty() && !uData.isEmpty() && !vData.isEmpty();
}

void Image_YUV420::clear()
{
    width = 0;
    height = 0;
    yData.clear();
    uData.clear();
    vData.clear();
}


ShaderYUV420::ShaderYUV420()
{

}
ShaderYUV420::~ShaderYUV420()
{
    QOpenGLFunctions_3_3_Core *gl = QOpenGLContext::currentContext()->versionFunctions<QOpenGLFunctions_3_3_Core>();
    if (m_VAO != 0) {
        if (gl) {
            gl->glDeleteVertexArrays(1, &m_VAO);
            gl->glDeleteBuffers(1, &m_VBO);
            gl->glDeleteBuffers(1, &m_EBO);
        }
    }
}

/* 初始化形状 */
GLuint ShaderYUV420::initShape()
{
    const int vertexCount = 4;    // 顶点数量
    const int vertexSize = 8;     // 顶点大小
    /* 矩形位置，铺满屏幕 */
    float vertices [vertexCount * vertexSize] {
        /* ----- 位置 -----|-------- 颜色 ----------|----- 纹理坐标 ------ */
        -1.0f, -1.0f, 0.0f,     1.0f, 1.0f, 0.0f,   0.0f, 0.0f,     // 左下
        1.0f, -1.0f, 0.0f,      0.0f, 0.0f, 1.0f,   1.0f, 0.0f,     // 右下
        1.0f,  1.0f, 0.0f,      1.0f, 0.0f, 0.0f,   1.0f, 1.0f,     // 右上
        -1.0f, 1.0f, 0.0f,      0.0f, 1.0f, 0.0f,   0.0f, 1.0f      // 右上
    };
    const int indexCount = 6; // 索引数量
    /* 顶点顺序 */
    uint32_t indices [indexCount] {
        0, 1, 2, // 第一个三角形
        2, 3, 0  // 第二个三角形
    };

    QOpenGLFunctions_3_3_Core *gl = QOpenGLContext::currentContext()->versionFunctions<QOpenGLFunctions_3_3_Core>();
    if (!gl) {
        SPDLOG_ERROR("Failed to get OpenGL functions");
        return 0;
    }
    /* -------------------------------------------------------------------------------------
     * 顶点数组对象
     * ------------------------------------------------------------------------------------ */
    gl->glGenVertexArrays(1, &m_VAO); // 生成一个VAO，返回的ID存储在VAO1中
    gl->glBindVertexArray(m_VAO); // 绑定VAO

    /* -------------------------------------------------------------------------------------
     * 创建顶点数组缓冲区
     * ------------------------------------------------------------------------------------ */
    gl->glGenBuffers(1, &m_VBO);
    gl->glBindBuffer(GL_ARRAY_BUFFER, m_VBO); // 绑定VBO
    gl->glBufferData(GL_ARRAY_BUFFER, vertexCount * vertexSize * sizeof(float), vertices, GL_STATIC_DRAW);

    /* -------------------------------------------------------------------------------------
     * 创建EBO
     * ------------------------------------------------------------------------------------ */
    gl->glGenBuffers(1, &m_EBO);
    gl->glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, m_EBO);
    gl->glBufferData(GL_ELEMENT_ARRAY_BUFFER, indexCount * sizeof(uint32_t), indices, GL_STATIC_DRAW); // 绑定EBO


    /* -------------------------------------------------------------------------------------
     * 链接顶点属性
     * ------------------------------------------------------------------------------------ */
    // 设置顶点属性指针，告诉OpenGL如何解析顶点数据
    gl->glVertexAttribPointer(
                    0,                  /* 顶点属性位置，0表示第一个属性 */
                    3,                  /* 每个顶点属性的大小，这里是3，表示x、y、z坐标 */
                    GL_FLOAT, GL_FALSE, /* 数据类型和是否标准化 */
                    vertexSize * sizeof(float), /* 步长，表示每个顶点属性之间的间隔 */
                    (void*)0            /* 偏移量，表示该属性（顶点坐标）从第0个字节开始 */
                ); 
    /* 启用顶点属性 */
    gl->glEnableVertexAttribArray(0);
 
    // 设置颜色位置，最后一个参数表示从第3个字节开始，倒数第二个参数是两个数据之间起始位置间隔8个字节
    gl->glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, vertexSize * sizeof(float), (void*)(3*sizeof(float))); 
    gl->glEnableVertexAttribArray(1);
 
    /* 设置纹理位置 */
    gl->glVertexAttribPointer(2, 2, GL_FLOAT, GL_FALSE, vertexSize * sizeof(float), (void*)(6*sizeof(float))); // 纹理坐标位置
    gl->glEnableVertexAttribArray(2); // 启用纹理坐标属性

    return m_VAO;
}

/* 初始化纹理对象 */
bool ShaderYUV420::initTextures()
{
    int unit = 0;
    unit = createTexture(m_textureYName, 0);
    if(unit < 0) {
        SPDLOG_ERROR("Failed to create Y texture");
        return false;
    }
    unit = createTexture(m_textureUName, 1);
    if(unit < 0) {
        SPDLOG_ERROR("Failed to create U texture");
        return false;
    }
    unit = createTexture(m_textureVName, 2);
    if(unit < 0) {
        SPDLOG_ERROR("Failed to create V texture");
        return false;
    }

    return true;
}

/* 绘制图形 */
void ShaderYUV420::drawShape()
{
    QOpenGLFunctions_3_3_Core *gl = QOpenGLContext::currentContext()->versionFunctions<QOpenGLFunctions_3_3_Core>();
    /* 需要在使用时手动绑定VAO，绘制图形 */
    gl->glBindVertexArray(m_VAO); // 绑定VAO

    gl->glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_INT, 0); // 绘制四边形

    gl->glBindVertexArray(0); // 解绑VAO
}

/* 刷新一帧 */
bool ShaderYUV420::refreshFrameYUV420(const Image_YUV420& yuvData)
{
    if (!yuvData.isValid()) {
        SPDLOG_ERROR("Invalid YUV420 data");
        return false;
    }
    QOpenGLTexture *textureY = m_mapTexture.find(0)->second;
    QOpenGLTexture *textureU = m_mapTexture.find(1)->second;
    QOpenGLTexture *textureV = m_mapTexture.find(2)->second;

    if(textureY == nullptr || textureU == nullptr || textureV == nullptr) {
        SPDLOG_ERROR("Failed to get YUV textures");
        return false;
    }

    if(yuvData.width != m_lastYUVSize.width() || yuvData.height != m_lastYUVSize.height()) 
    {
        /* 如果YUV数据的大小发生变化，则需要重新创建纹理 */
        textureY->destroy();
        textureU->destroy();
        textureV->destroy();
        textureY->create();
        textureU->create();
        textureV->create();

        textureY->setSize(yuvData.width, yuvData.height, 1); // 设置Y分量纹理大小
        textureY->setFormat(QOpenGLTexture::R8_UNorm); // 设置Y分量纹理格式
        textureY->allocateStorage(); // 分配Y分量纹理存储空间
        
        textureU->setSize(yuvData.width / 2, yuvData.height / 2, 1); // 设置U分量纹理大小
        textureU->setFormat(QOpenGLTexture::R8_UNorm); // 设置U分量纹理格式
        textureU->allocateStorage(); // 分配U分量纹理存储空间

        textureV->setSize(yuvData.width / 2, yuvData.height / 2, 1); // 设置V分量纹理大小
        textureV->setFormat(QOpenGLTexture::R8_UNorm); // 设置V分量纹理格式
        textureV->allocateStorage(); // 分配V分量纹理存储空间

        m_lastYUVSize = QSize(yuvData.width, yuvData.height);
    }

    /* 设置Y分量纹理 */
    textureY->setData(QOpenGLTexture::Red, QOpenGLTexture::UInt8, yuvData.yData.constData());
    textureY->generateMipMaps(); // 生成多级渐远纹理

    /* 设置U分量纹理 */
    textureU->setData(QOpenGLTexture::Red, QOpenGLTexture::UInt8, yuvData.uData.constData());
    textureU->generateMipMaps(); // 生成多级渐远纹理

    /* 设置V分量纹理 */
    textureV->setData(QOpenGLTexture::Red, QOpenGLTexture::UInt8, yuvData.vData.constData());
    textureV->generateMipMaps(); // 生成多级渐远纹理

    printGLerror("ShaderYUV420::refreshFrameYUV420");

    return true;
}

/* 刷新一帧 */
bool ShaderYUV420::refreshFrameRGBA(const QImage& image, int textureUnit)
{
    m_yuvData.clear(); // 清空YUV数据
    convertRGBAToYUV420(image, m_yuvData);
    
    refreshFrameYUV420(m_yuvData);
    return true;
}

/* RGBA转换成YUV420 */
bool ShaderYUV420::convertRGBAToYUV420(const QImage& image, Image_YUV420& yuvData)
{
    int width = image.width();
    int height = image.height();
    const int ySize = yuvData.width * yuvData.height;
    const int uvSize = ySize / 4;

    QImage image_rgb(width, height, QImage::Format_RGB888);

    yuvData.width = width;
    yuvData.height = height;

    yuvData.yData.resize(ySize);
    yuvData.uData.resize(uvSize);
    yuvData.vData.resize(uvSize);

    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            QColor color = image.pixelColor(x, y);
            int r = color.red();
            int g = color.green();
            int b = color.blue();

            int yIndex = y * width + x;
            yuvData.yData[yIndex] = static_cast<unsigned char>((0.257 * r) + (0.504 * g) + (0.098 * b) + 16);

            if (y % 2 == 0 && x % 2 == 0) {
                int uvIndex = (y / 2) * (width / 2) + (x / 2);
                yuvData.uData[uvIndex] = static_cast<unsigned char>((-0.148 * r) - (0.291 * g) + (0.439 * b) + 128);
                yuvData.vData[uvIndex] = static_cast<unsigned char>((0.439 * r) - (0.368 * g) - (0.071 * b) + 128);
            }
        }
    }

    return true;
}