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

#include <raylib.h>

#include <tiffio.h>

#include <lib/lib.h>
#include <lib/bool.h>
#include <lib/dir.h>
#include <lib/time.h>
#include <lib/color.h>
#include <lib/algo/flood_fill.h>

#define TAG 50839
#define OFFSET 16
#define ONE_ROW 1

char* full_path(char* dir, char* file)
{
  char* fpath = NULL;
  size_t dir_len = strlen(dir);
  size_t file_len = strlen(file);
  fpath = (char*)calloc(dir_len+file_len+2,sizeof(char));
  strcpy(fpath,dir);
  strcpy(fpath+dir_len+1,file);
  fpath[dir_len] = '/';
  return fpath;
}

bool_t is_tif_ext(char* file_name)
{
  size_t file_len = strlen(file_name);
  if ((file_name[file_len-3] == 't') && (file_name[file_len-2] == 'i') && (file_name[file_len-1] == 'f')) {
    return TRUE;
  }
  return FALSE;
}

void combine_masks(uint16_t *destination, uint16_t *extra_labels, size_t width, size_t height)
{
  if (destination == NULL) {
    destination = (uint16_t*)calloc(width*height, sizeof(uint16_t));
  }
  for (size_t y = 0; y < height; y++) {
    for (size_t x = 0; x < width; x++) {
      size_t coord = x + y*width;
      if (destination[coord] == 0) {
	destination[coord] = extra_labels[coord];
      }
    }
  }
}

uint16_t* tif_to_labels(char* tif_file_name, size_t *width, size_t *height)
{
  //-----------------------------------------------
  //-TIFF-IMAGE-OPEN-------------------------------
  //-----------------------------------------------
  TIFF *tif = TIFFOpen(tif_file_name, "r");
  if (!tif) {
    fprintf(stderr, "Failed to open TIFF file\n");
    return NULL;
  }
  //-----------------------------------------------

  //-----------------------------------------------
  //-TIFF-FIND-DIMENSIONS--------------------------
  //-----------------------------------------------
  size_t channels = 1;
  TIFFGetField(tif, TIFFTAG_IMAGEWIDTH, width);
  TIFFGetField(tif, TIFFTAG_IMAGELENGTH, height);
  //-----------------------------------------------
  tmsize_t STRIP_LENGTH = TIFFStripSize(tif);
  tmsize_t STRIP_COUNT = TIFFNumberOfStrips(tif);
  if ((*width)*(*height)*3 == STRIP_LENGTH*STRIP_COUNT) {
    channels = 3;
  } else if ((*width)*(*height)*4 == STRIP_LENGTH*STRIP_COUNT) {
    channels = 4;
  }
  //-----------------------------------------------

  //-----------------------------------------------
  //-TIFF-LOAD-DATA--------------------------------
  //-----------------------------------------------
  void* buffer = malloc(STRIP_LENGTH*sizeof(uint8_t));
  uint8_t* image_data = calloc((*width)*(*height)*channels,sizeof(uint8_t));
  if (image_data == NULL) {
    fprintf(stderr, "Memory allocation error\n");
    TIFFClose(tif);
    return NULL;
  }
  for (size_t y = 0; y < STRIP_COUNT; y++) {
    tmsize_t strip_size = TIFFReadRawStrip(tif, y, buffer, STRIP_LENGTH);
    assert(strip_size == STRIP_LENGTH);
    for (size_t x = 0; x < STRIP_LENGTH; x++) {
      image_data[x + y*STRIP_LENGTH] = ((uint8_t*)buffer)[x];
    }
  }
  free(buffer);
  //-----------------------------------------------

  //-----------------------------------------------
  //-FLOOD-FILL-SEGMENTATION-----------------------
  //-CONTIGUOUS-REGION-FINDING---------------------
  //-----------------------------------------------
  uint16_t starting_label = 1;
  uint16_t *labels = NULL;
  labels = (uint16_t*)calloc((*width)*(*height),sizeof(uint16_t));
  if (labels == NULL) {
    fprintf(stderr, "Memory allocation error\n");
    free(image_data);
    TIFFClose(tif);
    return NULL;
  }
  // Flood fill on each pixel
  //  Increase label for each success
  for (size_t y = 0; y < *height; y++) {
    for (size_t x = 0; x < *width; x++) {
      size_t coord = x + y*(*width);
      if(flood(image_data, labels, *width, *height, channels, x, y, &(image_data[coord*channels]), starting_label)) {
        starting_label += 1;
      }
    }
  }
  free(image_data);
  TIFFClose(tif);
  return labels;
}

