#include "drive.h"
#include "panic.h"
#include "util.h"
#include "ata.h"
#include "fat.h"
#include "pmap.h"

#define MAX_FAT_DRIVES 16
#define MAX_OPEN_FILES_PER_DRIVE 32

#define PATH_SEP_CHAR '/'
#define EXT_SEP_CHAR '.'

enum {
  FA_READ_ONLY = 0x01,
  FA_HIDDEN    = 0x02,
  FA_SYSTEM    = 0x04,
  FA_LABEL     = 0x08,
  FA_DIRECTORY = 0x10,
  FA_ARCHIVE   = 0x20,

  FA_LFN       = 0x0f
};

struct directory_entry {
  uint8_t name[11];
  uint8_t attrib;
  uint8_t name_case;
  uint8_t created_decimal;
  uint16_t created_time;
  uint16_t created_date;
  uint16_t accessed_date;
  uint16_t ignore;
  uint16_t modified_time;
  uint16_t modified_date;
  uint16_t first_cluster;
  uint32_t length;
} __attribute__ ((packed));

struct fat_info {
  //3 bytes jump
  uint8_t  oem[8];
  uint16_t bytes_per_sector;//Assumed to be 512
  uint8_t  sectors_per_cluster;//Assumed to be 1
  uint16_t reserved_sectors;
  uint8_t  fats;//Only first is used
  uint16_t root_entries;
  uint16_t sectors;//Assumed not to be 0
  uint8_t  media_type;
  uint16_t sectors_per_fat;
  uint16_t sectors_per_track;
  uint16_t heads;
  uint32_t hidden_sectors;
  uint32_t sectors_long;
  uint8_t  drive_number;
  uint8_t  reserved;
  uint8_t  ext_boot_marker;
  uint32_t volume_id;
  uint8_t  label[11];
  uint8_t  fs_type[8];
} __attribute__ ((packed));

#define CTOS(c, fdi) ((fdi)->data_start + (c) - 2)

struct open_file_info {
  //directory entry is the di_number'th entry in the di_sector'th sector
  //di_sector of 0 indicates an unused handle
  uint32_t di_sector;
  uint8_t di_number;

  uint16_t start_cluster;
  uint32_t length;
};

struct fat_drive_info {
  const struct fat_info *fi;
  uint16_t *fat;
  uint16_t root_start;
  uint16_t data_start;
  struct open_file_info open_files[MAX_OPEN_FILES_PER_DRIVE];
};

static struct fat_drive_info infos[MAX_FAT_DRIVES];
static uint8_t next_id;

static uint8_t fat_driver_buffer[512];
static struct fat_info *next_fi;

static void alloc_next_fi() {
  if (!((uint32_t)(next_fi = (struct fat_info *)((uint32_t)next_fi + 64)) & 0xfff))
    if (!(next_fi = allocate_kernel_pages(1)))
      panic("Out of memory in FAT driver.");
}

static const struct drive *cur_drive;
static fs_id_t cur_id;
static const struct fat_drive_info *cur_fdi;
static struct directory_entry *cur_dir;
static uint32_t cur_sect;

//loads cluster `c`
static void load_cluster(uint16_t c, void *to) {
  if (c == 0) {
    *(uint8_t *)to = 0;
    return;
  }

  uint32_t s = CTOS(c, cur_fdi);
  cur_drive->read_sectors(cur_drive, s, 1, to);
}

__attribute__ ((pure))
static uint16_t next_cluster(uint16_t c) {
  uint16_t found = infos[cur_id].fat[c];
  return (found < 2) || (found >= 0xfff0) ? 0 : found;
}

static const uint8_t this_dir[]   = {'.', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' '};
static const uint8_t parent_dir[] = {'.', '.', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' '};

static inline bool check_fat_names(const uint8_t *a, const uint8_t *b) {
  return (((uint32_t *)a)[0] == ((uint32_t *)b)[0]) &&
         (((uint32_t *)a)[1] == ((uint32_t *)b)[1]) &&
         (((uint16_t *)a)[4] == ((uint16_t *)b)[4]) &&
         (((uint8_t *)a)[10] == ((uint8_t *)b)[10]);
}

//after: cur_dir -> specified entry in root
static bool try_locate_root_entry(const uint8_t *fat_name) {
  cur_sect = cur_fdi->root_start - 1;
  cur_dir = (struct directory_entry *)(fat_driver_buffer + 512);
  while (true) {
    if (cur_dir == (struct directory_entry *)(fat_driver_buffer + 512)) {
      cur_dir = (struct directory_entry *)fat_driver_buffer;
      ++cur_sect;
      cur_drive->read_sectors(cur_drive, cur_sect, 1, cur_dir);
    }
    if (!*(uint8_t *)cur_dir)
      return false;
    if (check_fat_names(cur_dir->name, fat_name))
      return true;
    else
      ++cur_dir;
  }
}

