#include "composition_pipeline.h"

#include <android/log.h>
#include <algorithm>
#include <cstring>

#define TAG "LckComposition"
#define LOGI(...) __android_log_print(ANDROID_LOG_INFO, TAG, __VA_ARGS__)
#define LOGW(...) __android_log_print(ANDROID_LOG_WARN, TAG, __VA_ARGS__)
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, TAG, __VA_ARGS__)

// --- Shaders ---

static const char* BASE_VERTEX_SHADER = R"(#version 300 es
layout(location = 0) in vec2 aPos;
layout(location = 1) in vec2 aTexCoord;
out vec2 vTexCoord;
void main() {
    gl_Position = vec4(aPos, 0.0, 1.0);
    vTexCoord = aTexCoord;
}
)";

// Base pass: renders game frame (OES texture) full-screen to FBO
// - Flips V coordinate: Vulkan render targets have origin at top-left, GL at bottom-left
// - Swizzles BGR→RGB: Vulkan RGBA8 maps to BGRA in some AHardwareBuffer formats
static const char* BASE_FRAGMENT_SHADER = R"(#version 300 es
#extension GL_OES_EGL_image_external_essl3 : require
precision mediump float;
in vec2 vTexCoord;
out vec4 fragColor;
uniform samplerExternalOES uTexture;
void main() {
    vec2 flippedCoord = vec2(vTexCoord.x, 1.0 - vTexCoord.y);
    vec4 color = texture(uTexture, flippedCoord);
    fragColor = vec4(color.b, color.g, color.r, color.a);
}
)";

// Overlay pass: renders 2D layers with MVP transform and opacity
static const char* OVERLAY_VERTEX_SHADER = R"(#version 300 es
layout(location = 0) in vec2 aPos;
layout(location = 1) in vec2 aTexCoord;
uniform mat4 uMVP;
out vec2 vTexCoord;
void main() {
    gl_Position = uMVP * vec4(aPos, 0.0, 1.0);
    vTexCoord = aTexCoord;
}
)";

static const char* OVERLAY_FRAGMENT_SHADER = R"(#version 300 es
precision mediump float;
in vec2 vTexCoord;
out vec4 fragColor;
uniform sampler2D uTexture;
uniform float uOpacity;
void main() {
    vec4 color = texture(uTexture, vTexCoord);
    fragColor = vec4(color.rgb, color.a * uOpacity);
}
)";

// Standby pattern: dark gradient with color bars to prove the pipeline is alive
static const char* STANDBY_FRAGMENT_SHADER = R"(#version 300 es
precision mediump float;
in vec2 vTexCoord;
out vec4 fragColor;
void main() {
    // Vertical color bars
    float x = vTexCoord.x;
    vec3 color;
    if      (x < 0.125) color = vec3(0.75, 0.75, 0.75); // white-ish
    else if (x < 0.250) color = vec3(0.75, 0.75, 0.0);  // yellow
    else if (x < 0.375) color = vec3(0.0,  0.75, 0.75);  // cyan
    else if (x < 0.500) color = vec3(0.0,  0.75, 0.0);   // green
    else if (x < 0.625) color = vec3(0.75, 0.0,  0.75);  // magenta
    else if (x < 0.750) color = vec3(0.75, 0.0,  0.0);   // red
    else if (x < 0.875) color = vec3(0.0,  0.0,  0.75);  // blue
    else                 color = vec3(0.15, 0.15, 0.15);  // dark gray
    // Darken bottom third
    float brightness = vTexCoord.y > 0.33 ? 1.0 : 0.5;
    fragColor = vec4(color * brightness, 1.0);
}
)";

// --- Helpers ---

static GLuint CompileShader(GLenum type, const char* source) {
    GLuint shader = glCreateShader(type);
    glShaderSource(shader, 1, &source, nullptr);
    glCompileShader(shader);

    GLint status;
    glGetShaderiv(shader, GL_COMPILE_STATUS, &status);
    if (!status) {
        char log[512];
        glGetShaderInfoLog(shader, sizeof(log), nullptr, log);
        LOGE("Shader compile error: %s", log);
        glDeleteShader(shader);
        return 0;
    }
    return shader;
}

