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

#include <fs/fat32.h>

struct __attribute__((packed)) fat32_bpb {
    uint8_t jmp[3];
    uint8_t oem[8];
    uint16_t bytes_per_sector;
    uint8_t sectors_per_cluster;
    uint16_t reserved_sector_count;
    uint8_t num_fats;
    uint16_t root_entry_count;
    uint16_t total_sectors_16;
    uint8_t media;
    uint16_t fat_size_16;
    uint16_t sectors_per_track;
    uint16_t num_heads;
    uint32_t hidden_sectors;
    uint32_t total_sectors_32;
    uint32_t fat_size_32;
    uint16_t ext_flags;
    uint16_t fs_version;
    uint32_t root_cluster;
};

static int32_t fat32_drive = 0;
static struct fat32_bpb bpb;
static uint32_t fat_start_lba;
static uint32_t cluster_heap_lba;
static uint32_t current_directory_cluster;

static uint8_t sector[512];

int32_t fat32_init(int32_t drive)
{
  fat32_drive = drive;
  int err = ide_read48(drive, 0, 1, sector);
  if (err != 0) {
      printf("ide_read48 failed on sector 0 (err = %d)\n", err);
      return -1;
  }

  memcpy(&bpb, sector, sizeof(bpb));

  fat_start_lba = bpb.reserved_sector_count;
  cluster_heap_lba = fat_start_lba + bpb.num_fats * bpb.fat_size_32;
  current_directory_cluster = bpb.root_cluster;
  printf("Bytes per sector: %u\n", bpb.bytes_per_sector);
  printf("Sectors per cluster: %u\n", bpb.sectors_per_cluster);
  printf("Reserved sectors: %u\n", bpb.reserved_sector_count);
  printf("Number of FATs: %u\n", bpb.num_fats);
  printf("FAT size (32): %u\n", bpb.fat_size_32);
  printf("Root cluster: %u\n", bpb.root_cluster);

  return 0;
}

static uint32_t cluster_to_lba(uint32_t cluster)
{
  return cluster_heap_lba + (cluster - 2) * bpb.sectors_per_cluster;
}

static void format_83_name(const char *input, char out[11]) {
    memset(out, ' ', 11);
    int i = 0, j = 0;

    // Copy name part (up to 8 chars)
    while (input[i] && input[i] != '.' && j < 8) {
        out[j++] = toupper((unsigned char)input[i++]);
    }

    // Skip dot
    if (input[i] == '.') i++;

    // Copy extension (up to 3 chars)
    int k = 8;
    while (input[i] && j < 11) {
        out[k++] = toupper((unsigned char)input[i++]);
        j++;
    }
}

static uint32_t find_free_cluster() {
    // Read FAT sectors one by one
    for (uint32_t cluster = 2; cluster < bpb.total_sectors_32; cluster++) {
        uint32_t fat_sector = fat_start_lba + (cluster * 4) / 512;
        uint32_t fat_offset = (cluster * 4) % 512;

        if (ide_read48(fat32_drive, fat_sector, 1, sector) != 0)
            return 0;

        uint32_t entry = *(uint32_t *)(sector + fat_offset) & 0x0FFFFFFF;
        if (entry == 0x00000000) {
            return cluster;
        }
    }
    return 0; // No free cluster found
}



int32_t fat32_list_root(void)
{
  uint32_t lba = cluster_to_lba(current_directory_cluster);
  if (ide_read48(fat32_drive, lba, 1, sector) != 0)
  {
    return -1;
  }

  for (int i = 0; i < 512; i += 32)
  {
    char *entry = (char *)&sector[i];
    if ((uint8_t)entry[0] == 0x00)
    {
      break;
    }
      
    if ((uint8_t)entry[0] == 0xE5)
    {
      continue;
    }
  
      char name[12];
      memcpy(name, entry, 11);
      name[11] = 0;
      for (int j = 0; j < 11; j++)
      {
        if (name[j] == ' ')
        {
          name[j] = 0;
        }
      }
  
      uint8_t attr = entry[11];
      if (attr & FAT32_ATTR_DIRECTORY)
      {
        printf("<DIR> %s\n", name);
      }
      else
      {
        printf("      %s\n", name);
      }
    }
    return 0;
}