//before: cur_dir -> entry of dir to search
//after: cur_dir -> specified entry in dir
static bool try_locate_entry(const uint8_t *fat_name) {
  uint16_t cur_dir_cluster = cur_dir->first_cluster;
  load_cluster(cur_dir_cluster, fat_driver_buffer);
  cur_dir = (struct directory_entry *)fat_driver_buffer;
  while (true) {
    if (cur_dir == (struct directory_entry *)(fat_driver_buffer + 512)) {
      cur_dir = (struct directory_entry *)fat_driver_buffer;
      load_cluster(cur_dir_cluster = next_cluster(cur_dir_cluster), fat_driver_buffer);
    }
    if (!*(uint8_t *)cur_dir)
      return false;
    if (check_fat_names(cur_dir->name, fat_name)) {
      cur_sect = CTOS(cur_dir_cluster, cur_fdi);
      return true;
    }
    else
      ++cur_dir;
  }
}

//puts first path component's fat name into fat_name_buffer,
//returns rest of path
static const char *split_path(const char *path, uint8_t *fat_name_buffer) {
  uint8_t pi = 0, fi = 0;
  while (1) {
    if ((path[pi] == PATH_SEP_CHAR) || !path[pi]) {
      while (fi != 11)
        fat_name_buffer[fi++] = (uint8_t)' ';
      return path + (path[pi] ? pi + 1 : pi);
    }
    if (path[pi] == EXT_SEP_CHAR)
      if (fi <= 8) {
        while (fi != 8)
          fat_name_buffer[fi++] = (uint8_t)' ';
        ++pi;
      }
      else
        panic("Bad path in FAT16 driver");
    else if (((fi == 8) && (path[pi - 1] != EXT_SEP_CHAR)) || (fi == 11))
      panic("Bad path in FAT16 driver");
    else {
      fat_name_buffer[fi++] = (uint8_t)path[pi++];
    }
  }
}

//cur_dir -> specified entry
static bool try_load_from_path(const struct drive *d, const char *path) {
  cur_drive = d;
  cur_id = d->drive_id;
  cur_fdi = &infos[cur_id];

  uint8_t fat_name[11];
  path = split_path(path, fat_name);
  if (!try_locate_root_entry(fat_name))
    return false;
  while (*path) {
    path = split_path(path, fat_name);
    if (!try_locate_entry(fat_name))
      return false;
  }

  return true;
}

static file_id_t fat_get_file(const struct drive *d, const char *path) {
  d->ready(d);
  struct open_file_info *open_files = infos[d->drive_id].open_files - 1;
  for (file_id_t n = 1; n != MAX_OPEN_FILES_PER_DRIVE + 1; ++n)
    if (!open_files[n].di_sector) {
      if (!try_load_from_path(d, path)) {
        d->done(d);
        return 0;
      }

      d->done(d);
      open_files[n].di_sector = cur_sect;
      open_files[n].di_number = cur_dir - (struct directory_entry *)fat_driver_buffer;
      open_files[n].start_cluster = cur_dir->first_cluster;
      open_files[n].length = cur_dir->length;
      return n;
    }

  panic("Maximum number of files open reached for FAT drive.");
}

static void fat_free_file(const struct drive *d, file_id_t fid) {
  infos[d->drive_id].open_files[fid - 1].di_sector = 0;
}

static void fat_load_sector(const struct drive *d, file_id_t fid, uint32_t sector, void *at) {
  cur_drive = d;
  cur_id = d->drive_id;
  cur_fdi = &infos[cur_id];

  uint16_t c = cur_fdi->open_files[fid - 1].start_cluster;
  for (uint32_t i = 0; i < sector; ++i)
    c = next_cluster(c);

  d->ready(d);
  load_cluster(c, at);
  d->done(d);
}

__attribute__ ((pure))
static uint32_t fat_get_file_length(const struct drive *d, file_id_t fid) {
  return infos[d->drive_id].open_files[fid - 1].length;
}

__attribute__ ((pure))
static uint32_t fat_get_free_sectors(const struct drive *d) {
  uint16_t *start = infos[d->fs_id].fat + 2;
  uint16_t *end = start + d->n_sectors - infos[d->fs_id].data_start;
  uint32_t count = 0;
  for (uint16_t *i = start; i < end; ++i)
    if (!*i)
      ++count;
  return count;
}