static GLuint LinkProgram(GLuint vs, GLuint fs) {
    GLuint program = glCreateProgram();
    glAttachShader(program, vs);
    glAttachShader(program, fs);
    glLinkProgram(program);
    glDeleteShader(vs);
    glDeleteShader(fs);

    GLint status;
    glGetProgramiv(program, GL_LINK_STATUS, &status);
    if (!status) {
        char log[512];
        glGetProgramInfoLog(program, sizeof(log), nullptr, log);
        LOGE("Program link error: %s", log);
        glDeleteProgram(program);
        return 0;
    }
    return program;
}

// Identity 4x4 matrix
static void Mat4Identity(float* m) {
    memset(m, 0, 16 * sizeof(float));
    m[0] = m[5] = m[10] = m[15] = 1.0f;
}

// --- CompositionPipeline ---

CompositionPipeline::CompositionPipeline() {}

CompositionPipeline::~CompositionPipeline() {
    Release();
}

bool CompositionPipeline::Init(int width, int height) {
    if (initialized) Release();

    fboWidth = width;
    fboHeight = height;

    // Create FBO color attachment texture (GL_RGBA8)
    glGenTextures(1, &fboTexture);
    glBindTexture(GL_TEXTURE_2D, fboTexture);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, fboWidth, fboHeight, 0,
                 GL_RGBA, GL_UNSIGNED_BYTE, nullptr);
    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);

    // Create FBO
    glGenFramebuffers(1, &fbo);
    glBindFramebuffer(GL_FRAMEBUFFER, fbo);
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, fboTexture, 0);

    GLenum fboStatus = glCheckFramebufferStatus(GL_FRAMEBUFFER);
    if (fboStatus != GL_FRAMEBUFFER_COMPLETE) {
        LOGE("FBO incomplete: 0x%x", fboStatus);
        glBindFramebuffer(GL_FRAMEBUFFER, 0);
        Release();
        return false;
    }
    glBindFramebuffer(GL_FRAMEBUFFER, 0);

    // Compile base pass program (OES)
    {
        GLuint vs = CompileShader(GL_VERTEX_SHADER, BASE_VERTEX_SHADER);
        GLuint fs = CompileShader(GL_FRAGMENT_SHADER, BASE_FRAGMENT_SHADER);
        if (!vs || !fs) { Release(); return false; }
        baseProgram = LinkProgram(vs, fs);
        if (!baseProgram) { Release(); return false; }
    }

    // Compile overlay program (sampler2D + MVP + opacity)
    {
        GLuint vs = CompileShader(GL_VERTEX_SHADER, OVERLAY_VERTEX_SHADER);
        GLuint fs = CompileShader(GL_FRAGMENT_SHADER, OVERLAY_FRAGMENT_SHADER);
        if (!vs || !fs) { Release(); return false; }
        overlayProgram = LinkProgram(vs, fs);
        if (!overlayProgram) { Release(); return false; }
    }

    // Compile standby pattern program
    {
        GLuint vs = CompileShader(GL_VERTEX_SHADER, BASE_VERTEX_SHADER);
        GLuint fs = CompileShader(GL_FRAGMENT_SHADER, STANDBY_FRAGMENT_SHADER);
        if (!vs || !fs) { Release(); return false; }
        standbyProgram = LinkProgram(vs, fs);
        if (!standbyProgram) { Release(); return false; }
    }

    overlayMvpLoc = glGetUniformLocation(overlayProgram, "uMVP");
    overlayOpacityLoc = glGetUniformLocation(overlayProgram, "uOpacity");
    overlayTexLoc = glGetUniformLocation(overlayProgram, "uTexture");

    // Create shared full-screen quad VAO: pos(x,y) + texcoord(u,v)
    float quad[] = {
        -1.0f, -1.0f,  0.0f, 0.0f,
         1.0f, -1.0f,  1.0f, 0.0f,
        -1.0f,  1.0f,  0.0f, 1.0f,
         1.0f,  1.0f,  1.0f, 1.0f,
    };

    glGenVertexArrays(1, &quadVao);
    glGenBuffers(1, &quadVbo);
    glBindVertexArray(quadVao);
    glBindBuffer(GL_ARRAY_BUFFER, quadVbo);
    glBufferData(GL_ARRAY_BUFFER, sizeof(quad), quad, GL_STATIC_DRAW);
    glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), (void*)0);
    glEnableVertexAttribArray(0);
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), (void*)(2 * sizeof(float)));
    glEnableVertexAttribArray(1);
    glBindVertexArray(0);

    initialized = true;
    LOGI("Composition pipeline initialized: %dx%d", fboWidth, fboHeight);
    return true;
}

