#include "tiffread.h"
#include "aliases.h"
#include "endianness.h"
#include "image.h"
#include "mem_arena.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

// clang-format off
// clang-format on

TiffHdr read_tiff_header(FILE *fp);
TiffIFD read_ifd(FILE *fp, const TiffHdr *header, u32 offset, Arena *arena);
TiffImage read_fields(FILE *fp, const TiffHdr *header, const TiffIFD *ifd);
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) {
    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;
  }

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

  TiffHdr header = read_tiff_header(fp);
  if (IS_NULL_HEADER(header)) {
    goto READ_BASELINE_FILE_CLEANUP;
  }

  TiffIFD ifd = read_ifd(fp, &header, header.first_ifd_offset, arena);

  read_fields(fp, &header, &ifd);

READ_BASELINE_FILE_CLEANUP:
  fclose(fp);

READ_BASELINE_RETURN_IMG:
  return img_out;
}

TiffHdr read_tiff_header(FILE *fp) {
  if (!fp) {
    return NULL_TIFF_HEADER;
  }

  fseek(fp, 0, SEEK_SET);

  TiffHdr header = NULL_TIFF_HEADER;
  fread((void *)&header, sizeof(TiffHdr), 1, fp);

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

TiffIFD read_ifd(FILE *fp, const TiffHdr *header, u32 offset, Arena *arena) {
  if (!fp || !header) {
    return NULL_TIFF_IFD;
  }

  fseek(fp, offset, SEEK_SET);

  TiffIFD ifd = NULL_TIFF_IFD;

  fread(&(ifd.count), sizeof(ifd.count), 1, fp);

  switch (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, fp);
  fread(&(ifd.next_ifd), sizeof(ifd.next_ifd), 1, fp);

  return ifd;
}

TiffImage read_fields(FILE *fp, const TiffHdr *header, const TiffIFD *ifd) {
  TiffImage img_out = NULL_TIFF_IMAGE;
  u32 image_type_identifier = TIFF_IMAGE_TYPE_BILEVEL;

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

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

    switch (field->tag) {
    case TIFF_PUBLIC_TAG_IMAGE_WIDTH:
      if (field->count > 1) {
        img_out = NULL_TIFF_IMAGE;
        goto READ_FIELDS_RETURN_IMAGE;
      }

      img_out.image_width = field->value_offset;

      break;
    case TIFF_PUBLIC_TAG_IMAGE_LENGTH:
      if (field->count > 1) {
        img_out = NULL_TIFF_IMAGE;
        goto READ_FIELDS_RETURN_IMAGE;
      }

      img_out.image_length = field->value_offset;

      break;
    case TIFF_PUBLIC_TAG_BITS_PER_SAMPLE:
      image_type_identifier |= field->tag;

      if (img_out.sample_count == INVALID_SAMPLE_COUNT) {
        img_out.sample_count = field->count;
      } else if (img_out.sample_count != field->count) {
        img_out = NULL_TIFF_IMAGE;
        goto READ_FIELDS_RETURN_IMAGE;
      }

      img_out.bits_per_sample = field->value_offset;

      img_out.bits_per_sample_offset =
          img_out.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) {
        img_out = NULL_TIFF_IMAGE;
        goto READ_FIELDS_RETURN_IMAGE;
      }

      img_out.compression = field->value_offset;

      break;
    case TIFF_PUBLIC_TAG_PHOTOMETRIC_INTERPRETATION:
      if (field->count > 1) {
        img_out = NULL_TIFF_IMAGE;
        goto READ_FIELDS_RETURN_IMAGE;
      }

      img_out.photometric_interpretation = field->value_offset;

      break;
    case TIFF_PUBLIC_TAG_STRIP_OFFSETS:
      img_out.strip_offsets = field->value_offset;

      if (img_out.strip_count == INVALID_STRIP_COUNT) {
        img_out.strip_count = field->count;
      } else if (img_out.strip_count != field->count) {
        img_out = NULL_TIFF_IMAGE;
        goto READ_FIELDS_RETURN_IMAGE;
      }

      img_out.strip_offsets_offset =
          img_out.strip_count * field_types[field->type].byte_count > 4;

      break;
    case TIFF_PUBLIC_TAG_SAMPLES_PER_PIXEL:
      image_type_identifier |= field->tag;

      if (field->count > 1) {
        img_out = NULL_TIFF_IMAGE;
        goto READ_FIELDS_RETURN_IMAGE;
      }

      img_out.sample_count = field->value_offset;

      break;
    case TIFF_PUBLIC_TAG_ROWS_PER_STRIP:
      if (field->count > 1) {
        img_out = NULL_TIFF_IMAGE;
        goto READ_FIELDS_RETURN_IMAGE;
      }

      img_out.rows_per_strip = field->value_offset;

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

      break;
    case TIFF_PUBLIC_TAG_STRIP_BYTE_COUNTS:
      img_out.strip_byte_counts = field->value_offset;

      if (img_out.strip_count == INVALID_STRIP_COUNT) {
        img_out.strip_count = field->count;
      } else if (img_out.strip_count != field->count) {
        img_out = NULL_TIFF_IMAGE;
        goto READ_FIELDS_RETURN_IMAGE;
      }

      img_out.strip_byte_counts_offset =
          img_out.strip_count * field_types[field->type].byte_count > 4;

      break;
    case TIFF_PUBLIC_TAG_COLOR_MAP:
      image_type_identifier |= field->tag;
      break;
    case TIFF_PUBLIC_TAG_EXTRA_SAMPLES:
      img_out.extra_samples_count = field->count;

      img_out.extra_samples = field->value_offset;

      img_out.extra_samples_offset =
          img_out.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) {
        img_out = NULL_TIFF_IMAGE;
        goto READ_FIELDS_RETURN_IMAGE;
      }

      break;
    }
  }

  if (image_type_identifier < TIFF_IMAGE_TYPE_BILEVEL ||
      image_type_identifier > TIFF_IMAGE_TYPE_RGB) {
    img_out = NULL_TIFF_IMAGE;
    goto READ_FIELDS_RETURN_IMAGE;
  }

  switch (image_type_identifier) {
  case TIFF_IMAGE_TYPE_BILEVEL:
  case TIFF_IMAGE_TYPE_GRAYSCALE:
    if (img_out.photometric_interpretation >
        TIFF_PHOTOMETRIC_INTERPRETATION_BLACK_IS_ZERO) {
      img_out = NULL_TIFF_IMAGE;
      goto READ_FIELDS_RETURN_IMAGE;
    }

    break;
  case TIFF_IMAGE_TYPE_PALETTE:
    if (img_out.photometric_interpretation !=
        TIFF_PHOTOMETRIC_INTERPRETATION_RGB_PALETTE) {
      img_out = NULL_TIFF_IMAGE;
      goto READ_FIELDS_RETURN_IMAGE;
    }

    break;
  case TIFF_IMAGE_TYPE_RGB:
    if (img_out.photometric_interpretation !=
        TIFF_PHOTOMETRIC_INTERPRETATION_RGB) {
      img_out = NULL_TIFF_IMAGE;
      goto READ_FIELDS_RETURN_IMAGE;
    }

    break;
  default:
    img_out = NULL_TIFF_IMAGE;
    goto READ_FIELDS_RETURN_IMAGE;
  }

  img_out.type = image_type_identifier;

#if 1
  // 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, img_out.strip_offsets_offset);
  printf("STRIP BYTES (byte count, is offset)           :   %u, %u\n", img_out.strip_byte_counts, img_out.strip_byte_counts_offset);
  // clang-format on
#endif

READ_FIELDS_RETURN_IMAGE:
  return img_out;
}

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