#include "glwidget.h"
#include "videodecoder.h"
#include <QOpenGLContext>
#include <QDebug>

GLWidget::GLWidget(QWidget *parent)
    : QOpenGLWidget(parent)
    , m_videoDecoder(nullptr)
    , m_texY(0)
    , m_texUV(0)
    , m_vbo(0)
    , m_cudaResourceY(nullptr)
    , m_cudaResourceUV(nullptr)
    , m_allocatedWidth(0)
    , m_allocatedHeight(0)
{
    // Almacenamos el contenedor local del frame decodificado
    m_localFrame = av_frame_alloc();
}

GLWidget::~GLWidget() {
    // Nos aseguramos de tener el contexto de OpenGL activo antes de borrar recursos de la GPU
    makeCurrent();
    cleanupCudaGL();
    
    if (m_texY) {
        glDeleteTextures(1, &m_texY);
    }
    if (m_texUV) {
        glDeleteTextures(1, &m_texUV);
    }
    if (m_vbo) {
        glDeleteBuffers(1, &m_vbo);
    }
    
    // Liberamos la estructura local de FFmpeg
    if (m_localFrame) {
        av_frame_free(&m_localFrame);
        m_localFrame = nullptr;
    }
}

void GLWidget::setVideoDecoder(VideoDecoder *decoder) {
    m_videoDecoder = decoder;
    // Conectamos la señal frameReady del decodificador con la repintada de este Widget
    connect(m_videoDecoder, &VideoDecoder::frameReady, this, &GLWidget::onFrameReady);
}

void GLWidget::onFrameReady() {
    // Solicita a Qt redibujar el widget (desencadena una llamada interna a paintGL())
    update();
}

void GLWidget::clearVideo() {
    // Libera los recursos de interoperabilidad y actualiza el render
    cleanupCudaGL();
    update();
}

void GLWidget::initializeGL() {
    // Inicializa la tabla de funciones de OpenGL para poder llamarlas directamente
    initializeOpenGLFunctions();
    glClearColor(0.0f, 0.0f, 0.0f, 1.0f);

    // Compilar e inicializar los shaders desde los archivos fuente
    m_shaderProgram.addShaderFromSourceFile(QOpenGLShader::Vertex, "shader.vert");
    m_shaderProgram.addShaderFromSourceFile(QOpenGLShader::Fragment, "shader.frag");
    m_shaderProgram.link();

    // Crear VBO para un quad de pantalla completa (dos triángulos coordinados en 2D)
    // Nota: Espejamos el eje vertical (coordenada V) de textura para compensar la diferencia 
    // de origen de coordenadas: OpenGL asume esquina inferior-izquierda como (0,0) y CUDA asume superior-izquierda.
    float vertices[] = {
        // Posiciones   // Coordenadas UV
        -1.0f, -1.0f,   0.0f, 1.0f,
         1.0f, -1.0f,   1.0f, 1.0f,
        -1.0f,  1.0f,   0.0f, 0.0f,
        
        -1.0f,  1.0f,   0.0f, 0.0f,
         1.0f, -1.0f,   1.0f, 1.0f,
         1.0f,  1.0f,   1.0f, 0.0f
    };

    glGenBuffers(1, &m_vbo);
    glBindBuffer(GL_ARRAY_BUFFER, m_vbo);
    glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);

    // Generar texturas de OpenGL
    // Textura Y (Luminancia)
    glGenTextures(1, &m_texY);
    glBindTexture(GL_TEXTURE_2D, m_texY);
    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);

    // Textura UV (Crominancia entrelazada)
    glGenTextures(1, &m_texUV);
    glBindTexture(GL_TEXTURE_2D, m_texUV);
    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);
}

void GLWidget::resizeGL(int w, int h) {
    // Ajusta la ventana gráfica de OpenGL al tamaño del widget físico
    glViewport(0, 0, w, h);
}