int main(int argc, char** argv)
{
  //-----------------------------------------------
  //-LIST-FILES-IN-DIRECTORY-----------------------
  //-----------------------------------------------
  char** file_list = NULL;
  uint16_t *masks = NULL;
  if (argc > 1) {
    if (is_directory(argv[1])) {
      file_list = list_directory(argv[1]);
      if (file_list) {
	size_t index = 0;
	while (1) {
	  char* fname = file_list[index];
	  if (fname == NULL) {
	    break;
	  }
	  if (is_tif_ext(fname) == FALSE) {
	    free(file_list[index++]);
	    continue;
	  } else {
	    printf("Tiff found!\n");
	  }
	  char* fpath = full_path(argv[1], fname);
	  size_t width, height;
	  uint16_t *file_labels = tif_to_labels(fpath, &width, &height);
	  combine_masks(masks, file_labels, width, height);
	  free(file_labels);
	  printf("%s\n", fpath);
	  free(fpath);
	  printf("%s\n", fname);
	  free(file_list[index++]);
	}
	free(file_list);
      }
    }
  }
  free(masks);
  //-----------------------------------------------

  //-----------------------------------------------
  //-RAYLIB-INIT
  //-----------------------------------------------
  SetTraceLogLevel(LOG_ERROR);
  SetConfigFlags(FLAG_WINDOW_RESIZABLE);
  const char* gui_title = "Image Manip - Useful for segmentations!";
  InitWindow(SCREEN_WIDTH, SCREEN_HEIGHT, gui_title);
  //-----------------------------------------------

  //-----------------------------------------------
  //-TIFF-IMAGE-OPEN-------------------------------
  //-----------------------------------------------
  TIFF *tif = TIFFOpen("./data/test.tif", "r");
  if (!tif) {
    fprintf(stderr, "Failed to open TIFF file\n");
    return 1;
  }
  //-----------------------------------------------

  //-----------------------------------------------
  //-TIFF-FIND-DIMENSIONS--------------------------
  //-----------------------------------------------
  uint32_t width, height;
  size_t channels = 1;
  TIFFGetField(tif, TIFFTAG_IMAGEWIDTH, &width);
  TIFFGetField(tif, TIFFTAG_IMAGELENGTH, &height);
  //-----------------------------------------------
  tmsize_t STRIP_LENGTH = TIFFStripSize(tif);
  tmsize_t STRIP_COUNT = TIFFNumberOfStrips(tif);
  if (width*height*3 == STRIP_LENGTH*STRIP_COUNT) {
    channels = 3;
  } else if (width*height*4 == STRIP_LENGTH*STRIP_COUNT) {
    channels = 4;
  }
  //-----------------------------------------------

  printf("Total Strip Size: %llu\n", STRIP_LENGTH*STRIP_COUNT);
  printf("Total Image Size: %lu  in %hhu channel(s)\n", width*height*channels, channels);
  assert(STRIP_LENGTH*STRIP_COUNT == width*height*channels);

  //-----------------------------------------------
  //-TIFF-LOAD-DATA--------------------------------
  //-----------------------------------------------
  void* buffer = malloc(STRIP_LENGTH*sizeof(uint8_t));
  uint8_t* raster_8 = calloc(width*height*channels,sizeof(uint8_t));
  if (raster_8 == NULL) {
    fprintf(stderr, "Memory allocation error\n");
    TIFFClose(tif);
    return 1;
  }
  for (size_t y = 0; y < STRIP_COUNT; y++) {
    tmsize_t strip_size = TIFFReadRawStrip(tif, y, buffer, STRIP_LENGTH);
    assert(strip_size == STRIP_LENGTH);
    for (size_t x = 0; x < STRIP_LENGTH; x++) {
      raster_8[x + y*STRIP_LENGTH] = ((uint8_t*)buffer)[x];
    }
  }
  free(buffer);
  //-----------------------------------------------

  struct timespec ts_start, ts_end;
  get_time(&ts_start);
  //-----------------------------------------------
  //-FLOOD-FILL-SEGMENTATION-----------------------
  //-CONTIGUOUS-REGION-FINDING---------------------
  //-----------------------------------------------
  uint16_t starting_label = 1;
  uint16_t *labels = NULL;
  labels = (uint16_t*)calloc(width*height,sizeof(uint16_t));
  if (labels == NULL) {
    fprintf(stderr, "Memory allocation error\n");
    free(raster_8);
    TIFFClose(tif);
    return 1;
  }
  // Flood fill on each pixel
  //  Increase label for each success
  for (size_t y = 0; y < height; y++) {
    for (size_t x = 0; x < width; x++) {
      size_t coord = x + y*width;
      if(flood(raster_8, labels, width, height, channels, x, y, &(raster_8[coord*channels]), starting_label)) {
        starting_label += 1;
      }
    }
  }
  // (When treating this as RGBA, last bits should ensure opaque)
  //  This assumes 4096 (2^12) > labels
  for (size_t y = 0; y < height; y++) {
    for (size_t x = 0; x < width; x++) {
      labels[x + y*width] |= 0xFF000000;
    }
  }
  //-----------------------------------------------
  get_time(&ts_end);
  printf("Time difference: %.3fms\n", 1000*diff_time(&ts_end, &ts_start));
  printf("N_labels: %u\n", starting_label-1);

  uint32_t *raster = NULL;
  //-----------------------------------------------
  //-READ-TIFF-IMAGE-------------------------------
  //-----------------------------------------------
  /*
  raster = (uint32_t*)_TIFFmalloc(width*height*sizeof(uint32_t));
  if (raster == NULL) {
    fprintf(stderr, "Memory allocation error\n");
    free(labels);
    free(raster_8);
    TIFFClose(tif);
    return 1;
  }

  if (!TIFFReadRGBAImage(tif, width, height, raster, 0)) {
    fprintf(stderr, "Failed to read TIFF image\n");
    free(labels);
    free(raster_8);
    _TIFFfree(raster);
    TIFFClose(tif);
    return 1;
  }
  */
  //-----------------------------------------------

  //-----------------------------------------------
  //-RAYLIB-IMAGE-STRUCTURING----------------------
  //-----------------------------------------------
  Image RaylibImage;
  RaylibImage.width  = width;
  RaylibImage.height = height;
  RaylibImage.mipmaps = 1;
  // Use Loaded Image
  RaylibImage.data = raster;
  RaylibImage.format = PIXELFORMAT_UNCOMPRESSED_R8G8B8A8;
  // Use Strip Image Data
  RaylibImage.data = raster_8;
  if (channels == 1) {
    RaylibImage.format = PIXELFORMAT_UNCOMPRESSED_GRAYSCALE;
  } else if (channels == 3) {
    RaylibImage.format = PIXELFORMAT_UNCOMPRESSED_R8G8B8;
  } else if (channels == 4) {
    RaylibImage.format = PIXELFORMAT_UNCOMPRESSED_R8G8B8A8;
  } else {
    assert(0);
  }
  // Use Contiguous Labels
  RaylibImage.data = labels;
  RaylibImage.format = PIXELFORMAT_UNCOMPRESSED_R4G4B4A4;
  //-----------------------------------------------

  // Image to a Texture
  Texture2D RaylibTexture = LoadTextureFromImage(RaylibImage);

  // Scale the image to the viewport
  /// Source Rectangle: Original Size
  Rectangle sourceRec = { 0.0f, 0.0f, (float)width, (float)height };
  /// Destination Rectangle: Transformed Size
  Rectangle destRec = { 0.0f, 0.0f, (float)SCREEN_WIDTH, (float)(SCREEN_HEIGHT-OFFSET) };
  /// Location to begin drawing
  Vector2 origin = { (float)0, (float)-OFFSET };

  // Raylib boilerplate
  SetTargetFPS(60);
  Camera2D camera = { 0 };
  camera.zoom = 1.0f;

  // GUI Loop
  while (!WindowShouldClose()) {
    //-----------------------------------------------
    //-DRAWING---------------------------------------
    //-----------------------------------------------
    BeginDrawing();
    ClearBackground(RAYWHITE);
    BeginMode2D(camera);
    EndMode2D();
    DrawText("Image Manip", 0, 0, OFFSET, DARKGRAY);
    DrawTexturePro(RaylibTexture, sourceRec, destRec, origin, (float)0,
                   RAYWHITE);
    /*
    uint32_t x = 0x49, y = 0x4A;
    uint32_t dx = 0x69 - x, dy = 0x6E - y;
    x = (SCREEN_WIDTH*x)/width;
    y = SCREEN_HEIGHT-((SCREEN_HEIGHT-OFFSET)*y)/height;
    dx = (SCREEN_WIDTH*dx)/width;
    dy = SCREEN_HEIGHT-((SCREEN_HEIGHT-OFFSET)*dy)/height;
    DrawRectangleGradientH(x, y, dx, dy, BLUE, PURPLE);
    */
    EndDrawing();
    //-----------------------------------------------
  }

  if (labels != NULL) {
    free(labels);
  }
  if (raster_8 != NULL) {
    free(raster_8);
  }
  if (raster != NULL) {
    _TIFFfree(raster);
  }
  TIFFClose(tif);

  CloseWindow();
  return 0;
}