#include <wayland-client.h>
#include <wayland-egl.h>
#include <EGL/egl.h>
#include <GLES2/gl2.h>
#include <cstring>
#include <iostream>
#include <vector>
#include <fstream>
#include <sstream>
#include <chrono>
#include <cmath>

#include "xdg-shell-client-protocol.h"
#include "font8x8.h"

// --- Globals & State ---

struct wl_display *display = nullptr;
struct wl_registry *registry = nullptr;
struct wl_compositor *compositor = nullptr;
struct xdg_wm_base *wm_base = nullptr;
struct wl_seat *seat = nullptr;
struct wl_pointer *pointer = nullptr;

struct wl_surface *surface = nullptr;
struct xdg_surface *xdg_surface = nullptr;
struct xdg_toplevel *xdg_toplevel = nullptr;
struct wl_egl_window *egl_window = nullptr;

EGLDisplay egl_display;
EGLContext egl_context;
EGLSurface egl_surface;

// Text State
GLuint font_texture;
GLuint text_program;

bool running = true;
int width = 800;
int height = 600;

// Input state
double cursor_x = 0;
double cursor_y = 0;
bool button_pressed = false;

// Time state
auto start_time = std::chrono::high_resolution_clock::now();

// Shaders
struct ShaderProgram {
    GLuint program;
    GLint posAttrib;
    GLint timeUniform;
    GLint resolutionUniform;
    std::string name;
    float color[3]; // For sidebar button color
};

std::vector<ShaderProgram> shaders;
int current_shader_index = 0;
GLuint sidebar_program; // Simple shader for drawing UI quads

// UI Layout
const int SIDEBAR_WIDTH = 200;
const int BUTTON_HEIGHT = 60;
const int BUTTON_MARGIN = 10;

// Forward declarations
void render_frame();
GLuint load_shader_src(const char *source, GLenum type);

// --- Font & Text Utils ---

void init_font() {
    int tex_w = 128; // 16 chars * 8 px
    int tex_h = 64;  // 8 chars * 8 px
    std::vector<GLubyte> pixels(tex_w * tex_h, 0);

    for (int char_idx = 0; char_idx < 128; ++char_idx) {
        int grid_x = char_idx % 16;
        int grid_y = char_idx / 16;
        
        for (int row = 0; row < 8; ++row) {
            unsigned char bits = font8x8_basic[char_idx][row];
            for (int col = 0; col < 8; ++col) {
                if ((bits >> col) & 1) { 
                    int px = grid_x * 8 + col;
                    int py = grid_y * 8 + row;
                    pixels[py * tex_w + px] = 255;
                }
            }
        }
    }

    glGenTextures(1, &font_texture);
    glBindTexture(GL_TEXTURE_2D, font_texture);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_ALPHA, tex_w, tex_h, 0, GL_ALPHA, GL_UNSIGNED_BYTE, pixels.data());
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
}

void init_text_shader() {
    const char* vs = 
        "#version 100\n"
        "attribute vec2 position;\n"
        "attribute vec2 texcoord;\n"
        "varying vec2 v_tex;\n"
        "uniform vec2 u_screen_size;\n"
        "void main() {\n"
        "   vec2 ndc = (position / u_screen_size) * 2.0 - 1.0;\n"
        "   gl_Position = vec4(ndc.x, -ndc.y, 0.0, 1.0);\n"
        "   v_tex = texcoord;\n"
        "}";
    const char* fs = 
        "#version 100\n"
        "precision mediump float;\n"
        "varying vec2 v_tex;\n"
        "uniform sampler2D u_tex;\n"
        "uniform vec3 u_color;\n"
        "void main() {\n"
        "   float alpha = texture2D(u_tex, v_tex).a;\n"
        "   if (alpha < 0.5) discard;\n"
        "   gl_FragColor = vec4(u_color, 1.0);\n"
        "}";

    GLuint v = load_shader_src(vs, GL_VERTEX_SHADER);
    GLuint f = load_shader_src(fs, GL_FRAGMENT_SHADER);
    text_program = glCreateProgram();
    glAttachShader(text_program, v);
    glAttachShader(text_program, f);
    glLinkProgram(text_program);
}

