#include <errno.h>
#include <fcntl.h>
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

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

// Destroy shared memory buffer for an output
void chroma_buffer_destroy(chroma_output_t *output) {
  if (!output) {
    return;
  }

  if (output->shm_buffer) {
    wl_buffer_destroy(output->shm_buffer);
    output->shm_buffer = NULL;
  }

  if (output->shm_data) {
    if (munmap(output->shm_data, output->shm_size) < 0) {
      chroma_log("WARN", "munmap failed: %s", strerror(errno));
    }
    output->shm_data = NULL;
    output->shm_size = 0;
    output->shm_stride = 0;
  }
}

// Unmap client-side memory while keeping the wl_buffer proxy alive.
// The compositor retains its own fd reference and mapping, so the
// wallpaper remains visible after this call.
void chroma_buffer_unmap(chroma_output_t *output) {
  if (!output || !output->shm_data) {
    return;
  }

  if (munmap(output->shm_data, output->shm_size) < 0) {
    chroma_log("WARN", "munmap failed: %s", strerror(errno));
  }
  output->shm_data = NULL;
  output->shm_size = 0;
  output->shm_stride = 0;
}

// Create shared memory buffer for an output
int chroma_buffer_create(chroma_state_t *state, chroma_output_t *output) {
  if (!state || !output || !state->shm) {
    return CHROMA_ERROR_INIT;
  }

  int width = output->width;
  int height = output->height;
  if (width <= 0 || height <= 0) {
    return CHROMA_ERROR_INIT;
  }

  int stride = width * 4;
  size_t size = (size_t)stride * height;

  int fd = memfd_create("chroma-shm", MFD_CLOEXEC | MFD_ALLOW_SEALING);
  if (fd < 0) {
    // Fallback to shm_open if memfd_create is unavailable
    char name[64];
    snprintf(name, sizeof(name), "/chroma-shm-%d-%d", getpid(), rand());
    fd = shm_open(name, O_RDWR | O_CREAT | O_EXCL, 0600);
    if (fd >= 0) {
      shm_unlink(name);
    }
  }
  if (fd < 0) {
    chroma_log("ERROR", "Failed to create shared memory fd: %s",
               strerror(errno));
    return CHROMA_ERROR_MEMORY;
  }

  if (ftruncate(fd, (off_t)size) < 0) {
    chroma_log("ERROR", "ftruncate failed: %s", strerror(errno));
    close(fd);
    return CHROMA_ERROR_MEMORY;
  }

  void *data = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
  if (data == MAP_FAILED) {
    chroma_log("ERROR", "mmap failed: %s", strerror(errno));
    close(fd);
    return CHROMA_ERROR_MEMORY;
  }

  // Fill with black
  memset(data, 0, size);

  struct wl_shm_pool *pool = wl_shm_create_pool(state->shm, fd, (int32_t)size);
  if (!pool) {
    chroma_log("ERROR", "wl_shm_create_pool failed");
    munmap(data, size);
    close(fd);
    return CHROMA_ERROR_WAYLAND;
  }

  struct wl_buffer *buffer = wl_shm_pool_create_buffer(
      pool, 0, width, height, stride, WL_SHM_FORMAT_ARGB8888);
  wl_shm_pool_destroy(pool);
  close(fd);

  if (!buffer) {
    chroma_log("ERROR", "wl_shm_pool_create_buffer failed");
    munmap(data, size);
    return CHROMA_ERROR_WAYLAND;
  }

  output->shm_buffer = buffer;
  output->shm_data = data;
  output->shm_size = size;
  output->shm_stride = stride;

  return CHROMA_OK;
}

