#include "tiffread.h"
#include "aliases.h"
#include "endianness.h"
#include "image.h"
#include "mem_arena.h"
#include <assert.h>
#include <math.h>
#include <netinet/in.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>
#include <string.h>

#define TIFF_MAGIC 0x002a

#define TIFF_FILENAME_MIN_LENGTH 6
#define TIFF_EXT_MAX_LENGTH 5
#define IS_TIFF_EXTENSION(EXT)                                                 \
  (strncmp(EXT, ".tif", 4) == 0 || strncmp(EXT, ".tiff", 5) == 0)

#define NULL_TIFF_HEADER ((TiffHdr){0})
#define IS_NULL_HEADER(HDR) (HDR.order == 0)

#define NULL_TIFF_IFD ((TiffIFD){0})
#define IS_NULL_IFD(IFD) (IFD.count == 0)

#define NULL_TIFF_IMAGE ((TiffImage){0})
#define IS_NULL_IMAGE(IMG) (IMG.type == TIFF_IMAGE_TYPE_INVALID)

#define INVALID_SAMPLE_COUNT 0
#define INVALID_ROWS_PER_STRIP 0
#define INVALID_STRIP_COUNT 0

#define INVALID_ALPHA_OFFSET -1

#define TEMP_ARENA_CAPACITY (20 * 1024 * 1024)

typedef struct tiff_reader TiffReader;
struct tiff_reader {
  FILE *fp;
  TiffHdr header;
  TiffIFD ifd;
  TiffImage img;
  Pixel *pixels;
};

typedef struct strip_data_field StripDataField;
struct strip_data_field {
  u32 *strip_value;
  const ShortLongValue *value_from_file;
  u32 strip_index;
  u32 length_in_bytes;
  u16 type_byte_count;
};

internal TiffHdr read_tiff_header(const TiffReader *reader);
internal TiffIFD read_ifd(const TiffReader *reader, Arena *arena);
internal TiffImage read_ifd_fields(const TiffReader *reader);
internal Pixel *load_image_pixels(TiffReader *reader, Arena *arena);
internal bool read_strip_data(TiffReader *reader, Arena *arena);
internal void read_strips(const TiffReader *reader, Pixel *buf);
internal void read_strip_data_field(const TiffReader *reader,
                                    StripDataField *field);
internal bool read_field(const TiffField *field, TiffImage *img);
internal bool validate_image_type(const TiffImage *img);
internal void read_from_file_with_offset(FILE *fp, void *dst, u64 count,
                                         u64 offset);

Image *read_baseline_tiff(const char *file, Arena *arena) {
  Image *img_out = NULL;

  if (!file || !arena) {
    goto READ_BASELINE_RETURN_IMG;
  }

  u64 name_length = strlen(file);
  if (name_length < TIFF_FILENAME_MIN_LENGTH) {
    goto READ_BASELINE_RETURN_IMG;
  }

  const char *ext = NULL;
  for (u32 i = name_length - 1; i >= name_length - TIFF_FILENAME_MIN_LENGTH;
       --i) {
    if (file[i] == '.') {
      ext = &(file[i]);
      break;
    }
  }

  if (!ext || !IS_TIFF_EXTENSION(ext)) {
    goto READ_BASELINE_RETURN_IMG;
  }

  TiffReader reader = {0};

  reader.fp = fopen(file, "rb");
  if (!reader.fp) {
    goto READ_BASELINE_RETURN_IMG;
  }

  Arena *temp_arena = NULL;
  if (!wapp_mem_arena_init(&temp_arena, TEMP_ARENA_CAPACITY)) {
    goto READ_BASELINE_FILE_CLEANUP;
  }

  reader.header = read_tiff_header(&reader);
  if (IS_NULL_HEADER(reader.header)) {
    goto READ_BASELINE_DESTROY_ARENA;
  }

  reader.ifd = read_ifd(&reader, temp_arena);
  if (IS_NULL_IFD(reader.ifd)) {
    goto READ_BASELINE_DESTROY_ARENA;
  }

  reader.img = read_ifd_fields(&reader);
  if (IS_NULL_IMAGE(reader.img)) {
    goto READ_BASELINE_DESTROY_ARENA;
  }

  assert((reader.img.type == TIFF_IMAGE_TYPE_RGB) &&
         "Currently, only RGB images are supported");

  reader.pixels = load_image_pixels(&reader, temp_arena);
  if (!reader.pixels) {
    goto READ_BASELINE_DESTROY_ARENA;
  }

  img_out = create_image(reader.img.image_width, reader.img.image_length,
                         reader.pixels, arena);

READ_BASELINE_DESTROY_ARENA:
  wapp_mem_arena_destroy(&temp_arena);

READ_BASELINE_FILE_CLEANUP:
  fclose(reader.fp);

READ_BASELINE_RETURN_IMG:
  return img_out;
}