void draw_text(const std::string& text, float x, float y, float scale, float r, float g, float b) {
    glUseProgram(text_program);
    GLint posLoc = glGetAttribLocation(text_program, "position");
    GLint texLoc = glGetAttribLocation(text_program, "texcoord");
    GLint scrLoc = glGetUniformLocation(text_program, "u_screen_size");
    GLint colLoc = glGetUniformLocation(text_program, "u_color");

    glUniform2f(scrLoc, (float)width, (float)height);
    glUniform3f(colLoc, r, g, b);

    std::vector<float> verts;
    float cx = x;
    float cy = y;

    for (char c : text) {
        if (c < 0 || c > 127) continue;
        int grid_x = c % 16;
        int grid_y = c / 16;
        
        float u1 = (float)grid_x / 16.0;
        float v1 = (float)grid_y / 8.0;
        float u2 = u1 + 1.0/16.0;
        float v2 = v1 + 1.0/8.0;

        float w = 8 * scale;
        float h = 8 * scale;

        verts.push_back(cx);     verts.push_back(cy + h); verts.push_back(u1); verts.push_back(v2);
        verts.push_back(cx + w); verts.push_back(cy + h); verts.push_back(u2); verts.push_back(v2);
        verts.push_back(cx + w); verts.push_back(cy);     verts.push_back(u2); verts.push_back(v1);

        verts.push_back(cx);     verts.push_back(cy + h); verts.push_back(u1); verts.push_back(v2);
        verts.push_back(cx + w); verts.push_back(cy);     verts.push_back(u2); verts.push_back(v1);
        verts.push_back(cx);     verts.push_back(cy);     verts.push_back(u1); verts.push_back(v1);

        cx += w + 2.0f; 
    }

    glEnableVertexAttribArray(posLoc);
    glEnableVertexAttribArray(texLoc);
    glVertexAttribPointer(posLoc, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), verts.data());
    glVertexAttribPointer(texLoc, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float), verts.data() + 2);

    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
    
    glDrawArrays(GL_TRIANGLES, 0, verts.size() / 4);
    
    glDisable(GL_BLEND);
}

// --- Shader Utilities ---

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

    GLint compiled;
    glGetShaderiv(shader, GL_COMPILE_STATUS, &compiled);
    if (!compiled) {
        char log[512];
        glGetShaderInfoLog(shader, 512, nullptr, log);
        std::cerr << "Shader compile error: " << log << std::endl;
        return 0;
    }
    return shader;
}

GLuint load_shader_file(const char *path, GLenum type) {
    std::ifstream file(path);
    if (!file.is_open()) {
        std::cerr << "Failed to open shader: " << path << std::endl;
        return 0;
    }
    std::stringstream buffer;
    buffer << file.rdbuf();
    std::string source_str = buffer.str();
    return load_shader_src(source_str.c_str(), type);
}

GLuint create_program(const char* vertPath, const char* fragPath) {
    GLuint vert = load_shader_file(vertPath, GL_VERTEX_SHADER);
    GLuint frag = load_shader_file(fragPath, GL_FRAGMENT_SHADER);
    GLuint prog = glCreateProgram();
    glAttachShader(prog, vert);
    glAttachShader(prog, frag);
    glLinkProgram(prog);
    return prog;
}

// --- Input Handling (Pointer) ---

static void pointer_handle_enter(void *data, struct wl_pointer *wl_pointer, uint32_t serial, struct wl_surface *surface, wl_fixed_t sx, wl_fixed_t sy) {}
static void pointer_handle_leave(void *data, struct wl_pointer *wl_pointer, uint32_t serial, struct wl_surface *surface) {}

static void pointer_handle_motion(void *data, struct wl_pointer *wl_pointer, uint32_t time, wl_fixed_t sx, wl_fixed_t sy) {
    cursor_x = wl_fixed_to_double(sx);
    cursor_y = wl_fixed_to_double(sy);
}