// Sample a pixel from image data (grayscale, grayscale+alpha, RGB, or RGBA)
static inline uint32_t sample_pixel(const chroma_image_t *image, int x, int y) {
  if (!image || !image->data || x < 0 || y < 0 || x >= image->width ||
      y >= image->height) {
    return 0xFF000000; // black with full alpha
  }

  int idx = (y * image->width + x) * image->channels;
  unsigned char r, g, b, a;

  switch (image->channels) {
  case 1:
    r = g = b = image->data[idx];
    a = 0xFF;
    break;
  case 2:
    r = g = b = image->data[idx];
    a = image->data[idx + 1];
    break;
  case 3:
    r = image->data[idx];
    g = image->data[idx + 1];
    b = image->data[idx + 2];
    a = 0xFF;
    break;
  default: // 4 or more
    r = image->data[idx];
    g = image->data[idx + 1];
    b = image->data[idx + 2];
    a = image->data[idx + 3];
    break;
  }

  return ((uint32_t)a << 24) | ((uint32_t)r << 16) | ((uint32_t)g << 8) |
         (uint32_t)b;
}

// Draw scaled image into the shared memory buffer
static void draw_scaled_image(chroma_output_t *output, chroma_image_t *image) {
  if (!output || !image || !output->shm_data) {
    return;
  }

  int out_w = output->width;
  int out_h = output->height;
  int img_w = image->width;
  int img_h = image->height;
  uint32_t *dst = (uint32_t *)output->shm_data;

  // Pre-fill with black
  memset(output->shm_data, 0, output->shm_size);

  float anchor_x = output->anchor_x;
  float anchor_y = output->anchor_y;

  switch (output->scale_mode) {
  case CHROMA_SCALE_STRETCH: {
    for (int y = 0; y < out_h; y++) {
      for (int x = 0; x < out_w; x++) {
        int sx = (int)((float)x * img_w / out_w);
        int sy = (int)((float)y * img_h / out_h);
        sx = (sx < 0) ? 0 : (sx >= img_w) ? img_w - 1 : sx;
        sy = (sy < 0) ? 0 : (sy >= img_h) ? img_h - 1 : sy;
        dst[y * out_w + x] = sample_pixel(image, sx, sy);
      }
    }
    break;
  }

  case CHROMA_SCALE_FILL: {
    float img_aspect = (float)img_w / (float)img_h;
    float out_aspect = (float)out_w / (float)out_h;

    float crop_w, crop_h, crop_x, crop_y;

    if (img_aspect > out_aspect) {
      // Image is wider - crop left/right
      crop_h = (float)img_h;
      crop_w = crop_h * out_aspect;
    } else {
      // Image is taller - crop top/bottom
      crop_w = (float)img_w;
      crop_h = crop_w / out_aspect;
    }

    crop_x = (anchor_x / 100.0f) * ((float)img_w - crop_w);
    crop_y = (anchor_y / 100.0f) * ((float)img_h - crop_h);

    // Clamp crop region
    if (crop_x < 0)
      crop_x = 0;
    if (crop_y < 0)
      crop_y = 0;
    if (crop_x > (float)img_w - crop_w)
      crop_x = (float)img_w - crop_w;
    if (crop_y > (float)img_h - crop_h)
      crop_y = (float)img_h - crop_h;

    for (int y = 0; y < out_h; y++) {
      for (int x = 0; x < out_w; x++) {
        float sx = crop_x + (float)x * crop_w / (float)out_w;
        float sy = crop_y + (float)y * crop_h / (float)out_h;
        int ix = (int)sx;
        int iy = (int)sy;
        dst[y * out_w + x] = sample_pixel(image, ix, iy);
      }
    }
    break;
  }

  case CHROMA_SCALE_FIT: {
    float scale_x = (float)out_w / (float)img_w;
    float scale_y = (float)out_h / (float)img_h;
    float scale = (scale_x < scale_y) ? scale_x : scale_y;

    float scaled_w = (float)img_w * scale;
    float scaled_h = (float)img_h * scale;

    float offset_x = (anchor_x / 100.0f) * ((float)out_w - scaled_w);
    float offset_y = (anchor_y / 100.0f) * ((float)out_h - scaled_h);

    int start_x = (int)offset_x;
    int start_y = (int)offset_y;
    int end_x = (int)(offset_x + scaled_w + 0.5f);
    int end_y = (int)(offset_y + scaled_h + 0.5f);

    if (start_x < 0)
      start_x = 0;
    if (start_y < 0)
      start_y = 0;
    if (end_x > out_w)
      end_x = out_w;
    if (end_y > out_h)
      end_y = out_h;

    for (int y = start_y; y < end_y; y++) {
      for (int x = start_x; x < end_x; x++) {
        float sx = ((float)x - offset_x) / scale;
        float sy = ((float)y - offset_y) / scale;
        int ix = (int)sx;
        int iy = (int)sy;
        ix = (ix < 0) ? 0 : (ix >= img_w) ? img_w - 1 : ix;
        iy = (iy < 0) ? 0 : (iy >= img_h) ? img_h - 1 : iy;
        dst[y * out_w + x] = sample_pixel(image, ix, iy);
      }
    }
    break;
  }

  case CHROMA_SCALE_CENTER:
  default: {
    if (img_w <= out_w && img_h <= out_h) {
      // Image fits entirely - position with anchor
      float offset_x = (anchor_x / 100.0f) * ((float)out_w - (float)img_w);
      float offset_y = (anchor_y / 100.0f) * ((float)out_h - (float)img_h);

      for (int y = 0; y < img_h; y++) {
        for (int x = 0; x < img_w; x++) {
          int dx = (int)(offset_x + x);
          int dy = (int)(offset_y + y);
          if (dx >= 0 && dx < out_w && dy >= 0 && dy < out_h) {
            dst[dy * out_w + dx] = sample_pixel(image, x, y);
          }
        }
      }
    } else {
      // Image is larger than output - clip with anchor
      float src_offset_x = (anchor_x / 100.0f) * ((float)img_w - (float)out_w);
      float src_offset_y = (anchor_y / 100.0f) * ((float)img_h - (float)out_h);

      if (src_offset_x < 0)
        src_offset_x = 0;
      if (src_offset_y < 0)
        src_offset_y = 0;
      if (src_offset_x > (float)img_w - (float)out_w)
        src_offset_x = (float)img_w - (float)out_w;
      if (src_offset_y > (float)img_h - (float)out_h)
        src_offset_y = (float)img_h - (float)out_h;

      for (int y = 0; y < out_h; y++) {
        for (int x = 0; x < out_w; x++) {
          int sx = (int)(src_offset_x + x);
          int sy = (int)(src_offset_y + y);
          dst[y * out_w + x] = sample_pixel(image, sx, sy);
        }
      }
    }
    break;
  }
  }
}