void CompositionPipeline::Compose(GLuint srcOesTexture) {
    if (!initialized) return;

    glBindFramebuffer(GL_FRAMEBUFFER, fbo);
    glViewport(0, 0, fboWidth, fboHeight);
    glClearColor(0.05f, 0.05f, 0.05f, 1.0f);
    glClear(GL_COLOR_BUFFER_BIT);

    // Base pass: render game frame (or leave cleared dark background if no game input)
    if (srcOesTexture != 0) {
        RenderBasePass(srcOesTexture);
    }

    // Overlay pass: render layers sorted by zOrder
    RenderOverlayLayers();

    glBindFramebuffer(GL_FRAMEBUFFER, 0);
}

void CompositionPipeline::RenderBasePass(GLuint srcOesTexture) {
    glUseProgram(baseProgram);
    glActiveTexture(GL_TEXTURE0);
    glBindTexture(GL_TEXTURE_EXTERNAL_OES, srcOesTexture);
    glUniform1i(glGetUniformLocation(baseProgram, "uTexture"), 0);

    glBindVertexArray(quadVao);
    glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
    glBindVertexArray(0);
}

void CompositionPipeline::RenderStandbyPattern() {
    glUseProgram(standbyProgram);
    glBindVertexArray(quadVao);
    glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
    glBindVertexArray(0);
}

void CompositionPipeline::RenderOverlayLayers() {
    // Take a snapshot of layers under lock
    std::vector<CompositionLayer> snapshot;
    {
        std::lock_guard<std::mutex> lock(layerMutex);
        snapshot = layers;
    }

    // Sort by zOrder
    std::sort(snapshot.begin(), snapshot.end(),
              [](const CompositionLayer& a, const CompositionLayer& b) {
                  return a.zOrder < b.zOrder;
              });

    // Enable alpha blending
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

    glUseProgram(overlayProgram);

    for (const auto& layer : snapshot) {
        if (!layer.enabled || layer.textureId == 0 || layer.opacity <= 0.0f) continue;

        float mvp[16];
        BuildTransformMatrix(layer, mvp);

        glUniformMatrix4fv(overlayMvpLoc, 1, GL_FALSE, mvp);
        glUniform1f(overlayOpacityLoc, layer.opacity);

        glActiveTexture(GL_TEXTURE0);
        glBindTexture(GL_TEXTURE_2D, layer.textureId);
        glUniform1i(overlayTexLoc, 0);

        glBindVertexArray(quadVao);
        glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
        glBindVertexArray(0);
    }

    glDisable(GL_BLEND);
}

void CompositionPipeline::BuildTransformMatrix(const CompositionLayer& layer, float* m) {
    // Build 2D transform matrix in NDC space.
    // The quad is [-1,1] full-screen. We need to scale it to layer size
    // relative to FBO, then apply position/rotation.

    const auto& t = layer.transform;

    // Layer size in NDC (layer pixel size / FBO pixel size → fraction → *2 for NDC range)
    float layerW = (float)layer.texWidth / (float)fboWidth;
    float layerH = (float)layer.texHeight / (float)fboHeight;

    float sx = layerW * t.scaleX;
    float sy = layerH * t.scaleY;

    float cosR = cosf(t.rotation);
    float sinR = sinf(t.rotation);

    // Build: Translate(pos) * Rotate(r) * Scale(s)
    // Column-major 4x4
    Mat4Identity(m);
    m[0]  =  sx * cosR;
    m[1]  =  sx * sinR;
    m[4]  = -sy * sinR;
    m[5]  =  sy * cosR;
    m[12] =  t.posX;
    m[13] =  t.posY;
}