static void pointer_handle_button(void *data, struct wl_pointer *wl_pointer, uint32_t serial, uint32_t time, uint32_t button, uint32_t state) {
    if (state == WL_POINTER_BUTTON_STATE_PRESSED && button == 0x110) { // Left click
        if (cursor_x < SIDEBAR_WIDTH) {
            int clicked_index = -1;
            for (size_t i = 0; i < shaders.size(); ++i) {
                int y_start = BUTTON_MARGIN + i * (BUTTON_HEIGHT + BUTTON_MARGIN);
                int y_end = y_start + BUTTON_HEIGHT;
                if (cursor_y >= y_start && cursor_y <= y_end && cursor_x >= BUTTON_MARGIN && cursor_x <= SIDEBAR_WIDTH - BUTTON_MARGIN) {
                    clicked_index = i;
                    break;
                }
            }
            if (clicked_index != -1) {
                current_shader_index = clicked_index;
                std::cout << "Switched to shader: " << shaders[current_shader_index].name << std::endl;
            }
        }
    }
}

static void pointer_handle_axis(void *data, struct wl_pointer *wl_pointer, uint32_t time, uint32_t axis, wl_fixed_t value) {}

static const struct wl_pointer_listener pointer_listener = {
    .enter = pointer_handle_enter,
    .leave = pointer_handle_leave,
    .motion = pointer_handle_motion,
    .button = pointer_handle_button,
    .axis = pointer_handle_axis,
};

// --- Seat Listener ---

static void seat_handle_capabilities(void *data, struct wl_seat *wl_seat, uint32_t capabilities) {
    if (capabilities & WL_SEAT_CAPABILITY_POINTER) {
        pointer = wl_seat_get_pointer(wl_seat);
        wl_pointer_add_listener(pointer, &pointer_listener, nullptr);
    }
}

static void seat_handle_name(void *data, struct wl_seat *wl_seat, const char *name) {}

static const struct wl_seat_listener seat_listener = {
    .capabilities = seat_handle_capabilities,
    .name = seat_handle_name,
};

// --- Wayland/XDG Code ---

static void xdg_wm_base_ping(void *data, struct xdg_wm_base *xdg_wm_base, uint32_t serial) { xdg_wm_base_pong(xdg_wm_base, serial); }
static const struct xdg_wm_base_listener xdg_wm_base_listener = { .ping = xdg_wm_base_ping };

static void xdg_surface_configure(void *data, struct xdg_surface *xdg_surface, uint32_t serial) {
    xdg_surface_ack_configure(xdg_surface, serial);
}
static const struct xdg_surface_listener xdg_surface_listener = { .configure = xdg_surface_configure };

static void xdg_toplevel_configure(void *data, struct xdg_toplevel *xdg_toplevel, int32_t w, int32_t h, struct wl_array *states) {
    if (w > 0 && h > 0) {
        width = w;
        height = h;
        if (egl_window) wl_egl_window_resize(egl_window, width, height, 0, 0);
    }
}
static void xdg_toplevel_close(void *data, struct xdg_toplevel *xdg_toplevel) { running = false; }
static const struct xdg_toplevel_listener xdg_toplevel_listener = { .configure = xdg_toplevel_configure, .close = xdg_toplevel_close };

static void registry_handle_global(void *data, struct wl_registry *registry, uint32_t name, const char *interface, uint32_t version) {
    if (strcmp(interface, wl_compositor_interface.name) == 0) {
        compositor = (struct wl_compositor *)wl_registry_bind(registry, name, &wl_compositor_interface, 1);
    } else if (strcmp(interface, xdg_wm_base_interface.name) == 0) {
        wm_base = (struct xdg_wm_base *)wl_registry_bind(registry, name, &xdg_wm_base_interface, 1);
        xdg_wm_base_add_listener(wm_base, &xdg_wm_base_listener, nullptr);
    } else if (strcmp(interface, wl_seat_interface.name) == 0) {
        seat = (struct wl_seat *)wl_registry_bind(registry, name, &wl_seat_interface, 1);
        wl_seat_add_listener(seat, &seat_listener, nullptr);
    }
}
static const struct wl_registry_listener registry_listener = { .global = registry_handle_global, .global_remove = [](void*, wl_registry*, uint32_t){} };

