#include <string.h>
#include <drivers/ide.h>
#include <stdio.h>

#include <fs/fat16.h>

/*
  READ BELOW!!!
  Before any of you (or me,) use this code, please note that the functions below REQUIRE 8.3 formatted names.
  in fat16.h there is a helper function for convertion between normal filenames and 8.3 ones. ('filename_to_83', see static implementation in fat16.h)
*/

static fat16_fs_t fs;
static uint8_t boot_sector[SECTOR_SIZE];

int fat16_mount(uint8_t drive)
{
  if (read_sector(0, boot_sector) != 0)
  {
    return -1;
  }

  fs.bytes_per_sector    = boot_sector[11] | (boot_sector[12] << 8);
  fs.sectors_per_cluster = boot_sector[13];
  fs.reserved_sector_count = boot_sector[14] | (boot_sector[15] << 8);
  fs.num_fats            = boot_sector[16];
  fs.root_entry_count    = boot_sector[17] | (boot_sector[18] << 8);
  fs.total_sectors_16    = boot_sector[19] | (boot_sector[20] << 8);
  fs.sectors_per_fat     = boot_sector[22] | (boot_sector[23] << 8);
  fs.total_sectors = fs.total_sectors_16 != 0 ? fs.total_sectors_16 : boot_sector[32] | (boot_sector[33]<<8) | (boot_sector[34]<<16) | (boot_sector[35]<<24);
  fs.fat_start_lba = fs.reserved_sector_count;
  fs.root_dir_lba  = fs.fat_start_lba + (fs.num_fats * fs.sectors_per_fat);
  fs.data_start_lba = fs.root_dir_lba + ((fs.root_entry_count * 32 + (fs.bytes_per_sector - 1)) / fs.bytes_per_sector);

  return 0;
}

static uint32_t cluster_to_lba(uint16_t cluster)
{
  return fs.data_start_lba + (cluster - 2) * fs.sectors_per_cluster;
}

/* ----------------------- FAT helpers ------------------------- */
static int read_fat_entry(uint16_t cluster, uint16_t* out_value)
{
  uint32_t offset = (uint32_t)cluster * 2;
  uint32_t sector_index = offset / fs.bytes_per_sector;
  uint32_t idx = offset % fs.bytes_per_sector;
  uint8_t sector[SECTOR_SIZE];

  if (read_sector(fs.fat_start_lba + sector_index, sector) != 0)
  {
    return -1;
  }

  if (idx == fs.bytes_per_sector - 1)
  {
    /* entry spans boundary -> read next sector */
    uint8_t next_sector[SECTOR_SIZE];
    
    if (read_sector(fs.fat_start_lba + sector_index + 1, next_sector) != 0)
    {
      return -1;
    }
    
    *out_value = sector[idx] | (next_sector[0] << 8);
  }
  else
  {
    *out_value = sector[idx] | (sector[idx + 1] << 8);
  }
  return 0;
}


static int write_fat_entry(uint16_t cluster, uint16_t value)
{
  uint32_t offset = cluster * 2;
  uint32_t sector_lba = fs.fat_start_lba + (offset / fs.bytes_per_sector);
  uint8_t sector[SECTOR_SIZE];
  
  if (read_sector(sector_lba, sector) != 0)
  {
    return -1;
  }
  
  sector[offset % fs.bytes_per_sector] = value & 0xFF;
  sector[(offset % fs.bytes_per_sector) + 1] = (value >> 8) & 0xFF;

  /* write to all FAT copies */
  for (uint8_t f = 0; f < fs.num_fats; f++)
  {
    if (write_sector(fs.fat_start_lba + f * fs.sectors_per_fat + (offset / fs.bytes_per_sector), sector) != 0)
    {
      return -1;
    }
  }
  return 0;
}

static int find_free_cluster(uint16_t* out_cluster)
{
  uint16_t total_clusters = (fs.total_sectors - fs.data_start_lba) / fs.sectors_per_cluster;
  uint16_t value;
  for (uint16_t c = 2; c < total_clusters; c++)
  {
    if (read_fat_entry(c, &value) != 0)
    {
      return -1;
    }
    
    if (value == 0x0000)
    {
      *out_cluster = c;
      return 0;
    }
  }
  return -1; /* no free cluster */
}