void GLWidget::paintGL() {
    glClear(GL_COLOR_BUFFER_BIT);

    if (!m_videoDecoder) return;

    // 1. Obtener la referencia segura del fotograma decodificado por hardware (residiendo en VRAM)
    if (!m_videoDecoder->getCurrentFrameRef(m_localFrame)) {
        return; // No hay frame listo todavía para dibujar
    }

    int w = m_localFrame->width;
    int h = m_localFrame->height;

    // 2. Si la resolución del video cambió o no se ha asignado aún espacio en VRAM
    if (w != m_allocatedWidth || h != m_allocatedHeight) {
        cleanupCudaGL();

        // 2.1 Reservar espacio en VRAM para la textura Y (Luminancia - formato GL_RED, tamaño completo)
        glBindTexture(GL_TEXTURE_2D, m_texY);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, w, h, 0, GL_RED, GL_UNSIGNED_BYTE, nullptr);

        // 2.2 Reservar espacio en VRAM para la textura UV (Crominancia - formato GL_RG, tamaño a la mitad)
        glBindTexture(GL_TEXTURE_2D, m_texUV);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RG, w / 2, h / 2, 0, GL_RG, GL_UNSIGNED_BYTE, nullptr);

        // 2.3 Registrar los identificadores de textura de OpenGL en la API de interoperabilidad de CUDA.
        // La bandera 'cudaGraphicsRegisterFlagsWriteDiscard' le dice a CUDA que sobrescribiremos la textura
        // por completo y que descarte el contenido anterior para mejorar el rendimiento del driver.
        cudaError_t errY = cudaGraphicsGLRegisterImage(&m_cudaResourceY, m_texY, GL_TEXTURE_2D, cudaGraphicsRegisterFlagsWriteDiscard);
        cudaError_t errUV = cudaGraphicsGLRegisterImage(&m_cudaResourceUV, m_texUV, GL_TEXTURE_2D, cudaGraphicsRegisterFlagsWriteDiscard);

        if (errY != cudaSuccess || errUV != cudaSuccess) {
            qWarning() << "Error registrando texturas en CUDA:" << cudaGetErrorString(errY) << "/" << cudaGetErrorString(errUV);
            cleanupCudaGL();
            av_frame_unref(m_localFrame);
            return;
        }

        m_allocatedWidth = w;
        m_allocatedHeight = h;
    }

    // 3. Mapear las texturas de OpenGL hacia CUDA para poder escribir en ellas de manera exclusiva
    if (m_cudaResourceY && m_cudaResourceUV) {
        cudaError_t mapY = cudaGraphicsMapResources(1, &m_cudaResourceY, 0);
        cudaError_t mapUV = cudaGraphicsMapResources(1, &m_cudaResourceUV, 0);

        if (mapY == cudaSuccess && mapUV == cudaSuccess) {
            cudaArray_t arrayY = nullptr;
            cudaArray_t arrayUV = nullptr;

            // Extraemos los manejadores de arrays internos de CUDA asociados a nuestras texturas de OpenGL
            cudaGraphicsSubResourceGetMappedArray(&arrayY, m_cudaResourceY, 0, 0);
            cudaGraphicsSubResourceGetMappedArray(&arrayUV, m_cudaResourceUV, 0, 0);

            // Copia Directa interna de GPU (VRAM-to-VRAM) de la textura Y (Luminancia)
            // Usamos linesize[0] como el pitch (ancho de fila físico en memoria) del plano.
            cudaError_t cpyY = cudaMemcpy2DToArray(
                arrayY, 0, 0,
                (void*)m_localFrame->data[0], m_localFrame->linesize[0],
                w, h,
                cudaMemcpyDeviceToDevice
            );

            // Copia Directa interna de GPU (VRAM-to-VRAM) de la textura UV (Crominancia entrelazada)
            // Se le pasa el ancho físico en bytes ('w' bytes) debido a que tiene 2 componentes de 1 byte (U y V) por píxel.
            cudaError_t cpyUV = cudaMemcpy2DToArray(
                arrayUV, 0, 0,
                (void*)m_localFrame->data[1], m_localFrame->linesize[1],
                w, h / 2,
                cudaMemcpyDeviceToDevice
            );

            if (cpyY != cudaSuccess || cpyUV != cudaSuccess) {
                qWarning() << "Fallo en copias CUDA Memcpy2DToArray:" << cudaGetErrorString(cpyY) << "/" << cudaGetErrorString(cpyUV);
            }
        } else {
            qWarning() << "Fallo al mapear texturas en CUDA:" << cudaGetErrorString(mapY) << "/" << cudaGetErrorString(mapUV);
        }

        // Importante: desmapear los recursos para que CUDA libere el bloqueo y OpenGL vuelva a tomar el control del pintado
        cudaGraphicsUnmapResources(1, &m_cudaResourceY, 0);
        cudaGraphicsUnmapResources(1, &m_cudaResourceUV, 0);
    }

    // 4. Dibujar la escena en pantalla usando el Shader
    m_shaderProgram.bind();

    // Enlazar textura Y a la unidad de textura 0
    glActiveTexture(GL_TEXTURE0);
    glBindTexture(GL_TEXTURE_2D, m_texY);
    m_shaderProgram.setUniformValue("u_texY", 0);

    // Enlazar textura UV a la unidad de textura 1
    glActiveTexture(GL_TEXTURE1);
    glBindTexture(GL_TEXTURE_2D, m_texUV);
    m_shaderProgram.setUniformValue("u_texUV", 1);

    // Activar el VBO de vértices y configurar atributos para a_position y a_texCoord
    glBindBuffer(GL_ARRAY_BUFFER, m_vbo);
    
    int posLocation = m_shaderProgram.attributeLocation("a_position");
    m_shaderProgram.enableAttributeArray(posLocation);
    glVertexAttribPointer(posLocation, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), nullptr);

    int texLocation = m_shaderProgram.attributeLocation("a_texCoord");
    m_shaderProgram.enableAttributeArray(texLocation);
    glVertexAttribPointer(texLocation, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), reinterpret_cast<void*>(2 * sizeof(float)));

    // Dibujar los 6 vértices (2 triángulos que forman el quad)
    glDrawArrays(GL_TRIANGLES, 0, 6);

    // Desactivar y liberar
    m_shaderProgram.disableAttributeArray(posLocation);
    m_shaderProgram.disableAttributeArray(texLocation);
    m_shaderProgram.release();

    // Liberar nuestra referencia local del fotograma (decrementa el conteo de referencias del AVFrame de FFmpeg)
    av_frame_unref(m_localFrame);
}

void GLWidget::cleanupCudaGL() {
    makeCurrent();
    // Desregistra los recursos de CUDA asociados a las texturas de OpenGL
    if (m_cudaResourceY) {
        cudaGraphicsUnregisterResource(m_cudaResourceY);
        m_cudaResourceY = nullptr;
    }
    if (m_cudaResourceUV) {
        cudaGraphicsUnregisterResource(m_cudaResourceUV);
        m_cudaResourceUV = nullptr;
    }
    m_allocatedWidth = 0;
    m_allocatedHeight = 0;
}