#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <EGL/egl.h>
#include <GLES2/gl2.h>

#include "../include/chroma.h"
#include "../include/stb_image.h"

// Vertex shader for simple texture rendering
static const char *vertex_shader_source =
    "#version 120\n"
    "attribute vec2 position;\n"
    "attribute vec2 texcoord;\n"
    "varying vec2 v_texcoord;\n"
    "void main() {\n"
    "    gl_Position = vec4(position, 0.0, 1.0);\n"
    "    v_texcoord = texcoord;\n"
    "}\n";

// Fragment shader for simple texture rendering
static const char *fragment_shader_source =
    "#version 120\n"
    "varying vec2 v_texcoord;\n"
    "uniform sampler2D texture;\n"
    "void main() {\n"
    "    gl_FragColor = texture2D(texture, v_texcoord);\n"
    "}\n";

// Vertices for a fullscreen quad
static const float vertices[] = {
    // Position  Texcoord
    -1.0f, -1.0f, 0.0f, 1.0f, // bottom left
    1.0f,  -1.0f, 1.0f, 1.0f, // bottom right
    1.0f,  1.0f,  1.0f, 0.0f, // top right
    -1.0f, 1.0f,  0.0f, 0.0f, // top left
};

static const unsigned int indices[] = {
    0, 1, 2, // first triangle
    2, 3, 0  // second triangle
};

static GLuint compile_shader(GLenum type, const char *source) {
  GLuint shader = glCreateShader(type);
  if (!shader) {
    chroma_log("ERROR", "Failed to create shader");
    return 0;
  }

  glShaderSource(shader, 1, &source, NULL);
  glCompileShader(shader);

  GLint status;
  glGetShaderiv(shader, GL_COMPILE_STATUS, &status);
  if (status != GL_TRUE) {
    char log[512];
    glGetShaderInfoLog(shader, sizeof(log), NULL, log);
    chroma_log("ERROR", "Shader compilation failed: %s", log);
    glDeleteShader(shader);
    return 0;
  }

  return shader;
}

static GLuint create_shader_program(void) {
  GLuint vertex_shader = compile_shader(GL_VERTEX_SHADER, vertex_shader_source);
  if (!vertex_shader) {
    return 0;
  }

  GLuint fragment_shader =
      compile_shader(GL_FRAGMENT_SHADER, fragment_shader_source);
  if (!fragment_shader) {
    glDeleteShader(vertex_shader);
    return 0;
  }

  GLuint program = glCreateProgram();
  if (!program) {
    chroma_log("ERROR", "Failed to create shader program");
    glDeleteShader(vertex_shader);
    glDeleteShader(fragment_shader);
    return 0;
  }

  glAttachShader(program, vertex_shader);
  glAttachShader(program, fragment_shader);
  glLinkProgram(program);

  GLint status;
  glGetProgramiv(program, GL_LINK_STATUS, &status);
  if (status != GL_TRUE) {
    char log[512];
    glGetProgramInfoLog(program, sizeof(log), NULL, log);
    chroma_log("ERROR", "Shader program linking failed: %s", log);
    glDeleteProgram(program);
    program = 0;
  }

  glDeleteShader(vertex_shader);
  glDeleteShader(fragment_shader);

  return program;
}

// Initialize OpenGL resources for output
static int init_gl_resources(chroma_output_t *output) {
  if (!output || output->gl_resources_initialized) {
    return CHROMA_OK;
  }

  // Create shader prog
  output->shader_program = create_shader_program();
  if (!output->shader_program) {
    chroma_log("ERROR", "Failed to create shader program for output %u",
               output->id);
    return CHROMA_ERROR_EGL;
  }

  // Create and setup VBO/EBO
  glGenBuffers(1, &output->vbo);
  glGenBuffers(1, &output->ebo);

  glBindBuffer(GL_ARRAY_BUFFER, output->vbo);
  glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);

  glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, output->ebo);
  glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(indices), indices,
               GL_STATIC_DRAW);

  output->texture_id = 0; // will be created when image is assigned
  output->gl_resources_initialized = true;

  chroma_log("DEBUG", "Initialized GL resources for output %u", output->id);
  return CHROMA_OK;
}