/* -------------------- Root directory ------------------------- */
static int find_free_root_entry(uint32_t* out_sector, uint32_t* out_offset)
{
  uint32_t root_dir_sectors = ((fs.root_entry_count * 32) + (fs.bytes_per_sector - 1)) / fs.bytes_per_sector;
  uint8_t sector[SECTOR_SIZE];

  for (uint32_t i = 0; i < root_dir_sectors; i++)
  {
    if (read_sector(fs.root_dir_lba + i, sector) != 0)
    {
      return -1;
    }

    for (uint32_t offset = 0; offset < SECTOR_SIZE; offset += 32)
    {
      if (sector[offset] == 0x00 || sector[offset] == 0xE5)
      {
        *out_sector = fs.root_dir_lba + i;
        *out_offset = offset;
        return 0;
      }
    }
  }
  return -1;
}

/* -------------------- File operations ----------------------- */
int fat16_find_file(const char* name, fat16_file_t* out_file)
{
  uint32_t root_dir_sectors = ((fs.root_entry_count * 32) + (fs.bytes_per_sector - 1)) / fs.bytes_per_sector;
  uint8_t sector[SECTOR_SIZE];

  for (uint32_t i = 0; i < root_dir_sectors; i++)
  {
    if (read_sector(fs.root_dir_lba + i, sector) != 0)
    {
      return -1;
    }

    for (uint32_t offset = 0; offset < SECTOR_SIZE; offset += 32)
    {
      if (sector[offset] == 0x00)
      {
        return -1; /* no more entries */
      }
      else if (sector[offset] == 0xE5)
      {
        continue; /* deleted entry */
      }
      
      if (memcmp(sector + offset, name, 11) == 0)
      {
        memcpy(out_file->name, sector + offset, 12);
        out_file->attr = sector[offset + 11];
        out_file->first_cluster = sector[offset + 26] | (sector[offset + 27] << 8);
        out_file->size = sector[offset + 28] | (sector[offset + 29] << 8) | (sector[offset + 30] << 16) | (sector[offset + 31] << 24);
        
        return 0;
      }
    }
  }
  return -1;
}

int fat16_create_file(const char* name, fat16_file_t* out_file)
{
  uint16_t free_cluster;
  
  if (find_free_cluster(&free_cluster) != 0)
  {
    return -1;
  }

  uint32_t sector_lba, offset;
  
  if (find_free_root_entry(&sector_lba, &offset) != 0)
  {
    return -1;
  }

  uint8_t sector[SECTOR_SIZE];
  
  if (read_sector(sector_lba, sector) != 0)
  {
    return -1;
  }

  /* Fill root entry */
  memset(sector + offset, ' ', 11);
  memcpy(sector + offset, name, 11);
  sector[offset + 11] = 0x20; /* normal file */
  sector[offset + 26] = free_cluster & 0xFF;
  sector[offset + 27] = (free_cluster >> 8) & 0xFF;
  sector[offset + 28] = 0; sector[offset + 29] = 0;
  sector[offset + 30] = 0; sector[offset + 31] = 0;

  if (write_sector(sector_lba, sector) != 0)
  {
    return -1;
  }

  /* mark cluster as end-of-chain */
  if (write_fat_entry(free_cluster, 0xFFFF) != 0)
  {
    return -1;
  }

  if (out_file)
  {
    memset(out_file, 0, sizeof(fat16_file_t));
    strncpy(out_file->name, name, 11);
    out_file->first_cluster = free_cluster;
    out_file->size = 0;
    out_file->attr = 0x20;
  }

  return 0;
}