static void fat_name_to_path(const uint8_t *fat_name, char *path) {
  uint8_t last_visible = -1;
  for (uint8_t i = 0; i < 8; ++i) {
    if (fat_name[i] != (uint8_t)' ')
      last_visible = i;
    path[i] = (char)fat_name[i];
  }

  if (fat_name[8] || fat_name[9] || fat_name[10]) {
    path[last_visible + 1] = EXT_SEP_CHAR;
    for (uint8_t fi = 8, ti = last_visible + 2; fi < 11; ++fi, ++ti) {
      if (fat_name[fi] != (uint8_t)' ')
        last_visible = ti;
      path[ti] = (char)fat_name[fi];
    }
  }

  path[last_visible + 1] = '\0';
}

static uint32_t enumerate_root(const struct drive *d, struct directory_content_info *info, uint32_t max) {
  uint32_t sect = infos[d->drive_id].root_start - 1;
  struct directory_entry *entry = (struct directory_entry *)(fat_driver_buffer + 512);
  struct directory_content_info *fill = info;

  while (true) {
    if (entry == (struct directory_entry *)(fat_driver_buffer + 512)) {
      entry = (struct directory_entry *)fat_driver_buffer;
      ++sect;
      d->read_sectors(d, sect, 1, entry);
    }

    if (!*(uint8_t *)entry || (info == fill + max)) {
      d->done(d);
      return fill - info;
    }

    if (entry-> attrib & FA_LABEL) {
      ++entry;
      continue;
    }

    fill->is_dir = entry->attrib & FA_DIRECTORY;
    fill->size = entry->length;
    fat_name_to_path(entry->name, fill->name);

    ++entry;
    ++fill;
  }
}

static uint32_t fat_enumerate_dir(const struct drive *d, const char *path, struct directory_content_info *info, uint32_t max) {
  d->ready(d);

  if (!*path)
    return enumerate_root(d, info, max);

  if (!try_load_from_path(d, path)) {
    d->done(d);
    return 0;
  }

  uint16_t cluster = cur_dir->first_cluster;
  load_cluster(cluster, fat_driver_buffer);
  struct directory_entry *entry = (struct directory_entry *)fat_driver_buffer;
  struct directory_content_info *fill = info;

  while (true) {
    if (entry == (struct directory_entry *)(fat_driver_buffer + 512)) {
      entry = (struct directory_entry *)fat_driver_buffer;
      load_cluster(cluster = next_cluster(cluster), fat_driver_buffer);
    }

    if (!*(uint8_t *)entry || (fill == info + max)) {
      d->done(d);
      return fill - info;
    }

    if (check_fat_names(entry->name, this_dir) || check_fat_names(entry->name, parent_dir)) {
      ++entry;
      continue;
    }

    fill->is_dir = entry->attrib & FA_DIRECTORY;
    fill->size = entry->length;
    fat_name_to_path(entry->name, fill->name);

    ++entry;
    ++fill;
  }
}

void init_fat() {
  next_fi = allocate_kernel_pages(1);
  next_id = 0;
}

bool try_fat_init_drive(struct drive *d) {
  if (next_id >= MAX_FAT_DRIVES)
    panic("Maximum number of FAT drives reached.");

  if (!d->read_sectors(d, 0, 1, fat_driver_buffer))
    return false;
  memcpy(next_fi, fat_driver_buffer + 3, sizeof(struct fat_info));
  uint32_t *fs_type_32 = (uint32_t *)next_fi->fs_type;
  if ((fs_type_32[0] != ('F' + 'A' * 256 + 'T' * 65536 + '1' * 16777216)) ||
      (fs_type_32[1] != ('6' + ' ' * 256 + ' ' * 65536 + ' ' * 16777216)))
    return false;

  d->fs_type = "FAT16";
  d->get_file = &fat_get_file;
  d->free_file = &fat_free_file;
  d->load_sector = &fat_load_sector;
  d->get_file_length = &fat_get_file_length;
  d->enumerate_dir = &fat_enumerate_dir;
  d->get_free_sectors = &fat_get_free_sectors;

  d->fs_id = next_id;
  infos[next_id].fi = next_fi;
  infos[next_id].fat = allocate_kernel_pages(((next_fi->sectors_per_fat - 1) >> 3) + 1);
  infos[next_id].root_start = next_fi->reserved_sectors + 
                              next_fi->sectors_per_fat * next_fi->fats;
  infos[next_id].data_start = infos[next_id].root_start +
                              ((next_fi->root_entries - 1) >> 4) + 1;

  d->read_sectors(d, next_fi->reserved_sectors, next_fi->sectors_per_fat, infos[next_id].fat);
  struct open_file_info *open_files = infos[next_id].open_files - 1;
  for (file_id_t i = 0; i < MAX_OPEN_FILES_PER_DRIVE; ++i)
    open_files[i].di_sector = 0;

  alloc_next_fi();
  ++next_id;
  return true;
}