// Create or update texture from image data
static int update_texture_from_image(chroma_output_t *output,
                                     chroma_image_t *image) {
  if (!output || !image || !image->loaded) {
    return CHROMA_ERROR_INIT;
  }

  // If image data was already freed after previous GPU upload, we can't upload
  // again
  if (!image->data) {
    chroma_log("ERROR",
               "Cannot create texture: image data already freed for %s",
               image->path);
    return CHROMA_ERROR_IMAGE;
  }

  // Delete existing texture if it exists
  if (output->texture_id != 0) {
    glDeleteTextures(1, &output->texture_id);
    output->texture_id = 0;
  }

  // Create new texture
  glGenTextures(1, &output->texture_id);
  glBindTexture(GL_TEXTURE_2D, output->texture_id);

  // Set texture parameters
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);

  // Upload texture data (always RGBA now)
  glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, image->width, image->height, 0,
               GL_RGBA, GL_UNSIGNED_BYTE, image->data);

  glBindTexture(GL_TEXTURE_2D, 0);

  // Mark this output as having uploaded its texture
  output->texture_uploaded = true;

  // Free system RAM copy only when ALL outputs using this image have uploaded
  // to GPU
  if (image->data) {
    // Count total outputs using this image and how many have uploaded
    int total_using = 0;
    int uploaded_count = 0;

    chroma_state_t *state = output->state;
    for (int i = 0; i < state->output_count; i++) {
      if (state->outputs[i].active && state->outputs[i].image == image) {
        total_using++;
        if (state->outputs[i].texture_uploaded) {
          uploaded_count++;
        }
      }
    }

    // Only free image data when ALL outputs using it have uploaded
    if (total_using > 0 && uploaded_count >= total_using) {
      size_t freed_bytes =
          (size_t)image->width * image->height * image->channels;
      stbi_image_free(image->data);
      image->data = NULL;
      chroma_log("INFO",
                 "Freed %.2f MB of image data after all %d outputs uploaded to "
                 "GPU: %s",
                 (double)freed_bytes / (1024.0 * 1024.0), total_using,
                 image->path);
    }
  }

  chroma_log("DEBUG", "Updated texture for output %u (%dx%d)", output->id,
             image->width, image->height);
  return CHROMA_OK;
}

// Cleanup OpenGL resources for output
static void cleanup_gl_resources(chroma_output_t *output) {
  if (!output || !output->gl_resources_initialized) {
    return;
  }

  if (output->texture_id != 0) {
    glDeleteTextures(1, &output->texture_id);
    output->texture_id = 0;
  }

  if (output->shader_program != 0) {
    glDeleteProgram(output->shader_program);
    output->shader_program = 0;
  }

  if (output->vbo != 0) {
    glDeleteBuffers(1, &output->vbo);
    output->vbo = 0;
  }

  if (output->ebo != 0) {
    glDeleteBuffers(1, &output->ebo);
    output->ebo = 0;
  }

  output->gl_resources_initialized = false;
  chroma_log("DEBUG", "Cleaned up GL resources for output %u", output->id);
}

// EGL configuration selection
static int choose_egl_config(EGLDisplay display, EGLConfig *config) {
  EGLint attributes[] = {EGL_SURFACE_TYPE,
                         EGL_WINDOW_BIT,
                         EGL_RED_SIZE,
                         8,
                         EGL_GREEN_SIZE,
                         8,
                         EGL_BLUE_SIZE,
                         8,
                         EGL_ALPHA_SIZE,
                         8,
                         EGL_RENDERABLE_TYPE,
                         EGL_OPENGL_BIT,
                         EGL_NONE};

  EGLint num_configs;
  if (!eglChooseConfig(display, attributes, config, 1, &num_configs)) {
    chroma_log("ERROR", "Failed to choose EGL config: 0x%04x", eglGetError());
    return -1;
  }

  if (num_configs == 0) {
    chroma_log("ERROR", "No suitable EGL configs found");
    return -1;
  }

  return 0;
}

