#include "swill.h"
#include <math.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>

#define SIXEL_START_SEQUENCE "\ePq"
#define SIXEL_END_SEQUENCE "\e\\"

#define SIXEL_SUCCESS 0
const double BYTE_SCALE_FACTOR = 0.39215686;

typedef enum sort_by { SORT_RED, SORT_GREEN, SORT_BLUE } sort_by_t;

typedef struct bucket {
  unsigned int bufsize;
  unsigned int len;
  uint8_t low_r;
  uint8_t high_r;
  uint8_t low_g;
  uint8_t high_g;
  uint8_t low_b;
  uint8_t high_b;
  uint8_t *pixels;
} bucket_t;

bucket_t *new_bucket(unsigned int pixel_count) {
  bucket_t *output = malloc(sizeof(bucket_t));
  output->bufsize = pixel_count * 3; // Enough for a 256x256 image
  output->len = 0;
  output->low_r = 255;
  output->high_r = 0;
  output->low_g = 255;
  output->high_g = 0;
  output->low_b = 255;
  output->high_b = 0;
  output->pixels = malloc(sizeof(uint8_t[pixel_count * 3]));
  return output;
}

void free_bucket(bucket_t *bucket) {
  free(bucket->pixels);
  free(bucket);
}

void analyze_bucket(bucket_t *bucket) {
  for (unsigned int i = 0; i < bucket->len; i++) {
    uint8_t r = bucket->pixels[i * 3];
    uint8_t g = bucket->pixels[i * 3 + 1];
    uint8_t b = bucket->pixels[i * 3 + 2];

    if (r < bucket->low_r) {
      bucket->low_r = r;
    }
    if (bucket->high_r < r) {
      bucket->high_r = r;
    }
    if (g < bucket->low_g) {
      bucket->low_g = g;
    }
    if (bucket->high_g < g) {
      bucket->high_g = g;
    }
    if (b < bucket->low_b) {
      bucket->low_b = b;
    }
    if (bucket->high_b < b) {
      bucket->high_b = b;
    }
  }
}

void fill_bucket(bucket_t *bucket, uint8_t **pixel_data, unsigned int height,
                 unsigned int width, sort_by_t sort_by) {
  uint8_t sort_color;
  switch (sort_by) {
  case SORT_RED:
    sort_color = 0;
    break;
  case SORT_GREEN:
    sort_color = 1;
    break;
  case SORT_BLUE:
    sort_color = 2;
    break;
  }
  for (unsigned int y = 0; y < height; y++) {
    for (unsigned int x = 0; x < width; x++) {
      bucket->len++;

      int i;
      for (i = bucket->len; i > 0; i--) {
        if (bucket->len == 1) {
          i = 0;
          break;
        }
        int prev = (i - 1) * 3;
        if (pixel_data[y][x * 3 + sort_color] <=
            bucket->pixels[prev + sort_color]) {
          break;
        }
        // Copy RGB from previous to current; the pixel we're inserting goes
        // before it.
        bucket->pixels[i * 3] = bucket->pixels[prev];
        bucket->pixels[i * 3 + 1] = bucket->pixels[prev + 1];
        bucket->pixels[i * 3 + 2] = bucket->pixels[prev + 2];
      }
      // We've either found a suitable spot, or reached the beginning of the
      // list. Insert our pixel data.
      bucket->pixels[i * 3] = pixel_data[y][x * 3];
      bucket->pixels[i * 3 + 1] = pixel_data[y][x * 3 + 1];
      bucket->pixels[i * 3 + 2] = pixel_data[y][x * 3 + 2];
    }
  }
}

void swap_colors(uint8_t *a, uint8_t *b) {
  uint8_t buf;
  buf = a[0];
  a[0] = b[0];
  b[0] = buf;

  buf = a[1];
  a[1] = b[1];
  b[1] = buf;

  buf = a[2];
  a[2] = b[2];
  b[2] = buf;
}

