#include "opengl_video_widget.h"
#include <QDebug>

OpenGLVideoWidget::OpenGLVideoWidget(QWidget* parent)
    : QOpenGLWidget(parent)
    , m_program(nullptr)
    , m_textureId(0)
    , m_frameData(nullptr)
    , m_frameWidth(0)
    , m_frameHeight(0)
    , m_frameFormat(0)
    , m_frameUpdated(false)
    , m_initialized(false)
{
    // 设置顶点坐标
    m_vertices[0] = -1.0f; m_vertices[1] = -1.0f;
    m_vertices[2] = 1.0f;  m_vertices[3] = -1.0f;
    m_vertices[4] = -1.0f; m_vertices[5] = 1.0f;
    m_vertices[6] = 1.0f;  m_vertices[7] = 1.0f;

    // 设置纹理坐标
    m_texCoords[0] = 0.0f; m_texCoords[1] = 1.0f;
    m_texCoords[2] = 1.0f; m_texCoords[3] = 1.0f;
    m_texCoords[4] = 0.0f; m_texCoords[5] = 0.0f;
    m_texCoords[6] = 1.0f; m_texCoords[7] = 0.0f;
}

OpenGLVideoWidget::~OpenGLVideoWidget()
{
    Close();
}

bool OpenGLVideoWidget::Open(unsigned int width, unsigned int height)
{
    QMutexLocker locker(&m_mutex);
    
    m_frameWidth = width;
    m_frameHeight = height;
    
    // 如果已经有数据，释放它
    if (m_frameData) {
        delete[] m_frameData;
    }
    
    // 分配新的内存
    m_frameData = new unsigned char[width * height * 4]; // RGBA格式
    memset(m_frameData, 0, width * height * 4);
    
    return true;
}

void OpenGLVideoWidget::Close()
{
    makeCurrent();
    if (m_textureId) {
        glDeleteTextures(1, &m_textureId);
        m_textureId = 0;
    }
    if (m_program) {
        delete m_program;
        m_program = nullptr;
    }
    doneCurrent();
    
    // 释放帧数据
    QMutexLocker locker(&m_mutex);
    if (m_frameData) {
        delete[] m_frameData;
        m_frameData = nullptr;
    }
    
    m_frameWidth = 0;
    m_frameHeight = 0;
    m_frameUpdated = false;
    m_initialized = false;
}

void OpenGLVideoWidget::initializeGL()
{
    initializeOpenGLFunctions();
    glClearColor(0.0f, 0.0f, 0.0f, 1.0f);

    // 创建着色器程序
    m_program = new QOpenGLShaderProgram();
    m_program->addShaderFromSourceCode(QOpenGLShader::Vertex,
        "attribute vec2 vertexIn;\n"
        "attribute vec2 textureIn;\n"
        "varying vec2 textureOut;\n"
        "void main(void)\n"
        "{\n"
        "    gl_Position = vec4(vertexIn, 0.0, 1.0);\n"
        "    textureOut = textureIn;\n"
        "}\n");
    m_program->addShaderFromSourceCode(QOpenGLShader::Fragment,
        "varying vec2 textureOut;\n"
        "uniform sampler2D texture;\n"
        "void main(void)\n"
        "{\n"
        "    gl_FragColor = texture2D(texture, textureOut);\n"
        "}\n");
    m_program->bindAttributeLocation("vertexIn", 0);
    m_program->bindAttributeLocation("textureIn", 1);
    m_program->link();
    
    // 创建纹理
    glGenTextures(1, &m_textureId);
    glBindTexture(GL_TEXTURE_2D, m_textureId);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glBindTexture(GL_TEXTURE_2D, 0);
    
    m_initialized = true;
}

void OpenGLVideoWidget::resizeGL(int width, int height)
{
    glViewport(0, 0, width, height);
}

void OpenGLVideoWidget::paintGL()
{
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
    
    QMutexLocker locker(&m_mutex);
    
    if (!m_frameData || m_frameWidth <= 0 || m_frameHeight <= 0 || !m_frameUpdated)
        return;
    
    // 绑定纹理并更新数据
    glBindTexture(GL_TEXTURE_2D, m_textureId);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, m_frameWidth, m_frameHeight, 
                0, GL_RGBA, GL_UNSIGNED_BYTE, m_frameData);
    
    // 使用着色器程序
    m_program->bind();
    
    // 设置纹理单元
    m_program->setUniformValue("texture", 0);
    
    // 设置顶点和纹理坐标
    m_program->enableAttributeArray(0);
    m_program->enableAttributeArray(1);
    m_program->setAttributeArray(0, m_vertices, 2);
    m_program->setAttributeArray(1, m_texCoords, 2);
    
    // 绘制
    glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
    
    // 清理
    m_program->disableAttributeArray(0);
    m_program->disableAttributeArray(1);
    m_program->release();
    glBindTexture(GL_TEXTURE_2D, 0);
}

void OpenGLVideoWidget::updateFrame(const VideoFrame& frame)
{
    if (!frame.data || frame.width <= 0 || frame.height <= 0)
        return;
    
    QMutexLocker locker(&m_mutex);
    
    // 如果尺寸变化，重新分配内存
    if (m_frameWidth != frame.width || m_frameHeight != frame.height) {
        if (m_frameData) {
            delete[] m_frameData;
        }
        
        m_frameWidth = frame.width;
        m_frameHeight = frame.height;
        m_frameData = new unsigned char[m_frameWidth * m_frameHeight * 4]; // RGBA格式
    }
    
    // 复制帧数据
    memcpy(m_frameData, frame.data, m_frameWidth * m_frameHeight * 4);
    m_frameUpdated = true;
    
    // 请求重绘
    update();
}

