#include "src/fat16.h"
#include "src/ata.h"
#include "src/fs.h"
#include "src/memory.h"
#include "src/string.h"
#include "src/util.h"
#include <stdint.h>

#define SECTOR_SIZE 512
#define FIRST_PARTITION_SECTOR 2048
#define ATTRIBUTE_SUBDIRECTORY 0x10

typedef struct __attribute__((packed)) {
  uint16_t bytes_per_sector;
  uint8_t sectors_per_cluster;
  uint16_t reserved_sectors;
  uint8_t fats_count;
  uint16_t root_entry_count;
  uint16_t total_sectors_16;
  uint8_t media_type;
  uint16_t sectors_per_fat;
} fat16_bpb_t;

typedef struct __attribute__((packed)) {
  char name[8];
  char ext[3];
  uint8_t attributes;
  uint8_t reserved[10];
  uint16_t modified_time;
  uint16_t modified_date;
  uint16_t first_cluster;
  uint32_t size;
} fat16_dir_entry_t;

typedef struct {
  fs_node_t base;
  fat16_dir_entry_t entry;
} fat16_node_t;

static fat16_bpb_t bpb; // Assuming one partition.
static fs_node_t *fs = NUL;
static uint16_t *fat = NUL;

static void to_8_3(const char *name, char *output) {
  memory_set(' ', 11, output);
  output[11] = '\0';
  const char *c = name;
  uint64_t i = 0;
  uint8_t ext = 0;
  while (*c) {
    if (*c == '.') {
      i = 8;
      ext = 1;
    } else if (i < (!ext ? 8 : 11)) {
      output[i++] = *c >= 'a' && *c <= 'z' ? *c - 32 : *c;
    }
    c++;
  }
}

static void from_8_3(const char *name, const char *extension, char *output) {
  memory_set(0, 13, output);

  uint16_t ni = 0, ei = 0, oi = 0;
  while (ni < 8 && name[ni] != ' ') {
    output[oi++] = name[ni++];
  }

  if (extension[ei] != ' ') {
    output[oi++] = '.';
    while (ei < 3 && extension[ei] != ' ') {
      output[oi++] = extension[ei++];
    }
  }
}

fs_node_t *fat16_mount() { // Assuming one partition.
  uint8_t sector[512];
  ata_read_sectors(FIRST_PARTITION_SECTOR, 1, &sector);
  memory_copy(sector + 11, sizeof(fat16_bpb_t), &bpb);
  fat16_node_t *node = memory_allocate(sizeof(fat16_node_t));
  node->base.type = FAT16;
  node->base.name[0] = node->entry.name[0] = '/';
  node->base.size = node->entry.size = sizeof(fat16_dir_entry_t) * bpb.root_entry_count;
  node->base.is_dir = 1;
  return fs = (fs_node_t *)node;
}

static void ensure_fat() {
  ASSERT(bpb.sectors_per_fat < 256, "ensure_fat: big FAT not implemented")
  if (!fat) {
    fat = memory_allocate(bpb.sectors_per_fat * SECTOR_SIZE);
    ata_read_sectors(FIRST_PARTITION_SECTOR + bpb.reserved_sectors, (uint8_t)bpb.sectors_per_fat, fat);
  }
}

static fat16_dir_entry_t *load_directory(fat16_node_t *directory) {
  fat16_dir_entry_t *entries;

  if (!directory->entry.first_cluster) {
    uint8_t sectors = (uint8_t)((directory->base.size + SECTOR_SIZE - 1) / SECTOR_SIZE);
    entries = memory_allocate(sectors * SECTOR_SIZE + sizeof(fat16_dir_entry_t)); // One extra as null terminator.
    ata_read_sectors(FIRST_PARTITION_SECTOR + bpb.reserved_sectors + bpb.fats_count * bpb.sectors_per_fat, sectors, entries);
  } else {
    ensure_fat();

    uint16_t next_cluster = directory->entry.first_cluster;
    uint32_t cluster_count = 0;
    while (next_cluster < 0xFFF8) {
      cluster_count++;
      next_cluster = fat[next_cluster];
    }
    entries =
        memory_allocate(cluster_count * bpb.sectors_per_cluster * SECTOR_SIZE + sizeof(fat16_dir_entry_t)); // One extra as null terminator.

    uint8_t *chunk = (uint8_t *)entries;
    next_cluster = directory->entry.first_cluster;
    while (next_cluster < 0xFFF8) {
      ata_read_sectors(FIRST_PARTITION_SECTOR + bpb.reserved_sectors + bpb.sectors_per_fat * bpb.fats_count +
                           (bpb.root_entry_count * sizeof(fat16_dir_entry_t) + SECTOR_SIZE - 1) / SECTOR_SIZE +
                           bpb.sectors_per_cluster * (next_cluster - 2),
                       bpb.sectors_per_cluster, chunk);
      chunk += bpb.sectors_per_cluster * SECTOR_SIZE;
      next_cluster = fat[next_cluster];
    }
  }

  return entries;
}

static fs_node_t *open_entry(const char *name, const fat16_dir_entry_t *entry) {
  if (!entry->name[0]) {
    return NUL;
  }

  fat16_node_t *result = memory_allocate(sizeof(fat16_node_t));
  memory_copy((char *)name, string_length(name) + 1, &(result->base.name));
  result->base.type = FAT16;
  result->base.size = entry->size;
  result->base.is_dir = entry->attributes & ATTRIBUTE_SUBDIRECTORY;
  memory_copy((fat16_dir_entry_t *)entry, sizeof(fat16_dir_entry_t), (uint8_t *)result + sizeof(fs_node_t));

  return (fs_node_t *)result;
}