internal TiffHdr read_tiff_header(const TiffReader *reader) {
  TiffHdr header = NULL_TIFF_HEADER;
  read_from_file_with_offset(reader->fp, &header, sizeof(TiffHdr), 0);

  switch (header.order) {
  case TIFF_ORDER_LITTLE_ENDIAN:
    if (IS_BIG_ENDIAN) {
      header.magic = htons(header.magic);
      header.first_ifd_offset = htonl(header.first_ifd_offset);
    }

    break;
  case TIFF_ORDER_BIG_ENDIAN:
    if (IS_LITTLE_ENDIAN) {
      header.magic = ntohs(header.magic);
      header.first_ifd_offset = ntohl(header.first_ifd_offset);
    }

    break;
  default:
    return NULL_TIFF_HEADER;
  }

  if (header.magic != TIFF_MAGIC) {
    return NULL_TIFF_HEADER;
  }

  return header;
}

internal TiffIFD read_ifd(const TiffReader *reader, Arena *arena) {
  TiffIFD ifd = NULL_TIFF_IFD;
  read_from_file_with_offset(reader->fp, &(ifd.count), sizeof(ifd.count),
                             reader->header.first_ifd_offset);

  switch (reader->header.order) {
  case TIFF_ORDER_LITTLE_ENDIAN:
    if (IS_BIG_ENDIAN) {
      ifd.count = htons(ifd.count);
    }

    break;
  case TIFF_ORDER_BIG_ENDIAN:
    if (IS_LITTLE_ENDIAN) {
      ifd.count = ntohs(ifd.count);
    }

    break;
  }

  u64 field_byte_count = sizeof(TiffField) * ifd.count;
  ifd.fields = (TiffField *)wapp_mem_arena_alloc(arena, field_byte_count);
  if (!(ifd.fields)) {
    return NULL_TIFF_IFD;
  }

  fread(ifd.fields, field_byte_count, 1, reader->fp);
  fread(&(ifd.next_ifd), sizeof(ifd.next_ifd), 1, reader->fp);

  return ifd;
}