int fat16_write_file(fat16_file_t* file, const void* buffer, uint32_t size) {
    const uint8_t* buf = (const uint8_t*)buffer;
    uint32_t bytes_remaining = size;
    uint16_t cluster = file->first_cluster;

    if (cluster == 0) return -1; /* safety: file must have a cluster */

    while (bytes_remaining > 0) {
        uint32_t lba = cluster_to_lba(cluster);

        for (uint8_t i = 0; i < fs.sectors_per_cluster && bytes_remaining > 0; i++) {
            uint8_t sector[SECTOR_SIZE];
            memset(sector, 0, SECTOR_SIZE);
            uint32_t copy_bytes = bytes_remaining < SECTOR_SIZE ? bytes_remaining : SECTOR_SIZE;
            memcpy(sector, buf, copy_bytes);
            if (write_sector(lba + i, sector) != 0) return -1;
            buf += copy_bytes;
            bytes_remaining -= copy_bytes;
        }

        uint16_t next_cluster;
        if (bytes_remaining > 0) {
            if (find_free_cluster(&next_cluster) != 0) return -1;
            if (write_fat_entry(cluster, next_cluster) != 0) return -1;
            if (write_fat_entry(next_cluster, 0xFFFF) != 0) return -1;
            cluster = next_cluster;
        }
    }

    /* Update file size in root entry (your code already does this on disk) */
    /* Now update the in-memory struct so subsequent reads use the new size */
    file->size = size;

    /* optionally: refresh metadata from disk:
       fat16_find_file(file->name, file); */

    return 0;
}


int fat16_read_file(fat16_file_t* file, void* buffer)
{
  uint8_t sector[SECTOR_SIZE];
  uint16_t cluster = file->first_cluster;
  uint32_t bytes_remaining = file->size;
  uint8_t* buf = (uint8_t*)buffer;

  while (cluster >= 0x0002 && cluster < 0xFFF8)
  {
    uint32_t lba = cluster_to_lba(cluster);
    for (uint8_t i = 0; i < fs.sectors_per_cluster && bytes_remaining > 0; i++)
    {
      if (read_sector(lba + i, sector) != 0)
      {
        return -1;
      }
      uint32_t copy_bytes = bytes_remaining < SECTOR_SIZE ? bytes_remaining : SECTOR_SIZE;
      memcpy(buf, sector, copy_bytes);
      buf += copy_bytes;
      bytes_remaining -= copy_bytes;
    }

    /* read next cluster from FAT */
    uint32_t fat_offset = cluster * 2; /* FAT16 */
    uint32_t fat_sector = fs.fat_start_lba + (fat_offset / fs.bytes_per_sector);
    uint16_t fat_entry;
    if (read_sector(fat_sector, sector) != 0)
    {
      return -1;
    }
    fat_entry = sector[fat_offset % fs.bytes_per_sector] | (sector[(fat_offset % fs.bytes_per_sector) + 1] << 8);
    cluster = fat_entry;
  }

  return 0;
}

int fat16_delete_file(const char* filename)
{
    uint8_t sector[512];
    uint32_t root_dir_sectors = ((fs.root_entry_count * 32) + (fs.bytes_per_sector - 1)) / fs.bytes_per_sector;
    uint32_t root_dir_start = fs.root_dir_lba;

    for (uint32_t s = 0; s < root_dir_sectors; s++)
    {
        if (read_sector(root_dir_start + s, sector) != 0)
            return -1;

        for (int i = 0; i < fs.bytes_per_sector; i += 32)
        {
            uint8_t* entry = &sector[i];

            if (entry[0] == 0x00)
                return -1; // end of directory
            if (entry[0] == 0xE5)
                continue; // already deleted

            if (memcmp(entry, filename, 11) == 0)
            {
                // Found the file — mark as deleted
                entry[0] = 0xE5;

                if (write_sector(root_dir_start + s, sector) != 0)
                    return -1;

                // Free FAT chain
                uint16_t cluster = entry[26] | (entry[27] << 8);
                while (cluster >= 0x0002 && cluster < 0xFFF8)
                {
                    uint32_t fat_offset = cluster * 2;
                    uint32_t fat_sector = fs.fat_start_lba + (fat_offset / fs.bytes_per_sector);
                    uint32_t ent_offset = fat_offset % fs.bytes_per_sector;

                    if (read_sector(fat_sector, sector) != 0)
                        break;

                    uint16_t next_cluster = sector[ent_offset] | (sector[ent_offset + 1] << 8);
                    sector[ent_offset] = 0x00;
                    sector[ent_offset + 1] = 0x00;

                    if (write_sector(fat_sector, sector) != 0)
                        break;

                    cluster = next_cluster;
                }

                return 0; // success
            }
        }
    }

    return -1; // file not found
}