void sift_down(uint8_t *array, unsigned int element, unsigned int bottom,
               sort_by_t sort_by) {
  uint8_t sort_color;
  switch (sort_by) {
  case SORT_RED:
    sort_color = 0;
    break;
  case SORT_GREEN:
    sort_color = 1;
    break;
  case SORT_BLUE:
    sort_color = 2;
    break;
  }
  while (2 * element <= bottom) {
    unsigned int lower_child;
    if (bottom < 2 * element + 1) { // Only one child in the heap, and this must
                                    // be the 2 * element = bottom case.
      lower_child = 2 * element;
    } else if (array[(2 * element) * 3 + sort_color] <
               array[(2 * element + 1) * 3 + sort_color]) {
      lower_child = 2 * element;
    } else {
      lower_child = 2 * element + 1;
    }

    if (array[lower_child * 3 + sort_color] > array[element * 3 + sort_color]) {
      return;
    } else {
      swap_colors(&array[lower_child * 3], &array[element * 3]);
      element = lower_child;
    }
  }
}

void sort_bucket(bucket_t *bucket, sort_by_t sort_by) {
  // To simplify heap arithmetic, we want the start of the array to be at index
  // 1, so we define our pointer as one element behind the actual base.
  uint8_t *base = bucket->pixels - 3;
  unsigned int bottom = bucket->len;

  // Heapify
  unsigned int current_node = bottom / 2;
  while (current_node > 0) {
    sift_down(base, current_node, bottom, sort_by);
    current_node--;
  }

  // Sort
  while (bottom > 1) {
    swap_colors(&base[bottom], &base[1]);
    bottom--;
    sift_down(base, 1, bottom, sort_by);
  }
}

int color_quantize(uint8_t **pixel_data, uint8_t **colors, int color_count,
                   int height, int width) {
  int bucket_count = 1;
  bucket_t *buckets[256];
  buckets[0] = new_bucket(height * width);

  for (int y = 0; y < height; y++) {
    for (int x = 0; x < width; x++) {
      uint8_t r = pixel_data[y][x * 3];
      uint8_t g = pixel_data[y][x * 3 + 1];
      uint8_t b = pixel_data[y][x * 3 + 2];

      if (r < buckets[0]->low_r) {
        buckets[0]->low_r = r;
      } else if (buckets[0]->high_r < r) {
        buckets[0]->high_r = r;
      }
      if (g < buckets[0]->low_g) {
        buckets[0]->low_g = g;
      } else if (buckets[0]->high_g < g) {
        buckets[0]->high_g = g;
      }
      if (b < buckets[0]->low_b) {
        buckets[0]->low_b = b;
      } else if (buckets[0]->high_b < b) {
        buckets[0]->high_b = b;
      }
    }
  }

  for (int i = 0; i < 256; i++) {
    free_bucket(buckets[i]);
  }
  return SIXEL_SUCCESS;
}

// Converts the `pixel_data` input (as given to `sixel_encode_bytes`) from
// 0-255 to 0-100. Caller is responsible for freeing `normalized` and its
// members. Used internally and is not part of the public API.
int normalize_bytes(uint8_t **pixel_data, uint8_t ***normalized,
                    bool full_color, unsigned int height, unsigned int width) {
  uint8_t pixel_width = full_color ? 3 : 1;
  *normalized = malloc(height * sizeof(**normalized));
  for (unsigned int i = 0; i < height; i++) {
    (*normalized)[i] = malloc(width * pixel_width * sizeof(**normalized));
  }

  for (unsigned int y = 0; y < height; y++) {
    for (unsigned int x = 0; x < width; x++) {
      for (int color = 0; color < pixel_width; color++) {
        uint8_t color_in = pixel_data[y][x * pixel_width + color];
        (*normalized)[y][x * pixel_width + color] =
            round(color_in * BYTE_SCALE_FACTOR);
      }
    }
  }
  return SIXEL_SUCCESS;
}

// Takes a pointer to a two-dimensional byte array `pixel_data`, converts it to
// sixel, and writes the result into `sixel`. Each color is assumed to be a byte
// wide, ranging from 0-255. (Sixel normalizes it to a percentage 0-100, so
// there's not really any point in additional precision.)
int sixel_encode_bytes(uint8_t ***pixel_data, char **sixel, bool full_color,
                       unsigned int height, unsigned int width) {
  // Convert 8-bit color data to percentages
  uint8_t **normalized;
  normalize_bytes(*pixel_data, &normalized, full_color, height, width);

  // Clean up allocations.
  for (unsigned int i = 0; i < height; i++) {
    free(normalized[i]);
  }
  free(normalized);
  return SIXEL_SUCCESS;
}