fs_node_t *fat16_open_by(const fs_node_t *directory, const char *name) {
  ASSERT(directory->type == FAT16, "fat16_open_by: directory is not FAT16");
  ASSERT(directory->is_dir, "fat16_open_by: directory is not a directory");

  char name_8_3[12];
  to_8_3(name, name_8_3);

  fat16_dir_entry_t *entries = load_directory((fat16_node_t *)directory);

  fat16_dir_entry_t *entry = entries;
  while (entry->name[0]) {
    if ((uint8_t)entry->name[0] != 0xE5 && (uint8_t)entry->attributes != 0x0F && bytes_equal(name_8_3, (char *)entry, 11)) {
      break;
    }
    entry++;
  }

  fs_node_t *result = open_entry(name, entry);
  memory_free(entries);
  return result;
}

fs_node_t *fat16_open_at(const fs_node_t *directory, uint64_t index) {
  ASSERT(directory->type == FAT16, "fat16_open_at: directory is not FAT16");
  ASSERT(directory->is_dir, "fat16_open_at: directory is not a directory");

  fat16_dir_entry_t *entries = load_directory((fat16_node_t *)directory);

  uint64_t ei = 0, vi = 0;
  while (entries[ei].name[0]) {
    if ((uint8_t)entries[ei].name[0] != 0xE5 && (uint8_t)entries[ei].attributes != 0x0F) {
      if (vi == index) {
        break;
      }
      vi++;
    }
    ei++;
  }

  if (!entries[ei].name[0]) {
    return NUL;
  }

  char name[13];
  from_8_3(entries[ei].name, entries[ei].ext, name);

  fs_node_t *result = open_entry(name, entries + ei);
  memory_free(entries);
  return result;
}

fs_node_t *fat16_open_again(const fs_node_t *source) {
  fs_node_t *result = memory_allocate(sizeof(fat16_node_t));
  memory_copy(source, sizeof(fat16_node_t), result);
  return result;
}

void fat16_read(const fs_node_t *file, uint64_t offset, uint64_t size, void *to) {
  ASSERT(file->type == FAT16, "fat16_read: file is not FAT16");
  ASSERT(!file->is_dir, "fat16_read: can not read directory");
  ASSERT(file->size >= offset + size, "fat16_read: offset/size are out of bounds");

  ensure_fat();

  fat16_node_t *fat_file = (fat16_node_t *)file;

  uint32_t cluster_size = bpb.sectors_per_cluster * SECTOR_SIZE;
  uint32_t next_cluster = fat_file->entry.first_cluster;

  // Assuming filesystem is correct. TODO Check for real.

  while (offset >= cluster_size) {
    next_cluster = fat[next_cluster];
    offset -= cluster_size;
  }

  uint8_t *cursor = to;

  uint8_t *tmp = memory_allocate(cluster_size);
  ata_read_sectors(FIRST_PARTITION_SECTOR + bpb.reserved_sectors + bpb.sectors_per_fat * bpb.fats_count +
                       (bpb.root_entry_count * sizeof(fat16_dir_entry_t) + SECTOR_SIZE - 1) / SECTOR_SIZE +
                       bpb.sectors_per_cluster * (next_cluster - 2),
                   bpb.sectors_per_cluster, tmp);
  uint64_t prefix_size = size <= cluster_size - offset ? size : cluster_size - offset;
  memory_copy(tmp + offset, prefix_size, cursor);
  size -= prefix_size;
  cursor += prefix_size;
  next_cluster = fat[next_cluster];

  while (size >= cluster_size) {
    ata_read_sectors(FIRST_PARTITION_SECTOR + bpb.reserved_sectors + bpb.sectors_per_fat * bpb.fats_count +
                         (bpb.root_entry_count * sizeof(fat16_dir_entry_t) + SECTOR_SIZE - 1) / SECTOR_SIZE +
                         bpb.sectors_per_cluster * (next_cluster - 2),
                     bpb.sectors_per_cluster, cursor);
    size -= cluster_size;
    cursor += cluster_size;
    next_cluster = fat[next_cluster];
  }

  if (size) {
    ata_read_sectors(FIRST_PARTITION_SECTOR + bpb.reserved_sectors + bpb.sectors_per_fat * bpb.fats_count +
                         (bpb.root_entry_count * sizeof(fat16_dir_entry_t) + SECTOR_SIZE - 1) / SECTOR_SIZE +
                         bpb.sectors_per_cluster * (next_cluster - 2),
                     bpb.sectors_per_cluster, tmp);
    memory_copy(tmp, size, cursor);
  }

  memory_free(tmp);
}

void fat16_close(fs_node_t *node) {
  ASSERT(node->type == FAT16, "fat16_close: node is not FAT16");
  ASSERT(node != fs, "fat16_close: can not unmount FS");
  memory_free(node);
}

void fat16_unmount(fs_node_t *node) {
  ASSERT(node->type == FAT16, "fat16_unmount: node is not FAT16");
  ASSERT(node == fs, "fat16_unmount: node is not filesystem");
  memory_free(node);
  fs = NUL;
}