#include "effect.h"

#include <vector>
#include <regex>
#include <format>
#include <iostream>
#include <stdexcept>

#include "globals.h"

#include "texture.h"
#include "gstreamer_texture.h"

GLuint Effect::g_dummyVAO = 0;

static std::string shaderHeader = R"(#version 460 core

out vec4 FragColor;

in vec2 vUV;
in vec2 vPosition;

uniform float uTime;
uniform vec4 uMouse;
uniform vec2 uDesktopSize;
uniform vec4 uDisplay[16]; // support up to 16 displays, more can be added if needed
uniform vec4 uDisplayNorm[16]; // pre-normalized display rects for convenience
uniform int uNumDisplays;

#define MAX_DISPLAYS 16

// Background modes
#define BG_FIT_WIDTH 0
#define BG_FIT_HEIGHT 1
#define BG_COVER 2
#define BG_STRETCH 3
#define BG_CONTAIN 4
#define BG_TILE 5

vec2 GetDisplayUV(sampler2D tex, int displayIndex, int mode) {
    vec2 inUV = vec2(0.0);

    if (displayIndex < 0 || displayIndex >= uNumDisplays) {
        // Entire desktop as fallback
        inUV = vUV;
    } else {
        // Map global UV to display-local [0,1]
        inUV = (vUV - uDisplayNorm[displayIndex].xy) / uDisplayNorm[displayIndex].zw;
    }

    vec2 uv = inUV;

    if (mode == BG_STRETCH) {
        // No adjustment needed
        return uv;
    }

    vec4 display = uDisplay[displayIndex];
    float displayRatio = display.z / display.w;

    vec2 texSize = vec2(textureSize(tex, 0));
    float textureRatio = texSize.x / texSize.y;

    if (mode == BG_COVER) {
        if (displayRatio > textureRatio) {
            // Display wider: crop top/bottom
            float scale = textureRatio / displayRatio;
            uv.y = uv.y * scale + (1.0 - scale) * 0.5;
        } else {
            // Display taller: crop left/right
            float scale = displayRatio / textureRatio;
            uv.x = uv.x * scale + (1.0 - scale) * 0.5;
        }
    } else if (mode == BG_CONTAIN) {
        if (displayRatio > textureRatio) {
            // Display wider: pillarbox left/right
            float scale = displayRatio / textureRatio;
            uv.x = uv.x * scale + (1.0 - scale) * 0.5;
        } else {
            // Display taller: letterbox top/bottom
            float scale = textureRatio / displayRatio;
            uv.y = uv.y * scale + (1.0 - scale) * 0.5;
        }
    } else if (mode == BG_FIT_WIDTH) {
        // Scale to match display width, adjust Y
        float scale = textureRatio / displayRatio;
        uv.y = uv.y * scale + (1.0 - scale) * 0.5;
    } else if (mode == BG_FIT_HEIGHT) {
        // Scale to match display height, adjust X
        float scale = displayRatio / textureRatio;
        uv.x = uv.x * scale + (1.0 - scale) * 0.5;
    } else if (mode == BG_TILE) {
        // Tile at native texture resolution
        uv = uv * display.zw / texSize;
    }

    return uv;
}

// Default mode = BG_COVER
vec2 GetDisplayUV(sampler2D tex, int displayIndex) {
    return GetDisplayUV(tex, displayIndex, BG_COVER);
}

vec2 GetFullUV(sampler2D tex) {
    return GetDisplayUV(tex, -1, BG_COVER);
}

vec2 GetFullUV(sampler2D tex, int mode) {
    return GetDisplayUV(tex, -1, mode);
}

#line 1
)";

static std::vector<std::string> SplitLines(const std::string& str) {
    std::vector<std::string> lines;
    size_t start = 0, end = 0;
    while ((end = str.find_first_of("\r\n", start)) != std::string::npos) {
        lines.push_back(str.substr(start, end - start));
        start = end + 1;
    }
    if (start < str.size()) {
        lines.push_back(str.substr(start));
    }
    return lines;
}

Effect::Effect(const std::string& fragmentShaderSource)
{
    const std::string vertexShaderSource =
        #include "fx_vertex.glsl.h"
    ;

    GLuint vertexShader = CreateShader(vertexShaderSource, GL_VERTEX_SHADER);
    if (!vertexShader) {
        throw std::runtime_error("Failed to compile vertex shader");
        return;
    }

    GLuint fragmentShader = CreateShader(shaderHeader + fragmentShaderSource, GL_FRAGMENT_SHADER);
    if (!fragmentShader) {
        glDeleteShader(vertexShader);
        throw std::runtime_error("Failed to compile fragment shader");
        return;
    }

    m_program = glCreateProgram();
    glAttachShader(m_program, vertexShader);
    glAttachShader(m_program, fragmentShader);
    glLinkProgram(m_program);

    GLint success;
    glGetProgramiv(m_program, GL_LINK_STATUS, &success);
    if (!success) {
        GLint logLength;
        glGetProgramiv(m_program, GL_INFO_LOG_LENGTH, &logLength);
        std::string log(logLength, '\0');
        glGetProgramInfoLog(m_program, logLength, nullptr, log.data());
        throw std::runtime_error("Failed to link shader program: " + log);
    }

    glDeleteShader(vertexShader);
    glDeleteShader(fragmentShader);

    // Create a dummy VAO to allow rendering without a real VAO bound
    if (!g_dummyVAO) {
        glGenVertexArrays(1, &g_dummyVAO);
    }
}

