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

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

// GIF block types
#define GIF_EXT_INTRODUCER 0x21
#define GIF_GRAPHIC_CONTROL 0xF9
#define GIF_IMAGE_SEPARATOR 0x2C
#define GIF_TRAILER 0x3B

// Read a little-endian uint16 from a buffer
static uint16_t read_le16(const unsigned char *buf) {
  return (uint16_t)(buf[0] | (buf[1] << 8));
}

// Parse a GIF file to extract frame count and timing
static int gif_parse_header(const unsigned char *data, size_t size, int *width,
                            int *height, int *frame_count, int **frame_delays) {
  if (size < 13)
    return -1;

  // Verify GIF87a or GIF89a signature
  if (memcmp(data, "GIF87a", 6) != 0 && memcmp(data, "GIF89a", 6) != 0) {
    return -1;
  }

  *width = (int)read_le16(data + 6);
  *height = (int)read_le16(data + 8);

  // Skip logical screen descriptor fields
  size_t pos = 10;
  unsigned char packed = data[pos++];
  int has_gct = (packed & 0x80) >> 7;
  int bpp = (packed & 0x07) + 1;
  int gct_size = has_gct ? (3 << bpp) : 0;

  pos += 1; // background color index
  pos += 1; // pixel aspect ratio
  pos += (size_t)gct_size;

  // Count frames and parse frame delays
  int count = 0;
  int *delays = NULL;

  if (frame_count && frame_delays) {
    delays = malloc(1024 * sizeof(int));
    if (!delays)
      return -1;
  }

  while (pos < size) {
    unsigned char block_type = data[pos];

    if (block_type == GIF_TRAILER) {
      break;
    }

    if (block_type == GIF_EXT_INTRODUCER && pos + 1 < size) {
      unsigned char ext_type = data[pos + 1];
      pos += 2;

      if (ext_type == GIF_GRAPHIC_CONTROL && pos + 7 < size) {
        unsigned char block_size = data[pos];
        pos += 1 + (size_t)block_size + 1;

        if (delays && count < 1024) {
          // GCE block: [size:1][flags:1][delay:2][transparent:1][terminator:1]
          // pos has moved past the block, so data[pos - block_size] is the
          // start
          uint16_t delay_cs = read_le16(data + pos - block_size);
          delays[count] = (int)delay_cs * 10;
        }
      } else {
        // Skip other extensions
        while (pos < size && data[pos] != 0x00) {
          unsigned char sub_size = data[pos];
          pos += 1 + sub_size;
        }
        if (pos < size && data[pos] == 0x00) {
          pos++;
        }
      }
    } else if (block_type == GIF_IMAGE_SEPARATOR) {
      count++;
      pos += 10; // skip image descriptor header

      // Skip local color table if present
      unsigned char img_packed = data[pos - 1];
      int has_lct = (img_packed & 0x80) >> 7;
      int img_bpp = (img_packed & 0x07) + 1;
      int lct_size = has_lct ? (3 << img_bpp) : 0;
      pos += (size_t)lct_size;

      // Skip LZW minimum code size
      pos += 1;

      // Skip image data sub-blocks
      while (pos < size && data[pos] != 0x00) {
        unsigned char sub_size = data[pos];
        pos += 1 + sub_size;
      }
      if (pos < size && data[pos] == 0x00) {
        pos++;
      }
    } else {
      pos++;
    }
  }

  if (delays) {
    *frame_count = count;
    *frame_delays = delays;
  }

  return 0;
}