int CompositionPipeline::AddLayer(GLuint textureId, int texW, int texH,
                                   const CompositionTransform& transform,
                                   float opacity, int zOrder, const std::string& tag,
                                   bool ownsTexture) {
    std::lock_guard<std::mutex> lock(layerMutex);
    CompositionLayer layer;
    layer.id = nextLayerId++;
    layer.textureId = textureId;
    layer.texWidth = texW;
    layer.texHeight = texH;
    layer.transform = transform;
    layer.opacity = opacity;
    layer.zOrder = zOrder;
    layer.tag = tag;
    layer.ownsTexture = ownsTexture;
    layers.push_back(layer);
    LOGI("Added composition layer %d ('%s') z=%d", layer.id, tag.c_str(), zOrder);
    return layer.id;
}

void CompositionPipeline::RemoveLayer(int layerId) {
    std::lock_guard<std::mutex> lock(layerMutex);
    for (auto it = layers.begin(); it != layers.end(); ++it) {
        if (it->id == layerId) {
            if (it->ownsTexture && it->textureId) {
                glDeleteTextures(1, &it->textureId);
            }
            LOGI("Removed composition layer %d ('%s')", it->id, it->tag.c_str());
            layers.erase(it);
            return;
        }
    }
    LOGW("RemoveLayer: layer %d not found", layerId);
}

void CompositionPipeline::UpdateLayerTransform(int layerId, const CompositionTransform& transform) {
    std::lock_guard<std::mutex> lock(layerMutex);
    for (auto& layer : layers) {
        if (layer.id == layerId) {
            layer.transform = transform;
            return;
        }
    }
}

void CompositionPipeline::UpdateLayerOpacity(int layerId, float opacity) {
    std::lock_guard<std::mutex> lock(layerMutex);
    for (auto& layer : layers) {
        if (layer.id == layerId) {
            layer.opacity = opacity;
            return;
        }
    }
}

void CompositionPipeline::SetLayerEnabled(int layerId, bool enabled) {
    std::lock_guard<std::mutex> lock(layerMutex);
    for (auto& layer : layers) {
        if (layer.id == layerId) {
            layer.enabled = enabled;
            return;
        }
    }
}

GLuint CompositionPipeline::UploadTexture(const uint8_t* rgbaData, int w, int h) {
    if (!rgbaData || w <= 0 || h <= 0) return 0;

    GLuint tex;
    glGenTextures(1, &tex);
    glBindTexture(GL_TEXTURE_2D, tex);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, w, h, 0,
                 GL_RGBA, GL_UNSIGNED_BYTE, rgbaData);
    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);

    return tex;
}

void CompositionPipeline::Release() {
    // Delete owned layer textures
    {
        std::lock_guard<std::mutex> lock(layerMutex);
        for (auto& layer : layers) {
            if (layer.ownsTexture && layer.textureId) {
                glDeleteTextures(1, &layer.textureId);
            }
        }
        layers.clear();
    }

    if (quadVao) { glDeleteVertexArrays(1, &quadVao); quadVao = 0; }
    if (quadVbo) { glDeleteBuffers(1, &quadVbo); quadVbo = 0; }
    if (baseProgram) { glDeleteProgram(baseProgram); baseProgram = 0; }
    if (overlayProgram) { glDeleteProgram(overlayProgram); overlayProgram = 0; }
    if (standbyProgram) { glDeleteProgram(standbyProgram); standbyProgram = 0; }
    if (fbo) { glDeleteFramebuffers(1, &fbo); fbo = 0; }
    if (fboTexture) { glDeleteTextures(1, &fboTexture); fboTexture = 0; }

    initialized = false;
    LOGI("Composition pipeline released");
}