Effect::~Effect()
{
    if (m_program) {
        glDeleteProgram(m_program);
        m_program = 0;
    }
}

void Effect::Use() const
{
    glUseProgram(m_program);
}

GLuint Effect::GetUniformLocation(const std::string& name)
{
    auto it = m_uniformLocations.find(name);
    if (it != m_uniformLocations.end()) {
        return it->second;
    }

    GLuint location = glGetUniformLocation(m_program, name.c_str());
    m_uniformLocations[name] = location;
    return location;
}

void Effect::SetTexture(const std::string &name, const Texture& texture, size_t textureUnit)
{
    glActiveTexture(GL_TEXTURE0 + static_cast<GLuint>(textureUnit));
    texture.Bind();
    
    GStreamerTexture* gstTexture = dynamic_cast<GStreamerTexture*>(const_cast<Texture*>(&texture));
    if (gstTexture) {
        // For GStreamerTexture, we need to call Update() to upload the latest frame data
        gstTexture->Update();
    }

    glUniform1i(GetUniformLocation(name), static_cast<GLint>(textureUnit));
}

void Effect::SetInt(const std::string &name, int value)
{
    glUniform1i(GetUniformLocation(name), value);
}

void Effect::SetFloat(const std::string &name, float value)
{
    glUniform1f(GetUniformLocation(name), value);
}

void Effect::SetVector2(const std::string &name, const Vector2 &value)
{
    glUniform2f(GetUniformLocation(name), value.x, value.y);
}

void Effect::SetVector3(const std::string &name, const Vector3 &value)
{
    glUniform3f(GetUniformLocation(name), value.x, value.y, value.z);
}

void Effect::SetVector4(const std::string &name, const Vector4 &value)
{
    glUniform4f(GetUniformLocation(name), value.x, value.y, value.z, value.w);
}

void Effect::SetMatrix4(const std::string &name, const Matrix4 &value)
{
    glUniformMatrix4fv(GetUniformLocation(name), 1, GL_FALSE, value.data());
}

void Effect::Render()
{
    Use();

    // Set default uniforms
    SetFloat("uTime", g_Time);
    SetVector4("uMouse", g_Mouse);
    SetVector2("uDesktopSize", g_DesktopSize);
    for (size_t i = 0; i < g_Displays.size(); i++) {
        SetVector4("uDisplay[" + std::to_string(i) + "]", g_Displays[i]);
        SetVector4("uDisplayNorm[" + std::to_string(i) + "]", g_Displays[i] / Vector4(g_DesktopSize, g_DesktopSize));
    }
    SetInt("uNumDisplays", static_cast<int>(g_Displays.size()));

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

GLuint Effect::CreateShader(const std::string &source, GLenum shaderType)
{
    GLuint shader = glCreateShader(shaderType);

    const char* src = source.c_str();
    glShaderSource(shader, 1, &src, nullptr);
    glCompileShader(shader);
    GLint success;
    glGetShaderiv(shader, GL_COMPILE_STATUS, &success);

    if (!success) {
        std::vector<std::string> sourceLines = SplitLines(source);

        GLint logLength;
        glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &logLength);
        std::string log(logLength, '\0');
        glGetShaderInfoLog(shader, logLength, nullptr, log.data());

        std::string logText = "";
        std::string shaderTypeText = "??";

        switch (shaderType) {
            case GL_VERTEX_SHADER: shaderTypeText = "Vertex Shader"; break;
            case GL_FRAGMENT_SHADER: shaderTypeText = "Fragment Shader"; break;
            default: shaderTypeText = "Shader"; break;
        }

        std::smatch match;
        int lineNumber = 0;
        int column = -1;
        std::string errorMessage;

        // Mesa/Intel: "0:15(10): error: msg" or "ERROR: 0:15: msg"
        std::regex mesaRegex(R"((?:ERROR: )?\d+:(\d+)(?:\((\d+)\))?:\s*(.*))");
        // NVIDIA: "0(11) : error C7530: msg"
        std::regex nvidiaRegex(R"(\d+\((\d+)\)\s*:\s*(.*))");

        if (std::regex_search(log, match, mesaRegex)) {
            lineNumber = std::stoi(match[1].str());
            column = match[2].matched ? std::stoi(match[2].str()) : -1;
            errorMessage = match[3].str();
        } else if (std::regex_search(log, match, nvidiaRegex)) {
            lineNumber = std::stoi(match[1].str());
            errorMessage = match[2].str();
        }

        if (lineNumber > 0) {
            logText = std::format("{} error at line {}: {}\nOffending code:\n", shaderTypeText, lineNumber, errorMessage);
            if (lineNumber <= static_cast<int>(sourceLines.size())) {
                std::string line = sourceLines[lineNumber - 1];

                // add a ^ at column in a new line
                std::string arrow(line.size(), ' ');
                if (column >= 0 && column < static_cast<int>(arrow.size())) {
                    arrow[column] = '^';
                }

                if (lineNumber - 2 >= 0) {
                    logText += std::format("{}\n", sourceLines[lineNumber - 2]);
                }
                logText += std::format("{}\n{}", line, arrow);
                if (lineNumber < static_cast<int>(sourceLines.size())) {
                    logText += std::format("\n{}", sourceLines[lineNumber]);
                }
            }
        }

        std::cout << logText << std::endl;

        glDeleteShader(shader);
        return 0;
    }

    return shader;
}