// EGL initialization
int chroma_egl_init(chroma_state_t *state) {
  if (!state || !state->display) {
    return CHROMA_ERROR_INIT;
  }

  // Get EGL display
  state->egl_display = eglGetDisplay((EGLNativeDisplayType)state->display);
  if (state->egl_display == EGL_NO_DISPLAY) {
    chroma_log("ERROR", "Failed to get EGL display: 0x%04x", eglGetError());
    return CHROMA_ERROR_EGL;
  }

  // Initialize EGL
  EGLint major, minor;
  if (!eglInitialize(state->egl_display, &major, &minor)) {
    chroma_log("ERROR", "Failed to initialize EGL: 0x%04x", eglGetError());
    return CHROMA_ERROR_EGL;
  }

  chroma_log("INFO", "EGL initialized: version %d.%d", major, minor);

  // Bind OpenGL API
  if (!eglBindAPI(EGL_OPENGL_API)) {
    chroma_log("ERROR", "Failed to bind OpenGL API: 0x%04x", eglGetError());
    chroma_egl_cleanup(state);
    return CHROMA_ERROR_EGL;
  }

  // Choose EGL config
  if (choose_egl_config(state->egl_display, &state->egl_config) != 0) {
    chroma_egl_cleanup(state);
    return CHROMA_ERROR_EGL;
  }

  // Create EGL context
  EGLint context_attributes[] = {EGL_CONTEXT_MAJOR_VERSION, 2,
                                 EGL_CONTEXT_MINOR_VERSION, 1, EGL_NONE};

  state->egl_context = eglCreateContext(state->egl_display, state->egl_config,
                                        EGL_NO_CONTEXT, context_attributes);
  if (state->egl_context == EGL_NO_CONTEXT) {
    chroma_log("ERROR", "Failed to create EGL context: 0x%04x", eglGetError());
    chroma_egl_cleanup(state);
    return CHROMA_ERROR_EGL;
  }

  chroma_log("INFO", "EGL context created successfully");
  return CHROMA_OK;
}

// EGL cleanup
void chroma_egl_cleanup(chroma_state_t *state) {
  if (!state) {
    return;
  }

  if (state->egl_display != EGL_NO_DISPLAY) {
    eglMakeCurrent(state->egl_display, EGL_NO_SURFACE, EGL_NO_SURFACE,
                   EGL_NO_CONTEXT);

    if (state->egl_context != EGL_NO_CONTEXT) {
      eglDestroyContext(state->egl_display, state->egl_context);
      state->egl_context = EGL_NO_CONTEXT;
    }

    eglTerminate(state->egl_display);
    state->egl_display = EGL_NO_DISPLAY;
  }

  chroma_log("INFO", "EGL cleaned up");
}

// Create surface for output
int chroma_surface_create(chroma_state_t *state, chroma_output_t *output) {
  if (!state || !output || !state->compositor || !state->layer_shell) {
    return CHROMA_ERROR_INIT;
  }

  // Create Wayland surface
  output->surface = wl_compositor_create_surface(state->compositor);
  if (!output->surface) {
    chroma_log("ERROR", "Failed to create Wayland surface for output %u",
               output->id);
    return CHROMA_ERROR_WAYLAND;
  }

  // Create layer surface for wallpaper
  output->layer_surface = zwlr_layer_shell_v1_get_layer_surface(
      state->layer_shell, output->surface, output->wl_output,
      ZWLR_LAYER_SHELL_V1_LAYER_BACKGROUND, "chroma-wallpaper");

  if (!output->layer_surface) {
    chroma_log("ERROR", "Failed to create layer surface for output %u",
               output->id);
    chroma_surface_destroy(output);
    return CHROMA_ERROR_WAYLAND;
  }

  // Configure layer surface
  zwlr_layer_surface_v1_set_size(output->layer_surface, output->width,
                                 output->height);
  zwlr_layer_surface_v1_set_anchor(output->layer_surface,
                                   ZWLR_LAYER_SURFACE_V1_ANCHOR_TOP |
                                       ZWLR_LAYER_SURFACE_V1_ANCHOR_RIGHT |
                                       ZWLR_LAYER_SURFACE_V1_ANCHOR_BOTTOM |
                                       ZWLR_LAYER_SURFACE_V1_ANCHOR_LEFT);
  zwlr_layer_surface_v1_set_exclusive_zone(output->layer_surface, -1);
  zwlr_layer_surface_v1_set_keyboard_interactivity(
      output->layer_surface, ZWLR_LAYER_SURFACE_V1_KEYBOARD_INTERACTIVITY_NONE);

  // Add layer surface listener
  zwlr_layer_surface_v1_add_listener(
      output->layer_surface, &chroma_layer_surface_listener_impl, output);

  // Commit surface to trigger configure event
  wl_surface_commit(output->surface);

  // Wait for configure event
  wl_display_roundtrip(state->display);

  // Create EGL window
  output->egl_window =
      wl_egl_window_create(output->surface, output->width, output->height);
  if (!output->egl_window) {
    chroma_log("ERROR", "Failed to create EGL window for output %u",
               output->id);
    chroma_surface_destroy(output);
    return CHROMA_ERROR_EGL;
  }

  // Create EGL surface
  output->egl_surface =
      eglCreateWindowSurface(state->egl_display, state->egl_config,
                             (EGLNativeWindowType)output->egl_window, NULL);
  if (output->egl_surface == EGL_NO_SURFACE) {
    chroma_log("ERROR", "Failed to create EGL surface for output %u: 0x%04x",
               output->id, eglGetError());
    chroma_surface_destroy(output);
    return CHROMA_ERROR_EGL;
  }

  chroma_log("INFO", "Created surface for output %u (%dx%d)", output->id,
             output->width, output->height);

  return CHROMA_OK;
}