int32_t fat32_read_file(const char *name, uint8_t *buffer, uint32_t max_len)
{
  uint32_t lba = cluster_to_lba(current_directory_cluster);
  if (ide_read48(fat32_drive, lba, 1, sector) != 0)
  {
    return -1;
  }

  for (int32_t i = 0; i < 512; i += 32)
  {
    char *entry = (char *)&sector[i];
    if ((uint8_t)entry[0] == 0x00)
    {
      break;
    }
    if ((uint8_t)entry[0] == 0xE5)
    {
      continue;
    }

    char entry_name[12];
    memcpy(entry_name, entry, 11);
    entry_name[11] = 0;
    for (int32_t j = 0; j < 11; j++)
    {
      if (entry_name[j] == ' ')
      {
        entry_name[j] = 0;
      }
    }

    if (strncmp(entry_name, name, 11) != 0)
    {
      continue;
    }

    uint32_t cluster = (*(uint16_t *)(entry + 26)) | ((*(uint16_t *)(entry + 20)) << 16);
    uint32_t size = *(uint32_t *)(entry + 28);
    uint32_t bytes_read = 0;

    while (cluster < 0x0FFFFFF8 && bytes_read < max_len && bytes_read < size)
    {
      uint32_t lba = cluster_to_lba(cluster);
      if (ide_read48(fat32_drive, lba, bpb.sectors_per_cluster, buffer + bytes_read) != 0)
      {
        return -1;
      }

      bytes_read += bpb.sectors_per_cluster * 512;

      uint32_t fat_sector = cluster * 4 / 512;
      uint32_t fat_offset = cluster * 4 % 512;
      if (ide_read48(fat32_drive, fat_start_lba + fat_sector, 1, sector) != 0)
      {
        return -1;
      }

      cluster = *(uint32_t *)(sector + fat_offset) & 0x0FFFFFFF;
    }
    
    return bytes_read;
  }
  return -2;
}

int32_t fat32_write_file(const char *name, const uint8_t *buffer, uint32_t len)
{
  uint32_t lba = cluster_to_lba(current_directory_cluster);
  if (ide_read48(fat32_drive, lba, 1, sector) != 0)
  {
    return -1;
  }

  for (int32_t i = 0; i < 512; i += 32)
  {
    char *entry = (char *)&sector[i];
    if ((uint8_t)entry[0] == 0x00) break;
    if ((uint8_t)entry[0] == 0xE5) continue;

    char entry_name[12];
    memcpy(entry_name, entry, 11);
    entry_name[11] = 0;
    for (int j = 0; j < 11; j++)
    {
      if (entry_name[j] == ' ')
      {
        entry_name[j] = 0;
      }
    }
    
    if (strncmp(entry_name, name, 11) != 0)
    {
      continue;
    }

    uint32_t cluster = (*(uint16_t *)(entry + 26)) | ((*(uint16_t *)(entry + 20)) << 16);
    uint32_t bytes_written = 0;

    while (cluster < 0x0FFFFFF8 && bytes_written < len)
    {
      uint32_t lba = cluster_to_lba(cluster);
      if (ide_write48(fat32_drive, lba, bpb.sectors_per_cluster, buffer + bytes_written) != 0)
      {
        return -1;
      }

      bytes_written += bpb.sectors_per_cluster * 512;

      if (bytes_written < len)
      {
        uint32_t fat_sector = cluster * 4 / 512;
        uint32_t fat_offset = cluster * 4 % 512;

        if (ide_read48(fat32_drive, fat_start_lba + fat_sector, 1, sector) != 0)
        {
          return -1;
        }

        uint32_t next_cluster = find_free_cluster();
        if (next_cluster == 0)
        {
          return -3;
        }

        *(uint32_t *)(sector + fat_offset) = next_cluster & 0x0FFFFFFF;

        if (ide_write48(fat32_drive, fat_start_lba + fat_sector, 1, sector) != 0)
        {
          return -1;
        }

        cluster = next_cluster;
      }
      else
      {
        uint32_t fat_sector = cluster * 4 / 512;
        uint32_t fat_offset = cluster * 4 % 512;

        if (ide_read48(fat32_drive, fat_start_lba + fat_sector, 1, sector) != 0)
        {
          return -1;
        }

        *(uint32_t *)(sector + fat_offset) = 0x0FFFFFFF;

        if (ide_write48(fat32_drive, fat_start_lba + fat_sector, 1, sector) != 0)
        {
          return -1;
        }

        break;
      }
    }

    *(uint32_t *)(entry + 28) = len;
    if (ide_write48(fat32_drive, lba, 1, sector) != 0)
    {
      return -1;
    }

    return bytes_written;
  }
  return -2;
}