internal TiffImage read_ifd_fields(const TiffReader *reader) {
  TiffImage img_out = NULL_TIFF_IMAGE;
  img_out.type = TIFF_IMAGE_TYPE_BILEVEL;

  for (u64 i = 0; i < reader->ifd.count; ++i) {
    TiffField *field = &(reader->ifd.fields[i]);

    switch (reader->header.order) {
    case TIFF_ORDER_BIG_ENDIAN:
      if (IS_LITTLE_ENDIAN) {
        field->tag = ntohs(field->tag);
        field->type = ntohs(field->type);
        field->count = ntohl(field->count);

        if (field_types[field->type].byte_count > 2 ||
            field_types[field->type].byte_count * field->count > 4) {
          field->value_offset = ntohl(field->value_offset);
        } else {
          field->value_offset = ntohs(field->value_offset);
        }
      }
      break;
    case TIFF_ORDER_LITTLE_ENDIAN:
      if (IS_BIG_ENDIAN) {
        field->tag = htons(field->tag);
        field->type = htons(field->type);
        field->count = htonl(field->count);

        if (field_types[field->type].byte_count > 2 ||
            field_types[field->type].byte_count * field->count > 4) {
          field->value_offset = htonl(field->value_offset);
        } else {
          field->value_offset = htons(field->value_offset);
        }
      }
      break;
    }

    if (!read_field(field, &img_out)) {
      img_out = NULL_TIFF_IMAGE;
      goto READ_FIELDS_RETURN_IMAGE;
    }
  }

  if (!validate_image_type(&img_out)) {
    img_out = NULL_TIFF_IMAGE;
    goto READ_FIELDS_RETURN_IMAGE;
  }

#ifdef DEBUG
  // clang-format off
  printf("SIZE (width, height)                          :   %u, %u\n", img_out.image_width, img_out.image_length);
  printf("SAMPLES (bits per sample, count, is offset)   :   %u, %u, %u\n", img_out.bits_per_sample, img_out.sample_count, img_out.bits_per_sample_offset);
  printf("EXTRA SAMPLES (samples, count, is offset)     :   %u, %u, %u\n", img_out.extra_samples, img_out.extra_samples_count, img_out.extra_samples_offset);
  printf("SAMPLES PER PIXEL                             :   %u\n", img_out.sample_count);
  printf("PHOTOMETRIC INTERPRETATION                    :   %u\n", img_out.photometric_interpretation);
  printf("ROWS PER STRIP (rows, strip count)            :   %u, %u\n", img_out.rows_per_strip, img_out.strip_count);
  printf("STRIP OFFSETS (offsets, is offset)            :   %u, %u\n", img_out.strip_offsets.long_val, img_out.strip_offsets_offset);
  printf("STRIP BYTES (byte count, is offset)           :   %u, %u\n", img_out.strip_byte_counts.long_val, img_out.strip_byte_counts_offset);
  // clang-format on
#endif

READ_FIELDS_RETURN_IMAGE:
  return img_out;
}

internal Pixel *load_image_pixels(TiffReader *reader, Arena *arena) {
  Pixel *buf = NULL;
  u64 img_byte_count =
      sizeof(Pixel) * reader->img.image_width * reader->img.image_length;

  if (!read_strip_data(reader, arena)) {
    goto LOAD_IMAGE_PIXELS_RETURN;
  }

#ifdef DEBUG
  for (u32 i = 0; i < reader->img.strip_count; ++i) {
    printf("%u, %u\n", reader->img.strips[i].offset,
           reader->img.strips[i].byte_count);
  }
#endif

  buf = wapp_mem_arena_alloc(arena, img_byte_count);
  if (!buf) {
    goto LOAD_IMAGE_PIXELS_RETURN;
  }

  read_strips(reader, buf);

LOAD_IMAGE_PIXELS_RETURN:
  return buf;
}

internal bool read_strip_data(TiffReader *reader, Arena *arena) {
  reader->img.strips =
      wapp_mem_arena_alloc(arena, sizeof(TiffStrip) * reader->img.strip_count);
  if (!(reader->img.strips)) {
    return false;
  }

  u32 offsets_total_bytes =
      reader->img.strip_count * reader->img.strip_offsets_type_byte_count;
  u32 byte_count_total_bytes =
      reader->img.strip_count * reader->img.strip_byte_count_type_byte_count;

  if ((!(reader->img.strip_offsets_offset) &&
       offsets_total_bytes > sizeof(u32)) ||
      (!(reader->img.strip_byte_counts_offset) &&
       byte_count_total_bytes > sizeof(u32))) {
    return false;
  }

  u16 offset_size = reader->img.strip_offsets_type_byte_count;
  u16 counts_size = reader->img.strip_byte_count_type_byte_count;

  for (u32 i = 0; i < reader->img.strip_count; ++i) {
    TiffStrip *strip = &(reader->img.strips[i]);
    StripDataField offsets = {
        .strip_value = &(strip->offset),
        .value_from_file = &(reader->img.strip_offsets),
        .strip_index = i,
        .length_in_bytes = offsets_total_bytes,
        .type_byte_count = offset_size,
    };
    StripDataField counts = {
        .strip_value = &(strip->byte_count),
        .value_from_file = &(reader->img.strip_byte_counts),
        .strip_index = i,
        .length_in_bytes = byte_count_total_bytes,
        .type_byte_count = counts_size,
    };

    read_strip_data_field(reader, &offsets);
    read_strip_data_field(reader, &counts);
  }

  return true;
}