// --- Rendering ---

static void frame_handle_done(void *data, struct wl_callback *callback, uint32_t time);
static const struct wl_callback_listener frame_listener = { .done = frame_handle_done };

static void frame_handle_done(void *data, struct wl_callback *callback, uint32_t time) {
    wl_callback_destroy(callback);
    render_frame();
}

void draw_quad(float x, float y, float w, float h, float screen_w, float screen_h, float r, float g, float b) {
    float x1 = (x / screen_w) * 2.0f - 1.0f;
    float y1 = 1.0f - (y / screen_h) * 2.0f;
    float x2 = ((x + w) / screen_w) * 2.0f - 1.0f;
    float y2 = 1.0f - ((y + h) / screen_h) * 2.0f;

    GLfloat vertices[] = { x1, y1, x1, y2, x2, y2, x2, y2, x2, y1, x1, y1 };

    GLint posLoc = glGetAttribLocation(sidebar_program, "position");
    GLint colLoc = glGetUniformLocation(sidebar_program, "u_color");
    
    glUseProgram(sidebar_program);
    glUniform3f(colLoc, r, g, b);
    glVertexAttribPointer(posLoc, 2, GL_FLOAT, GL_FALSE, 0, vertices);
    glEnableVertexAttribArray(posLoc);
    glDrawArrays(GL_TRIANGLES, 0, 6);
}

void render_frame() {
    if (!running) return;

    // 1. Sidebar BG
    glViewport(0, 0, width, height);
    glUseProgram(sidebar_program);
    draw_quad(0, 0, SIDEBAR_WIDTH, height, width, height, 0.2f, 0.2f, 0.2f);

    // 2. Buttons
    for (size_t i = 0; i < shaders.size(); ++i) {
        float r = shaders[i].color[0];
        float g = shaders[i].color[1];
        float b = shaders[i].color[2];
        
        if ((int)i == current_shader_index) {
            r = std::min(r + 0.3f, 1.0f);
            g = std::min(g + 0.3f, 1.0f);
            b = std::min(b + 0.3f, 1.0f);
        }

        float by = BUTTON_MARGIN + i * (BUTTON_HEIGHT + BUTTON_MARGIN);
        
        // Quad
        glUseProgram(sidebar_program);
        draw_quad(BUTTON_MARGIN, by, SIDEBAR_WIDTH - 2 * BUTTON_MARGIN, BUTTON_HEIGHT, width, height, r, g, b);

        // Text
        float text_scale = 2.0f;
        float text_h = 8 * text_scale;
        float text_y = by + (BUTTON_HEIGHT - text_h) / 2;
        float text_x = BUTTON_MARGIN + 10;
        draw_text(shaders[i].name, text_x, text_y, text_scale, 0.0f, 0.0f, 0.0f);
    }

    // 3. Shader
    glViewport(SIDEBAR_WIDTH, 0, width - SIDEBAR_WIDTH, height);
    
    ShaderProgram &active = shaders[current_shader_index];
    glUseProgram(active.program);

    auto now = std::chrono::high_resolution_clock::now();
    float time = std::chrono::duration<float>(now - start_time).count();
    
    if (active.timeUniform != -1) glUniform1f(active.timeUniform, time);
    if (active.resolutionUniform != -1) glUniform2f(active.resolutionUniform, (float)(width - SIDEBAR_WIDTH), (float)height);

    GLfloat vertices[] = { -1, 1, -1, -1, 1, -1, 1, -1, 1, 1, -1, 1 };
    glVertexAttribPointer(active.posAttrib, 2, GL_FLOAT, GL_FALSE, 0, vertices);
    glEnableVertexAttribArray(active.posAttrib);
    glDrawArrays(GL_TRIANGLES, 0, 6);

    struct wl_callback *callback = wl_surface_frame(surface);
    wl_callback_add_listener(callback, &frame_listener, nullptr);

    eglSwapBuffers(egl_display, egl_surface);
}