// Destroy surface
void chroma_surface_destroy(chroma_output_t *output) {
  if (!output) {
    return;
  }

  // Clean up OpenGL resources first
  cleanup_gl_resources(output);

  if (output->egl_surface != EGL_NO_SURFACE) {
    eglDestroySurface(eglGetCurrentDisplay(), output->egl_surface);
    output->egl_surface = EGL_NO_SURFACE;
  }

  if (output->egl_window) {
    wl_egl_window_destroy(output->egl_window);
    output->egl_window = NULL;
  }

  if (output->layer_surface) {
    zwlr_layer_surface_v1_destroy(output->layer_surface);
    output->layer_surface = NULL;
  }

  if (output->surface) {
    wl_surface_destroy(output->surface);
    output->surface = NULL;
  }

  chroma_log("DEBUG", "Destroyed surface for output %u", output->id);
}

// Render wallpaper to output using cached resources
int chroma_render_wallpaper(chroma_state_t *state, chroma_output_t *output) {
  if (!state || !output || !output->image || !output->image->loaded) {
    return CHROMA_ERROR_INIT;
  }

  // Make context current
  if (!eglMakeCurrent(state->egl_display, output->egl_surface,
                      output->egl_surface, state->egl_context)) {
    chroma_log("ERROR", "Failed to make EGL context current: 0x%04x",
               eglGetError());
    return CHROMA_ERROR_EGL;
  }

  if (init_gl_resources(output) != CHROMA_OK) {
    return CHROMA_ERROR_EGL;
  }

  if (output->texture_id == 0) {
    if (update_texture_from_image(output, output->image) != CHROMA_OK) {
      chroma_log("ERROR", "Failed to update texture for output %u", output->id);
      return CHROMA_ERROR_EGL;
    }
  }

  // Set viewport
  glViewport(0, 0, output->width, output->height);

  // Clear screen
  glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
  glClear(GL_COLOR_BUFFER_BIT);

  // Use cached shader program
  glUseProgram(output->shader_program);

  // Bind cached texture
  glActiveTexture(GL_TEXTURE0);
  glBindTexture(GL_TEXTURE_2D, output->texture_id);
  glUniform1i(glGetUniformLocation(output->shader_program, "texture"), 0);

  // Use cached VBO/EBO
  glBindBuffer(GL_ARRAY_BUFFER, output->vbo);
  glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, output->ebo);

  // Set vertex attributes
  GLint position_attr = glGetAttribLocation(output->shader_program, "position");
  GLint texcoord_attr = glGetAttribLocation(output->shader_program, "texcoord");

  glEnableVertexAttribArray(position_attr);
  glVertexAttribPointer(position_attr, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float),
                        (void *)0);

  glEnableVertexAttribArray(texcoord_attr);
  glVertexAttribPointer(texcoord_attr, 2, GL_FLOAT, GL_FALSE, 4 * sizeof(float),
                        (void *)(2 * sizeof(float)));

  // Draw
  glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_INT, 0);

  // Unbind resources
  glBindBuffer(GL_ARRAY_BUFFER, 0);
  glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, 0);
  glBindTexture(GL_TEXTURE_2D, 0);
  glUseProgram(0);

  // Swap buffers
  if (!eglSwapBuffers(state->egl_display, output->egl_surface)) {
    chroma_log("ERROR", "Failed to swap buffers for output %u: 0x%04x",
               output->id, eglGetError());
    return CHROMA_ERROR_EGL;
  }

  // Commit surface
  wl_surface_commit(output->surface);

  chroma_log("DEBUG", "Rendered wallpaper to output %u (cached resources)",
             output->id);
  return CHROMA_OK;
}

// Invalidate texture cache for output
void chroma_output_invalidate_texture(chroma_output_t *output) {
  if (!output || !output->gl_resources_initialized) {
    return;
  }

  if (output->texture_id != 0) {
    glDeleteTextures(1, &output->texture_id);
    output->texture_id = 0;
    output->texture_uploaded = false; // reset upload flag
    chroma_log("DEBUG", "Invalidated texture cache for output %u", output->id);
  }
}