int32_t fat32_create_file(const char *name) {
    uint32_t root_lba = cluster_to_lba(bpb.root_cluster);

    // Read root directory cluster sector (assuming 1 sector here for simplicity)
    if (ide_read48(fat32_drive, root_lba, 1, sector) != 0)
        return -1;

    char formatted_name[11];
    format_83_name(name, formatted_name);

    for (int32_t i = 0; i < 512; i += 32) {
        uint8_t *entry = &sector[i];
        if (entry[0] == 0x00 || entry[0] == 0xE5) {
            // Found free directory entry, clear it
            memset(entry, 0, 32);

            // Copy formatted filename
            memcpy(entry, formatted_name, 11);

            // Set attributes to Archive (normal file)
            entry[11] = 0x20;

            // Allocate first cluster for the file
            uint32_t free_cluster = find_free_cluster();
            if (free_cluster == 0) {
                return -2; // No free cluster
            }

            // Mark cluster as end-of-chain in FAT
            uint32_t fat_sector = fat_start_lba + (free_cluster * 4) / 512;
            uint32_t fat_offset = (free_cluster * 4) % 512;
            if (ide_read48(fat32_drive, fat_sector, 1, sector) != 0)
                return -1;

            *(uint32_t *)(sector + fat_offset) = 0x0FFFFFFF; // end of cluster chain
            if (ide_write48(fat32_drive, fat_sector, 1, sector) != 0)
                return -1;

            // Set first cluster in directory entry (split high and low 16 bits)
            entry[20] = (free_cluster >> 16) & 0xFF;
            entry[21] = (free_cluster >> 24) & 0xFF;
            entry[26] = free_cluster & 0xFF;
            entry[27] = (free_cluster >> 8) & 0xFF;

            // Set file size = 0
            entry[28] = 0;
            entry[29] = 0;
            entry[30] = 0;
            entry[31] = 0;

            // Write back directory sector
            if (ide_write48(fat32_drive, root_lba, 1, sector) != 0)
                return -1;

            return 0; // Success
        }
    }

    return -3; // No free directory entry found
}

int32_t fat32_create_directory(const char *name) {
    uint32_t lba = cluster_to_lba(current_directory_cluster);
    if (ide_read48(fat32_drive, lba, 1, sector) != 0)
        return -1;

    for (int32_t i = 0; i < 512; i += 32) {
        char *entry = (char *)&sector[i];
        if ((uint8_t)entry[0] == 0x00 || (uint8_t)entry[0] == 0xE5) {
            memset(entry, 0, 32);
            memcpy(entry, name, strlen(name) > 11 ? 11 : strlen(name));
            entry[11] = FAT32_ATTR_DIRECTORY;

            uint32_t cluster = 6;
            *(uint16_t *)(entry + 26) = cluster & 0xFFFF;
            *(uint16_t *)(entry + 20) = (cluster >> 16) & 0xFFFF;
            *(uint32_t *)(entry + 28) = 0;

            if (ide_write48(fat32_drive, lba, 1, sector) != 0)
                return -1;

            return 0;
        }
    }
    return -2;
}

int32_t fat32_change_directory(const char *name) {
    uint32_t lba = cluster_to_lba(current_directory_cluster);
    if (ide_read48(fat32_drive, lba, 1, sector) != 0)
        return -1;

    for (int32_t i = 0; i < 512; i += 32) {
        char *entry = (char *)&sector[i];
        if ((uint8_t)entry[0] == 0x00) break;
        if ((uint8_t)entry[0] == 0xE5) continue;

        char entry_name[12];
        memcpy(entry_name, entry, 11);
        entry_name[11] = 0;
        for (int32_t j = 0; j < 11; j++)
            if (entry_name[j] == ' ') entry_name[j] = 0;

        if (strncmp(entry_name, name, 11) != 0) continue;

        if (!(entry[11] & FAT32_ATTR_DIRECTORY))
            return -2;

        current_directory_cluster = (*(uint16_t *)(entry + 26)) | ((*(uint16_t *)(entry + 20)) << 16);
        return 0;
    }
    return -3;
}