internal void read_strips(const TiffReader *reader, Pixel *buf) {
  u64 position = 0;
  for (u64 i = 0; i < reader->img.strip_count; ++i) {
    const TiffStrip *strip = &(reader->img.strips[i]);

    for (u64 j = 0; j < strip->byte_count / reader->img.sample_count; ++j) {
      Pixel *p = &(buf[position]);

      read_from_file_with_offset(reader->fp, p, reader->img.sample_count,
                                 strip->offset + j * reader->img.sample_count);

      if (reader->img.sample_count == 3) {
        p->a = 255;
      }

      ++position;
    }
  }
}

internal void read_strip_data_field(const TiffReader *reader,
                                    StripDataField *field) {
  u16 tiff_long_byte_count = field_types[TIFF_FIELD_TYPE_LONG].byte_count;

  bool value_is_file_offset = field->length_in_bytes > tiff_long_byte_count;
  u32 offset = field->type_byte_count * field->strip_index;

  if (!value_is_file_offset) {
    u8 *value = (u8 *)&(field->value_from_file->long_val);
    memcpy(field->strip_value, value + offset, field->type_byte_count);

    return;
  }

  read_from_file_with_offset(reader->fp, field->strip_value,
                             field->type_byte_count,
                             field->value_from_file->long_val + offset);

  u16 tiff_short_byte_count = field_types[TIFF_FIELD_TYPE_SHORT].byte_count;
  switch (reader->header.order) {
  case TIFF_ORDER_BIG_ENDIAN:
    if (IS_LITTLE_ENDIAN) {
      *(field->strip_value) = field->type_byte_count > tiff_short_byte_count
                                  ? ntohl(*(field->strip_value))
                                  : ntohs(*(field->strip_value));
    }
    break;
  case TIFF_ORDER_LITTLE_ENDIAN:
    if (IS_BIG_ENDIAN) {
      *(field->strip_value) = field->type_byte_count > tiff_short_byte_count
                                  ? htonl(*(field->strip_value))
                                  : htons(*(field->strip_value));
    }
    break;
  }
}