// --- Init & Main ---

void init_shaders() {
    const char* sidebar_vert = 
        "#version 100\n"
        "attribute vec2 position;\n"
        "void main() { gl_Position = vec4(position, 0.0, 1.0); }";
    const char* sidebar_frag = 
        "#version 100\n"
        "precision mediump float;\n"
        "uniform vec3 u_color;\n"
        "void main() { gl_FragColor = vec4(u_color, 1.0); }";
    
    GLuint sv = load_shader_src(sidebar_vert, GL_VERTEX_SHADER);
    GLuint sf = load_shader_src(sidebar_frag, GL_FRAGMENT_SHADER);
    sidebar_program = glCreateProgram();
    glAttachShader(sidebar_program, sv);
    glAttachShader(sidebar_program, sf);
    glLinkProgram(sidebar_program);

    auto add_shader = [&](const char* name, const char* fpath, float r, float g, float b) {
        GLuint prog = create_program("shaders/vert.glsl", fpath);
        if (prog) {
            ShaderProgram sp;
            sp.program = prog;
            sp.name = name; // Now used for text!
            sp.posAttrib = glGetAttribLocation(prog, "position");
            sp.timeUniform = glGetUniformLocation(prog, "u_time");
            sp.resolutionUniform = glGetUniformLocation(prog, "u_resolution");
            sp.color[0] = r; sp.color[1] = g; sp.color[2] = b;
            shaders.push_back(sp);
        }
    };

    add_shader("PLASMA", "shaders/plasma.frag", 0.8f, 0.2f, 0.2f);
    add_shader("TUNNEL", "shaders/tunnel.frag", 0.2f, 0.8f, 0.2f);
    add_shader("FRACTAL", "shaders/fractal.frag", 0.2f, 0.2f, 0.8f);
    add_shader("MANDEL", "shaders/mandelbrot.frag", 0.5f, 0.0f, 0.5f);
}

int main() {
    display = wl_display_connect(nullptr);
    if (!display) return 1;
    registry = wl_display_get_registry(display);
    wl_registry_add_listener(registry, &registry_listener, nullptr);
    wl_display_roundtrip(display);
    
    egl_display = eglGetDisplay((EGLNativeDisplayType)display);
    eglInitialize(egl_display, nullptr, nullptr);
    EGLint attribs[] = { EGL_RED_SIZE, 8, EGL_GREEN_SIZE, 8, EGL_BLUE_SIZE, 8, EGL_RENDERABLE_TYPE, EGL_OPENGL_ES2_BIT, EGL_NONE };
    EGLConfig config; EGLint num; eglChooseConfig(egl_display, attribs, &config, 1, &num);
    EGLint ctx_attribs[] = { EGL_CONTEXT_CLIENT_VERSION, 2, EGL_NONE };
    egl_context = eglCreateContext(egl_display, config, EGL_NO_CONTEXT, ctx_attribs);

    surface = wl_compositor_create_surface(compositor);
    xdg_surface = xdg_wm_base_get_xdg_surface(wm_base, surface);
    xdg_surface_add_listener(xdg_surface, &xdg_surface_listener, nullptr);
    xdg_toplevel = xdg_surface_get_toplevel(xdg_surface);
    xdg_toplevel_add_listener(xdg_toplevel, &xdg_toplevel_listener, nullptr);
    xdg_toplevel_set_title(xdg_toplevel, "Shader Studio");
    wl_surface_commit(surface);
    wl_display_roundtrip(display);

    egl_window = wl_egl_window_create(surface, width, height);
    egl_surface = eglCreateWindowSurface(egl_display, config, (EGLNativeWindowType)egl_window, nullptr);
    eglMakeCurrent(egl_display, egl_surface, egl_surface, egl_context);

    // Initialize Shaders & Font
    init_font();
    init_text_shader();
    init_shaders();
    
    render_frame();

    while (running && wl_display_dispatch(display) != -1) {}

    return 0;
}