// SPDX-FileCopyrightText: Copyright 2025 Citron Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

// Shader headers generated by CMake
#include "video_core/host_shaders/taa_frag.h"
#include "video_core/host_shaders/taa_vert.h"
#include "video_core/renderer_opengl/gl_shader_manager.h"
#include "video_core/renderer_opengl/gl_shader_util.h"
#include "video_core/renderer_opengl/present/taa.h"
#include "video_core/renderer_opengl/present/util.h"
#include "common/logging/log.h"

namespace OpenGL {

TAA::TAA(u32 render_width, u32 render_height) : width(render_width), height(render_height), current_frame(0) {
    // Validate dimensions
    if (width == 0 || height == 0) {
        LOG_ERROR(Render_OpenGL, "TAA: Invalid dimensions {}x{}", width, height);
        return;
    }

    vert_shader = CreateProgram(HostShaders::TAA_VERT, GL_VERTEX_SHADER);
    frag_shader = CreateProgram(HostShaders::TAA_FRAG, GL_FRAGMENT_SHADER);

    sampler = CreateBilinearSampler();

    // Create uniform buffer
    uniform_buffer.Create();
    glNamedBufferStorage(uniform_buffer.handle, sizeof(TaaParams), nullptr,
                        GL_DYNAMIC_STORAGE_BIT | GL_MAP_WRITE_BIT);

    // Initialize TAA parameters
    params.frame_count = 0.0f;
    params.blend_factor = 0.25f; // Increased blend factor to reduce ghosting
    params.inv_resolution[0] = 1.0f / static_cast<float>(width);
    params.inv_resolution[1] = 1.0f / static_cast<float>(height);
    params.motion_scale = 1.0f;
    params.jitter_offset[0] = 0.0f;
    params.jitter_offset[1] = 0.0f;

    CreateFramebuffers();
}

TAA::~TAA() = default;

void TAA::CreateFramebuffers() {
    framebuffer.Create();

    // Current frame texture (RGBA16F for HDR support)
    current_texture.Create(GL_TEXTURE_2D);
    glTextureStorage2D(current_texture.handle, 1, GL_RGBA16F, width, height);
    glNamedFramebufferTexture(framebuffer.handle, GL_COLOR_ATTACHMENT0, current_texture.handle, 0);

    // Previous frame texture
    previous_texture.Create(GL_TEXTURE_2D);
    glTextureStorage2D(previous_texture.handle, 1, GL_RGBA16F, width, height);

    // Motion vector texture (RG16F for 2D motion vectors) - initialize with zeros
    motion_texture.Create(GL_TEXTURE_2D);
    glTextureStorage2D(motion_texture.handle, 1, GL_RG16F, width, height);
    // Clear motion texture to zero motion vectors
    glClearTexImage(motion_texture.handle, 0, GL_RG, GL_FLOAT, nullptr);

    // Depth texture (R32F for depth buffer) - initialize with far depth
    depth_texture.Create(GL_TEXTURE_2D);
    glTextureStorage2D(depth_texture.handle, 1, GL_R32F, width, height);
    // Clear depth texture to far depth (1.0)
    const float far_depth = 1.0f;
    glClearTexImage(depth_texture.handle, 0, GL_RED, GL_FLOAT, &far_depth);

    // Check framebuffer completeness
    GLenum status = glCheckNamedFramebufferStatus(framebuffer.handle, GL_FRAMEBUFFER);
    if (status != GL_FRAMEBUFFER_COMPLETE) {
        LOG_ERROR(Render_OpenGL, "TAA framebuffer incomplete: 0x{:04X}", status);
    }
}

void TAA::UpdateJitter(u32 frame_count) {
    // Halton sequence (2,3) for low-discrepancy sampling
    constexpr float halton_2[8] = {0.0f, 0.5f, 0.25f, 0.75f, 0.125f, 0.625f, 0.375f, 0.875f};
    constexpr float halton_3[8] = {0.0f, 0.333333f, 0.666667f, 0.111111f, 0.444444f, 0.777778f, 0.222222f, 0.555556f};

    // Ensure safe array access
    const size_t index = static_cast<size_t>(frame_count) % 8;
    // Reduce jitter intensity to minimize visible jittering
    const float jitter_scale = 0.5f; // Reduce jitter by 50%
    params.jitter_offset[0] = (halton_2[index] - 0.5f) * jitter_scale * params.inv_resolution[0];
    params.jitter_offset[1] = (halton_3[index] - 0.5f) * jitter_scale * params.inv_resolution[1];
}

GLuint TAA::Draw(ProgramManager& program_manager, GLuint input_texture, GLuint prev_texture,
                 GLuint motion_text, GLuint depth_text, u32 frame_count) {
    // Validate input texture
    if (input_texture == 0) {
        LOG_ERROR(Render_OpenGL, "TAA: Invalid input texture");
        return input_texture;
    }

    // Update parameters
    params.frame_count = static_cast<float>(frame_count);
    UpdateJitter(frame_count);

    // Update uniform buffer
    void* buffer_data = glMapNamedBuffer(uniform_buffer.handle, GL_WRITE_ONLY);
    if (buffer_data) {
        memcpy(buffer_data, &params, sizeof(TaaParams));
        glUnmapNamedBuffer(uniform_buffer.handle);
    }

    // Bind TAA program
    program_manager.BindPresentPrograms(vert_shader.handle, frag_shader.handle);

    // Bind framebuffer
    glBindFramebuffer(GL_DRAW_FRAMEBUFFER, framebuffer.handle);

    // Bind textures - use fallback textures if motion/depth are not available
    glBindTextureUnit(0, input_texture);      // Current frame
    glBindTextureUnit(1, previous_texture.handle); // Previous frame (from TAA)

    // Use motion texture if provided, otherwise use a dummy texture
    if (motion_text != 0 && motion_text != GL_NONE) {
        glBindTextureUnit(2, motion_text);
    } else {
        glBindTextureUnit(2, motion_texture.handle);   // Use our dummy motion texture
    }

    // Use depth texture if provided, otherwise use a dummy texture
    if (depth_text != 0 && depth_text != GL_NONE) {
        glBindTextureUnit(3, depth_text);
    } else {
        glBindTextureUnit(3, depth_texture.handle);    // Use our dummy depth texture
    }

    // Bind samplers
    glBindSampler(0, sampler.handle);
    glBindSampler(1, sampler.handle);
    glBindSampler(2, sampler.handle);
    glBindSampler(3, sampler.handle);

    // Bind uniform buffer
    glBindBufferBase(GL_UNIFORM_BUFFER, 5, uniform_buffer.handle);

    // Draw full-screen triangle
    glFrontFace(GL_CCW);
    glDrawArrays(GL_TRIANGLES, 0, 3);
    glFrontFace(GL_CW);

    // Copy current frame to previous frame for next iteration
    glCopyImageSubData(current_texture.handle, GL_TEXTURE_2D, 0, 0, 0, 0,
                       previous_texture.handle, GL_TEXTURE_2D, 0, 0, 0, 0,
                       width, height, 1);

    return current_texture.handle;
}

void TAA::SwapBuffers() {
    // Swap current and previous textures
    std::swap(current_texture, previous_texture);
    current_frame++;
}

} // namespace OpenGL