// 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, (uint32_t)output->width,
                                 (uint32_t)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);

  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;
  }

  chroma_buffer_destroy(output);

  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;
  }

  output->configured = false;

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

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

  if (!output->configured) {
    chroma_log("WARN", "Output %u not yet configured, skipping render",
               output->id);
    return CHROMA_ERROR_INIT;
  }

  // Clean up old buffer
  chroma_buffer_destroy(output);

  // Create new shared memory buffer
  int ret = chroma_buffer_create(state, output);
  if (ret != CHROMA_OK) {
    chroma_log("ERROR", "Failed to create shm buffer for output %u",
               output->id);
    return ret;
  }

  // Draw scaled image into buffer
  draw_scaled_image(output, output->image);

  // Attach buffer to surface
  wl_surface_attach(output->surface, output->shm_buffer, 0, 0);
  wl_surface_damage_buffer(output->surface, 0, 0, output->width,
                           output->height);
  wl_surface_commit(output->surface);
  if (wl_display_flush(state->display) < 0) {
    chroma_log("ERROR", "Failed to flush Wayland display for output %u: %s",
               output->id, strerror(errno));
    return CHROMA_ERROR_WAYLAND;
  }

  chroma_log("INFO", "Rendered wallpaper to output %u (%dx%d)", output->id,
             output->width, output->height);

  // Unmap client-side memory; compositor keeps its own fd reference
  chroma_buffer_unmap(output);

  // Release image reference and free CPU data if no longer needed
  if (output->image) {
    chroma_image_t *image = output->image;
    chroma_image_release(image);
    output->image = NULL;
  }

  return CHROMA_OK;
}