bool OpenGLVideoWidget::convertFromAVFrame(AVFrame* frame)
{
    if (!frame || frame->width <= 0 || frame->height <= 0)
        return false;

    QMutexLocker locker(&m_mutex);

    // 如果尺寸变化，重新分配内存
    if (m_frameWidth != frame->width || m_frameHeight != frame->height) {
        if (m_frameData) {
            delete[] m_frameData;
        }

        m_frameWidth = frame->width;
        m_frameHeight = frame->height;
        m_frameData = new unsigned char[m_frameWidth * m_frameHeight * 4]; // RGBA格式
    }

    // 根据不同的像素格式进行转换
    switch (frame->format) {
    case AV_PIX_FMT_RGBA: {
        // 直接复制RGBA数据
        for (int y = 0; y < frame->height; y++) {
            memcpy(m_frameData + y * m_frameWidth * 4,
                   frame->data[0] + y * frame->linesize[0],
                   frame->width * 4);
        }
    } break;
    case AV_PIX_FMT_RGB24: {
        // RGB24转RGBA
        for (int y = 0; y < frame->height; y++) {
            uint8_t* src = frame->data[0] + y * frame->linesize[0];
            uint8_t* dst = m_frameData + y * m_frameWidth * 4;

            for (int x = 0; x < frame->width; x++) {
                *dst++ = *src++; // R
                *dst++ = *src++; // G
                *dst++ = *src++; // B
                *dst++ = 255;    // A
            }
        }
    } break;
    case AV_PIX_FMT_BGR0:
    case AV_PIX_FMT_BGRA: {
        // BGRA转RGBA
        for (int y = 0; y < frame->height; y++) {
            uint8_t* src = frame->data[0] + y * frame->linesize[0];
            uint8_t* dst = m_frameData + y * m_frameWidth * 4;

            for (int x = 0; x < frame->width; x++) {
                uint8_t b = *src++;
                uint8_t g = *src++;
                uint8_t r = *src++;
                uint8_t a = *src++;

                *dst++ = r;
                *dst++ = g;
                *dst++ = b;
                *dst++ = a;
            }
        }
    } break;
    case AV_PIX_FMT_YUV420P: // 添加对YUV420P格式的支持
    {
        // YUV420P转RGBA
        for (int y = 0; y < frame->height; y++) {
            uint8_t* dst = m_frameData + y * m_frameWidth * 4;

            for (int x = 0; x < frame->width; x++) {
                int Y = frame->data[0][y * frame->linesize[0] + x];
                int U = frame->data[1][(y / 2) * frame->linesize[1] + (x / 2)];
                int V = frame->data[2][(y / 2) * frame->linesize[2] + (x / 2)];

                // YUV转RGB公式
                int C = Y - 16;
                int D = U - 128;
                int E = V - 128;

                int R = (298 * C + 409 * E + 128) >> 8;
                int G = (298 * C - 100 * D - 208 * E + 128) >> 8;
                int B = (298 * C + 516 * D + 128) >> 8;

                // 限制RGB值在0-255范围内
                R = R < 0 ? 0 : (R > 255 ? 255 : R);
                G = G < 0 ? 0 : (G > 255 ? 255 : G);
                B = B < 0 ? 0 : (B > 255 ? 255 : B);

                *dst++ = R;   // R
                *dst++ = G;   // G
                *dst++ = B;   // B
                *dst++ = 255; // A
            }
        }
    } break;
    case AV_PIX_FMT_NV12: {
        // NV12转RGBA
        for (int y = 0; y < frame->height; y++) {
            uint8_t* dst = m_frameData + y * m_frameWidth * 4;

            for (int x = 0; x < frame->width; x++) {
                int Y = frame->data[0][y * frame->linesize[0] + x];
                int U = frame->data[1][(y / 2) * frame->linesize[1] + (x / 2) * 2];
                int V = frame->data[1][(y / 2) * frame->linesize[1] + (x / 2) * 2 + 1];

                // YUV转RGB公式
                int C = Y - 16;
                int D = U - 128;
                int E = V - 128;

                int R = (298 * C + 409 * E + 128) >> 8;
                int G = (298 * C - 100 * D - 208 * E + 128) >> 8;
                int B = (298 * C + 516 * D + 128) >> 8;

                // 限制RGB值在0-255范围内
                R = R < 0 ? 0 : (R > 255 ? 255 : R);
                G = G < 0 ? 0 : (G > 255 ? 255 : G);
                B = B < 0 ? 0 : (B > 255 ? 255 : B);

                *dst++ = R;   // R
                *dst++ = G;   // G
                *dst++ = B;   // B
                *dst++ = 255; // A
            }
        }
    } break;
    default:
        // 对于其他格式，可以考虑使用FFmpeg的sws_scale函数
        qDebug() << "Unsupported pixel format:" << frame->format;
        return false;
    }

    m_frameUpdated = true;
    update();
    return true;
}

bool OpenGLVideoWidget::Render(AVFrame* frame)
{
    if (!m_initialized && isValid()) {
        makeCurrent();
        initializeGL();
        doneCurrent();
    }
    
    if (!frame) {
        update(); // 仅刷新显示
        return true;
    }
    
    return convertFromAVFrame(frame);
}

void OpenGLVideoWidget::clearFrame()
{
    QMutexLocker locker(&m_mutex);
    m_frameUpdated = false;
    update();
}