internal bool read_field(const TiffField *field, TiffImage *img) {
  switch (field->tag) {
  case TIFF_PUBLIC_TAG_IMAGE_WIDTH:
    if (field->count > 1) {
      return false;
    }

    img->image_width = field->value_offset;

    break;
  case TIFF_PUBLIC_TAG_IMAGE_LENGTH:
    if (field->count > 1) {
      return false;
    }

    img->image_length = field->value_offset;

    break;
  case TIFF_PUBLIC_TAG_BITS_PER_SAMPLE:
    img->type |= field->tag;

    if (img->sample_count == INVALID_SAMPLE_COUNT) {
      img->sample_count = field->count;
    } else if (img->sample_count != field->count) {
      return false;
    }

    img->bits_per_sample = field->value_offset;

    img->bits_per_sample_offset =
        img->sample_count * field_types[field->type].byte_count > 4;

    break;
  case TIFF_PUBLIC_TAG_COMPRESSION:
    if (field->count > 1 ||
        field->value_offset != TIFF_COMPRESSION_UNCOMPRESSED) {
      return false;
    }

    img->compression = field->value_offset;

    break;
  case TIFF_PUBLIC_TAG_PHOTOMETRIC_INTERPRETATION:
    if (field->count > 1) {
      return false;
    }

    img->photometric_interpretation = field->value_offset;

    break;
  case TIFF_PUBLIC_TAG_STRIP_OFFSETS:
    img->strip_offsets_type_byte_count = field_types[field->type].byte_count;

    if (img->strip_offsets_type_byte_count == sizeof(u16)) {
      img->strip_offsets.short_val = field->value_offset;
    } else if (img->strip_offsets_type_byte_count == sizeof(u32)) {
      img->strip_offsets.long_val = field->value_offset;
    }

    if (img->strip_count == INVALID_STRIP_COUNT) {
      img->strip_count = field->count;
    } else if (img->strip_count != field->count) {
      return false;
    }

    img->strip_offsets_offset =
        img->strip_count * img->strip_offsets_type_byte_count > 4;

    break;
  case TIFF_PUBLIC_TAG_SAMPLES_PER_PIXEL:
    img->type |= field->tag;

    if (field->count > 1) {
      return false;
    }

    img->sample_count = field->value_offset;

    break;
  case TIFF_PUBLIC_TAG_ROWS_PER_STRIP:
    if (field->count > 1) {
      return false;
    }

    img->rows_per_strip = field->value_offset;

    if (img->image_length > 0) {
      img->strip_count =
          floor((f64)(img->image_length + img->rows_per_strip - 1) /
                img->rows_per_strip);
    } else {
      img->strip_count = INVALID_STRIP_COUNT;
    }

    break;
  case TIFF_PUBLIC_TAG_STRIP_BYTE_COUNTS:
    img->strip_byte_count_type_byte_count = field_types[field->type].byte_count;

    if (img->strip_byte_count_type_byte_count == sizeof(u16)) {
      img->strip_byte_counts.short_val = field->value_offset;
    } else if (img->strip_byte_count_type_byte_count == sizeof(u32)) {
      img->strip_byte_counts.long_val = field->value_offset;
    }

    if (img->strip_count == INVALID_STRIP_COUNT) {
      img->strip_count = field->count;
    } else if (img->strip_count != field->count) {
      return false;
    }

    img->strip_byte_counts_offset =
        img->strip_count * img->strip_byte_count_type_byte_count > 4;

    break;
  case TIFF_PUBLIC_TAG_COLOR_MAP:
    img->type |= field->tag;
    break;
  case TIFF_PUBLIC_TAG_EXTRA_SAMPLES:
    img->extra_samples_count = field->count;

    img->extra_samples = field->value_offset;

    img->extra_samples_offset =
        img->extra_samples_count * field_types[field->type].byte_count > 4;

    break;
  case TIFF_PUBLIC_TAG_PLANAR_CONFIGURATION:
    if (field->count > 1 || field->value_offset != TIFF_PLANAR_CONFIG_CHUNKY) {
      return false;
    }

    break;
  }

  return true;
}

internal bool validate_image_type(const TiffImage *img) {
  if (img->type < TIFF_IMAGE_TYPE_BILEVEL || img->type > TIFF_IMAGE_TYPE_RGB) {
    return false;
  }

  switch (img->type) {
  case TIFF_IMAGE_TYPE_BILEVEL:
  case TIFF_IMAGE_TYPE_GRAYSCALE:
    if (img->photometric_interpretation >
        TIFF_PHOTOMETRIC_INTERPRETATION_BLACK_IS_ZERO) {
      return false;
    }

    break;
  case TIFF_IMAGE_TYPE_PALETTE:
    if (img->photometric_interpretation !=
        TIFF_PHOTOMETRIC_INTERPRETATION_RGB_PALETTE) {
      return false;
    }

    break;
  case TIFF_IMAGE_TYPE_RGB:
    if (img->photometric_interpretation !=
        TIFF_PHOTOMETRIC_INTERPRETATION_RGB) {
      return false;
    }

    break;
  default:
    return false;
  }

  return true;
}

internal void read_from_file_with_offset(FILE *fp, void *dst, u64 count,
                                         u64 offset) {
  if (!fp || !dst) {
    return;
  }

  fseek(fp, offset, SEEK_SET);

  fread(dst, count, 1, fp);
}