#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("========= Start FAT info =========\n");
  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);
  printf("========= End FAT info   =========\n");

  return 0;
}

void print_83(const char *s) {
  for (int i = 0; i < 11; i++) {
    char c = s[i];
    if (c >= 32 && c <= 126) {
      printf("%c", c); // or your kernel's character output function
    } else {
      printf("\\x%02X", (unsigned char)c); // hex-escape non-printables
    }
  }
  printf("\n");
}


static uint32_t cluster_to_lba(uint32_t cluster)
{
  printf("cluster_heap_lba: %u\n", cluster_heap_lba);
  return cluster_heap_lba + (cluster - 2) * bpb.sectors_per_cluster;
}

void format_83_name(const char* input, char* output)
{
  printf("Input name: %s, Input length: %i\n", input, strlen(input));

  memset(output, ' ', 11);
  int32_t i = 0, j = 0;
  while (input[i] && j < 11)
  {
    if (input[i] == '.')
    {
      j = 8;
      i++;
      continue;
    }
    if (j < 11)
    {
      output[j++] = toupper((unsigned char)input[i++]);
    }
  }
}


static uint32_t find_free_cluster(void)
{
  /* 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)
{
  char formatted[11];
  memset(formatted, 0xAB, sizeof(formatted));
  format_83_name(name, formatted);

  uint32_t dir_cluster = current_directory_cluster;

  while (dir_cluster < 0x0FFFFFF8)
  {
    uint32_t lba = cluster_to_lba(dir_cluster);

    for (uint8_t sector_idx = 0; sector_idx < bpb.sectors_per_cluster; sector_idx++)
    {
      if (ide_read48(fat32_drive, lba + sector_idx, 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 || (uint8_t)entry[11] == 0x0F)
        {
          continue;
        }
        
        printf("here\n");
        
        char entry_name[12] = { '\0' };
        memcpy((void*)entry_name, (void*)entry, 11);
        
        int32_t err = memcmp(entry_name, formatted, (size_t)11);
        
        if (err != 0)
        {
          printf("memcmp: %i\n", err);
          continue;
        }

        /* File found */
        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 file_lba = cluster_to_lba(cluster);
          if (ide_read48(fat32_drive, file_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;
      }
    }

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

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

  return -2; /* File not found */
}


int32_t fat32_write_file(const char *name, const uint8_t *buffer, uint32_t len)
{
  uint8_t fat_sector_buf[512];

  char formatted[11];
  memset(formatted, 0xAB, sizeof(formatted));
  format_83_name(name, formatted);
  
  uint32_t dir_cluster = current_directory_cluster;
  
  while (dir_cluster < 0x0FFFFFF8)
  {
    uint32_t lba = cluster_to_lba(dir_cluster);

    for (uint8_t sector_idx = 0; sector_idx < bpb.sectors_per_cluster; sector_idx++)
    {
      // Read directory sector
      if (ide_read48(fat32_drive, lba + sector_idx, 1, sector) != 0)
      {
        return -1;
      }

      // Find matching entry offset inside this sector
      int32_t entry_offset = -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 || (uint8_t)entry[11] == 0x0F)
        {
          continue;
        }
        
        char entry_name[12] = {0};
        memcpy(entry_name, entry, 11);

        if (memcmp(entry_name, formatted, 11) == 0)
        {
          entry_offset = i;
          break;
        }
      }

      if (entry_offset == -1)
      {
        // Not found in this sector, continue to next
        continue;
      }

      // Entry found, write file data clusters
      uint8_t *entry = &sector[entry_offset];
      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 file_lba = cluster_to_lba(cluster);
        if (ide_write48(fat32_drive, file_lba, bpb.sectors_per_cluster, buffer + bytes_written) != 0)
        {
          return -1;
        }

        bytes_written += bpb.sectors_per_cluster * 512;

        if (bytes_written < len)
        {
          // Read FAT sector for current cluster
          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, fat_sector_buf) != 0)
          {
            return -1;
          }

          uint32_t next_cluster = find_free_cluster();
          if (next_cluster == 0)
          {
            return -3; // No free cluster
          }

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

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

          cluster = next_cluster;
        }
        else
        {
          // Mark end of cluster chain
          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, fat_sector_buf) != 0)
          {
            return -1;
          }

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

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

          break;
        }
      }

      // **Re-read the directory sector before updating file size**
      if (ide_read48(fat32_drive, lba + sector_idx, 1, sector) != 0)
      {
        return -1;
      }

      // Update file size in directory entry again after re-read
      entry = &sector[entry_offset];
      *(uint32_t *)(entry + 28) = len;

      // Write back updated directory sector
      if (ide_write48(fat32_drive, lba + sector_idx, 1, sector) != 0)
      {
        return -1;
      }
      
      // Debug print updated directory entry
      printf("Directory entry after write (offset %d):\n", entry_offset);
      for (int j = 0; j < 32; j++) {
          uint8_t b = (uint8_t)entry[j];
          printf("%02X ", b);
          if ((j + 1) % 16 == 0) printf("  ");
      }
      printf("\n");
      for (int j = 0; j < 32; j++) {
          char c = entry[j];
          printf("%c", (c >= 32 && c <= 126) ? c : '.');
      }
      printf("\n");

      return bytes_written;
    }

    // Move to next cluster of the directory
    uint32_t fat_sector = dir_cluster * 4 / 512;
    uint32_t fat_offset = dir_cluster * 4 % 512;

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

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

  return -2; // File not found
}



int32_t fat32_create_file(const char *name)
{
  uint8_t dir_sector[512];
  uint8_t fat_sector_buf[512];

  uint32_t root_lba = cluster_to_lba(current_directory_cluster);

  if (ide_read48(fat32_drive, root_lba, 1, dir_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 = &dir_sector[i];
    if (entry[0] == 0x00 || entry[0] == 0xE5)
    {
      printf("Writing entry at offset %d\n", i);
      memset(entry, 0, 32);
      memcpy(entry, formatted_name, 11);
      entry[11] = 0x20;

      uint32_t free_cluster = find_free_cluster();
      if (free_cluster == 0)
      {
        return -2;
      }

      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, fat_sector_buf) != 0)
      {
        return -1;
      }

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

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

      /* Set first cluster */
      entry[20] = (free_cluster >> 16) & 0xFF;
      entry[21] = (free_cluster >> 24) & 0xFF;
      entry[26] = free_cluster & 0xFF;
      entry[27] = (free_cluster >> 8) & 0xFF;

      /* File size = 0 */
      memset(entry + 28, 0, 4);

      if (ide_write48(fat32_drive, root_lba, 1, dir_sector) != 0)
      {
        return -1;
      }

      uint8_t temp[512];
      ide_read48(fat32_drive, cluster_to_lba(bpb.root_cluster), 1, temp);
      printf("Root directory cluster 2 dump (first 64 bytes):\n");
      for (int i = 0; i < 64; i++) {
          printf("%02X ", temp[i]);
          if ((i + 1) % 16 == 0) printf("\n");
      }

      
      return 0;
    }
  }
  
  return -8;
}

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 dir_cluster = current_directory_cluster;

  while (dir_cluster < 0x0FFFFFF8)
  {
    uint32_t lba = cluster_to_lba(dir_cluster);

    for (uint8_t sector_idx = 0; sector_idx < bpb.sectors_per_cluster; sector_idx++)
    {
      if (ide_read48(fat32_drive, lba + sector_idx, 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;
        }

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

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

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

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

  return -3;
}