// Load an animated GIF into chroma_image_t
// Uses stb_image to decode individual frames
int chroma_anim_load_gif(chroma_image_t *image, const char *path) {
  if (!image || !path) {
    return CHROMA_ERROR_INIT;
  }

  // Read the entire GIF file into memory
  FILE *fp = fopen(path, "rb");
  if (!fp) {
    chroma_log("ERROR", "Cannot open GIF file: %s", path);
    return CHROMA_ERROR_IMAGE;
  }

  fseek(fp, 0, SEEK_END);
  long file_size = ftell(fp);
  fseek(fp, 0, SEEK_SET);

  if (file_size <= 0 || file_size > 100 * 1024 * 1024) {
    chroma_log("ERROR", "GIF file too large or invalid: %s", path);
    fclose(fp);
    return CHROMA_ERROR_IMAGE;
  }

  unsigned char *gif_data = malloc((size_t)file_size);
  if (!gif_data) {
    fclose(fp);
    return CHROMA_ERROR_MEMORY;
  }

  if (fread(gif_data, 1, (size_t)file_size, fp) != (size_t)file_size) {
    chroma_log("ERROR", "Failed to read GIF file: %s", path);
    free(gif_data);
    fclose(fp);
    return CHROMA_ERROR_IMAGE;
  }
  fclose(fp);

  int gif_width, gif_height, frame_count;
  int *frame_delays = NULL;

  if (gif_parse_header(gif_data, (size_t)file_size, &gif_width, &gif_height,
                       &frame_count, &frame_delays) != 0) {
    chroma_log("ERROR", "Failed to parse GIF header: %s", path);
    free(gif_data);
    return CHROMA_ERROR_IMAGE;
  }

  if (frame_count < 2) {
    // Single-frame GIF - use stb_image to load as regular image
    free(gif_data);
    free(frame_delays);
    return 1; // signal to caller: not animated
  }

  chroma_log("DEBUG", "Animated GIF: %dx%d, %d frames", gif_width, gif_height,
             frame_count);

  // Load frames using stb_image
  int *delays_stbi = NULL;
  int frame_w, frame_h, comp;
  unsigned char *frames_data =
      stbi_load_gif_from_memory(gif_data, (int)file_size, &delays_stbi,
                                &frame_w, &frame_h, &frame_count, &comp, 4);

  free(gif_data);

  if (!frames_data || frame_count < 2) {
    chroma_log("ERROR", "Failed to decode GIF frames: %s", path);
    if (frames_data)
      stbi_image_free(frames_data);
    free(frame_delays);
    if (delays_stbi)
      stbi_image_free(delays_stbi);
    return CHROMA_ERROR_IMAGE;
  }

  image->is_animated = true;
  image->frame_count = frame_count;
  image->current_frame = 0;
  image->width = frame_w;
  image->height = frame_h;
  image->channels = 4;

  image->frames = malloc((size_t)frame_count * sizeof(unsigned char *));
  image->frame_delays_ms = malloc((size_t)frame_count * sizeof(int));
  image->frame_widths = malloc((size_t)frame_count * sizeof(int));
  image->frame_heights = malloc((size_t)frame_count * sizeof(int));

  if (!image->frames || !image->frame_delays_ms || !image->frame_widths ||
      !image->frame_heights) {
    free(image->frames);
    free(image->frame_delays_ms);
    free(image->frame_widths);
    free(image->frame_heights);
    stbi_image_free(frames_data);
    free(frame_delays);
    if (delays_stbi)
      stbi_image_free(delays_stbi);
    return CHROMA_ERROR_MEMORY;
  }

  size_t frame_size = (size_t)frame_w * (size_t)frame_h * 4;

  for (int i = 0; i < frame_count; i++) {
    image->frames[i] = malloc(frame_size);
    if (!image->frames[i]) {
      for (int j = 0; j < i; j++)
        free(image->frames[j]);
      free(image->frames);
      free(image->frame_delays_ms);
      free(image->frame_widths);
      free(image->frame_heights);
      free(frames_data);
      free(frame_delays);
      if (delays_stbi)
        stbi_image_free(delays_stbi);
      return CHROMA_ERROR_MEMORY;
    }
    memcpy(image->frames[i], frames_data + (size_t)i * frame_size, frame_size);
    image->frame_delays_ms[i] =
        delays_stbi ? delays_stbi[i] : (frame_delays ? frame_delays[i] : 100);
    image->frame_widths[i] = frame_w;
    image->frame_heights[i] = frame_h;
  }

  stbi_image_free(frames_data);
  free(frame_delays);
  if (delays_stbi)
    stbi_image_free(delays_stbi);

  // Set initial data to first frame
  image->data = image->frames[0];
  image->loaded = true;
  image->ref_count = 1;

  chroma_log("INFO", "Loaded animated GIF: %s (%d frames, %dx%d)", path,
             frame_count, frame_w, frame_h);

  return CHROMA_OK;
}

// Advance to the next frame if enough time has passed
void chroma_anim_advance_frame(chroma_image_t *image,
                               long long current_time_ms) {
  if (!image || !image->is_animated || image->frame_count < 2) {
    return;
  }

  int delay = image->frame_delays_ms[image->current_frame];
  if (delay < 16)
    delay = 100;

  if (image->last_frame_ms == 0) {
    image->last_frame_ms = current_time_ms;
    return;
  }

  long long elapsed = current_time_ms - image->last_frame_ms;
  if (elapsed < (long long)delay) {
    return;
  }

  image->current_frame = (image->current_frame + 1) % image->frame_count;
  image->data = image->frames[image->current_frame];
  image->last_frame_ms = current_time_ms;
}

// Free animated frame data
void chroma_anim_free_frames(chroma_image_t *image) {
  if (!image || !image->is_animated) {
    return;
  }

  for (int i = 0; i < image->frame_count; i++) {
    free(image->frames[i]);
  }
  free(image->frames);
  free(image->frame_delays_ms);
  free(image->frame_widths);
  free(image->frame_heights);

  image->frames = NULL;
  image->frame_delays_ms = NULL;
  image->frame_widths = NULL;
  image->frame_heights = NULL;
  image->frame_count = 0;
  image->is_animated = false;
}

// Check if any output has an animated image that needs frame advance
bool chroma_anim_needs_redraw(chroma_state_t *state) {
  if (!state || !state->config.animation_enabled) {
    return false;
  }

  for (int i = 0; i < state->output_count; i++) {
    chroma_output_t *output = &state->outputs[i];
    if (output->active && output->image && output->image->is_animated) {
      return true;
    }
  }

  return false;
}

// Advance animation frames and re-render
int chroma_anim_update_outputs(chroma_state_t *state) {
  if (!state || !state->config.animation_enabled) {
    return CHROMA_OK;
  }

  long long now = chroma_get_time_ms();
  int updated = 0;

  for (int i = 0; i < state->output_count; i++) {
    chroma_output_t *output = &state->outputs[i];
    if (!output->active || !output->image || !output->image->is_animated) {
      continue;
    }

    if (!output->surface) {
      continue;
    }

    int old_frame = output->image->current_frame;
    chroma_anim_advance_frame(output->image, now);

    if (output->image->current_frame != old_frame) {
      int ret = chroma_render_wallpaper(state, output);
      if (ret == CHROMA_OK) {
        updated++;
      }
    }
  }

  if (updated > 0) {
    chroma_log("DEBUG", "Animation frames advanced: %d outputs", updated);
  }

  state->last_anim_ms = now;
  return CHROMA_OK;
}