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

#include <fs/duckfs.h>


static duckfs_superblock_t super;

extern int32_t disk_write(uint64_t lba, const void* buffer, uint32_t count);
extern int32_t disk_read(uint64_t lba, void* buffer, uint32_t count);

int32_t duckfs_mount(void)
{
  if (disk_read(8, &super, 1) != 0)
  {
    return -1;
  }

  if (memcmp(super.magic, DUCKFS_MAGIC, 7) != 0)
  {
    return -2;
  }

  return 0;
}

static int32_t read_header(uint64_t lba, duckfs_file_header_t* hdr)
{
  uint8_t sector[512];
  if (disk_read(lba, sector, 1) != 0)
  {
    return -1;
  }
  memcpy(hdr, sector, sizeof(duckfs_file_header_t));
  return 0;
}

static const char* next_token(const char* path, char* token)
{
  while (*path == '/')
  {
    path++;
  }
  
  if (*path == 0)
  {
    return NULL;
  }

  size_t len = 0;
  while (*path && *path != '/')
  {
    token[len++] = *path++;
  }
  token[len] = 0;

  return path;
}

static int32_t duckfs_alloc_block(uint64_t* out_lba)
{
  if (super.free_list_lba == 0)
  {
    return -1;
  }

  *out_lba = super.free_list_lba;

  uint64_t next;
  if (disk_read(*out_lba, &next, 1) != 0)
  {
    return -2;
  }

  super.free_list_lba = next;

  /* Persist the updated superblock */
  if (disk_write(8, &super, 1) != 0)
  {
    return -3;
  }

  return 0;
}

static int32_t duckfs_alloc_file_data(const void* content, size_t size, uint64_t* first_lba)
{
  const uint8_t* src = (const uint8_t*)content;
  size_t remaining = size;
  uint64_t prev_lba = 0;
  uint64_t head_lba = 0;

  while (remaining > 0)
  {
    uint64_t block_lba;
    if (duckfs_alloc_block(&block_lba) != 0)
    {
      return -1;
    }

    /* Write next block LBA at start (for linking) */
    uint64_t next = 0;
    if (remaining > super.block_size - sizeof(uint64_t))
    {
      if (duckfs_alloc_block(&next) != 0)
      {
        return -2;
      }
    }

    uint8_t buffer[512] = {0};
    memcpy(buffer, &next, sizeof(uint64_t));

    size_t copy_len = super.block_size - sizeof(uint64_t);
    if (copy_len > remaining)
    {
      copy_len = remaining;
    }
    memcpy(buffer + sizeof(uint64_t), src, copy_len);

    if (disk_write(block_lba, buffer, 1) != 0)
    {
      return -3;
    }

    if (prev_lba != 0) {
      if (disk_write(prev_lba, &block_lba, 1) != 0)
      {
        return -4;
      }
    }
    else
    {
      head_lba = block_lba;
    }

    prev_lba = block_lba;
    src += copy_len;
    remaining -= copy_len;
  }

  *first_lba = head_lba;
  return 0;
}


int32_t duckfs_find(const char* path, duckfs_file_header_t* out)
{
  char token[DUCKFS_NAME_LEN];
  duckfs_file_header_t current;

  if (read_header(super.root_header_lba, &current) != 0)
  {
    return -1;
  }

  const char* p = path;
  while ((p = next_token(p, token)))
  {
    if (current.type != DUCKFS_TYPE_DIR)
    {
      return -2;
    }

    uint64_t child = current.child_lba;
    int32_t found = 0;
    while (child != 0)
    {
      duckfs_file_header_t entry;
      if (read_header(child, &entry) != 0)
      {
        return -3;
      }

      if (strncmp(entry.name, token, DUCKFS_NAME_LEN) == 0)
      {
        current = entry;
        found = 1;
        break;
      }
      child = entry.next_file_lba;
    }

    if (!found)
    {
      return -4;
    }
  }

  *out = current;
  return 0;
}

int32_t duckfs_create_file(const char* name, const void* content, size_t size, uint16_t uid, uint16_t gid, uint8_t perms, duckfs_file_header_t* parent_dir)
{
  if (!parent_dir || parent_dir->type != DUCKFS_TYPE_DIR)
  {
    return -1;
  }

  /* Check for existing file */
  duckfs_file_header_t child;
  uint64_t current = parent_dir->child_lba;
  uint64_t prev_lba = 0;

  while (current != 0)
  {
    if (disk_read(current, &child, 1) != 0)
    {
      return -2;
    }
    
    if (strncmp(child.name, name, DUCKFS_NAME_LEN) == 0)
    {
      /* Overwrite content */
      uint64_t data_lba;
      if (duckfs_alloc_file_data(content, size, &data_lba) != 0)
      {
        return -3;
      }
      
      child.first_block_lba = data_lba;
      child.size = size;
      child.modified = 0; /* TODO: timestamp */
      return disk_write(current, &child, 1);
    }
    prev_lba = current;
    current = child.next_file_lba;
  }

  /* Allocate header */
  uint64_t file_hdr_lba;
  if (duckfs_alloc_block(&file_hdr_lba) != 0)
  {
    return -4;
  }

  /* Write content */
  uint64_t data_lba;
  if (duckfs_alloc_file_data(content, size, &data_lba) != 0)
  {
    return -5;
  }

  duckfs_file_header_t new_file = {0};
  strncpy(new_file.name, name, DUCKFS_NAME_LEN);
  new_file.type = DUCKFS_TYPE_FILE;
  new_file.permissions = perms;
  new_file.uid = uid;
  new_file.gid = gid;
  new_file.size = size;
  new_file.first_block_lba = data_lba;
  new_file.modified = 0;

  /* Add to dir */
  if (prev_lba == 0)
  {
    parent_dir->child_lba = file_hdr_lba;
    if (disk_write(parent_dir->first_block_lba, parent_dir, 1) != 0)
    {
      return -6;
    }
  }
  else
  {
    child.next_file_lba = file_hdr_lba;
    if (disk_write(prev_lba, &child, 1) != 0)
    {
      return -7;
    }
  }

  return disk_write(file_hdr_lba, &new_file, 1);
}

int32_t duckfs_create_dir(const char* name, uint16_t uid, uint16_t gid, uint8_t perms, duckfs_file_header_t* parent_dir)
{
  if (!parent_dir || parent_dir->type != DUCKFS_TYPE_DIR)
  {
    return -1;
  }

  /* Allocate header */
  uint64_t hdr_lba;
  if (duckfs_alloc_block(&hdr_lba) != 0)
  {
    return -2;
  }

  duckfs_file_header_t new_dir = {0};
  strncpy(new_dir.name, name, DUCKFS_NAME_LEN);
  new_dir.type = DUCKFS_TYPE_DIR;
  new_dir.permissions = perms;
  new_dir.uid = uid;
  new_dir.gid = gid;
  new_dir.modified = 0;
  new_dir.size = 0;

  /* Write header */
  if (disk_write(hdr_lba, &new_dir, 1) != 0)
  {
    return -3;
  }

  /* Add to parent */
  duckfs_file_header_t sibling;
  uint64_t current = parent_dir->child_lba;
  uint64_t prev = 0;

  while (current != 0)
  {
    if (disk_read(current, &sibling, 1) != 0)
    {
      return -4;
    }
    prev = current;
    current = sibling.next_file_lba;
  }

  if (prev == 0)
  {
    parent_dir->child_lba = hdr_lba;
    if (disk_write(parent_dir->first_block_lba, parent_dir, 1) != 0)
    {
      return -5;
    }
  }
  else
  {
    sibling.next_file_lba = hdr_lba;
    if (disk_write(prev, &sibling, 1) != 0)
    {
      return -6;
    }
  }

  return 0;
}


int32_t duckfs_read_file(const duckfs_file_header_t* file, void* buffer, size_t size)
{
  if (file->type != DUCKFS_TYPE_FILE)
  {
    return -1;
  }

  size_t to_read = size;
  if (to_read > file->size)
  {
     to_read = file->size;
  }

  uint64_t blocks = (to_read + super.block_size - 1) / super.block_size;
  return disk_read(file->first_block_lba, buffer, blocks);
}

int32_t duckfs_write_data(duckfs_file_header_t* file, const void* data, size_t offset, size_t length)
{
  if (!file || file->type != DUCKFS_TYPE_FILE)
  {
    return -1;
  }

  size_t final_size = offset + length;
  if (final_size > file->size)
  {
    /* Grow */
    if (duckfs_append_data(file, ((const uint8_t*)data) + (file->size - offset), final_size - file->size) != 0)
    {
      return -2;
    }
  }

  size_t block_size = super.block_size;
  size_t remaining = length;
  const uint8_t* src = (const uint8_t*)data;

  uint64_t current_lba = file->first_block_lba;
  uint64_t block_offset = 0;

  /* Traverse to block containing `offset` */
  size_t logical = 0;
  while (logical + (block_size - 8) <= offset)
  {
    uint64_t next;
    if (disk_read(current_lba, &next, 1) != 0)
    {
      return -3;
    }
    
    current_lba = next;
    if (current_lba == 0)
    {
      return -4; /* ran out */
    }
    
    logical += block_size - 8;
  }

  /* Start writing */
  while (remaining > 0)
  {
    uint8_t buffer[512];
    if (disk_read(current_lba, buffer, 1) != 0)
    {
      return -5;
    }

    uint64_t next_lba;
    memcpy(&next_lba, buffer, sizeof(uint64_t));

    size_t block_start = (offset > logical) ? (offset - logical) : 0;
    size_t to_copy = block_size - 8 - block_start;
    if (to_copy > remaining) to_copy = remaining;

    memcpy(buffer + 8 + block_start, src, to_copy);
    if (disk_write(current_lba, buffer, 1) != 0)
    {
      return -6;
    }

    src += to_copy;
    remaining -= to_copy;
    offset += to_copy;
    logical += block_size - 8;
    current_lba = next_lba;
  }

  return 0;
}

int32_t duckfs_append_data(duckfs_file_header_t* file, const void* data, size_t length)
{
  if (!file || file->type != DUCKFS_TYPE_FILE)
  {
    return -1;
  }

  size_t block_size = super.block_size;
  size_t data_per_block = block_size - sizeof(uint64_t);

  uint64_t last_lba = file->first_block_lba;
  if (last_lba == 0)
  {
    /* Allocate first block */
    if (duckfs_alloc_block(&last_lba) != 0)
    {
      return -2;
    }
    
    file->first_block_lba = last_lba;
  }
  else
  {
    /* Traverse to last block */
    uint64_t next;
    while (1)
    {
      if (disk_read(last_lba, &next, 1) != 0)
      {
        return -3;
      }
      
      memcpy(&next, &next, sizeof(uint64_t));
      if (next == 0)
      {
        break;
      }
      
      last_lba = next;
    }
  }

  size_t offset_in_block = file->size % data_per_block;
  size_t remaining = length;
  const uint8_t* src = (const uint8_t*)data;

  /* Write to the partially filled last block */
  if (offset_in_block != 0)
  {
    uint8_t buffer[512];
    if (disk_read(last_lba, buffer, 1) != 0)
    {
      return -4;
    }
    
    size_t to_copy = data_per_block - offset_in_block;
    if (to_copy > remaining)
    {
      to_copy = remaining;
    }
    
    memcpy(buffer + sizeof(uint64_t) + offset_in_block, src, to_copy);
    if (disk_write(last_lba, buffer, 1) != 0)
    {
      return -5;
    }
    
    file->size += to_copy;
    remaining -= to_copy;
    src += to_copy;
  }

  /* Write remaining data to new blocks */
  while (remaining > 0)
  {
    uint64_t new_lba;
    if (duckfs_alloc_block(&new_lba) != 0)
    {
      return -6;
    }

    /* Link previous block to this one */
    if (disk_write(last_lba, &new_lba, 1) != 0)
    {
      return -7;
    }

    uint8_t buffer[512] = {0};
    uint64_t next = 0;
    memcpy(buffer, &next, sizeof(uint64_t));

    size_t to_copy = (remaining > data_per_block) ? data_per_block : remaining;
    memcpy(buffer + sizeof(uint64_t), src, to_copy);
    if (disk_write(new_lba, buffer, 1) != 0)
    {
      return -8;
    }

    last_lba = new_lba;
    src += to_copy;
    remaining -= to_copy;
    file->size += to_copy;
  }

  return 0;
}

int32_t duckfs_truncate(duckfs_file_header_t* file, size_t new_size)
{
  if (!file || file->type != DUCKFS_TYPE_FILE)
  {
    return -1;
  }

  size_t block_size = super.block_size;
  size_t data_per_block = block_size - sizeof(uint64_t);

  size_t needed_blocks = (new_size + data_per_block - 1) / data_per_block;
  uint64_t current = file->first_block_lba;
  uint64_t prev = 0;

  for (size_t i = 0; i < needed_blocks; i++)
  {
    if (current == 0)
    {
      break;
    }
    
    uint64_t next;
    if (disk_read(current, &next, 1) != 0)
    {
      return -2;
    }
    
    memcpy(&next, &next, sizeof(uint64_t));
    prev = current;
    current = next;
  }

  /* `current` is now first block to free */
  while (current != 0)
  {
    uint64_t next;
    if (disk_read(current, &next, 1) != 0)
    {
      return -3;
    }
    
    memcpy(&next, &next, sizeof(uint64_t));

    /* Write freelist entry */
    if (disk_write(current, &super.free_list_lba, 1) != 0)
    {
      return -4;
    }
    
    super.free_list_lba = current;
    current = next;
  }

  /* Finalize */
  if (disk_write(8, &super, 1) != 0)
  {
    return -5;
  }

  /* Cut block chain */
  if (prev != 0)
  {
    uint64_t zero = 0;
    if (disk_write(prev, &zero, 1) != 0)
    {
      return -6;
    }
  }

  file->size = new_size;
  return 0;
}

int32_t duckfs_delete_file(const char* name, duckfs_file_header_t* parent_dir)
{
  if (!parent_dir || parent_dir->type != DUCKFS_TYPE_DIR)
  {
    return -1;
  }

  duckfs_file_header_t child;
  uint64_t prev_lba = 0;
  uint64_t current_lba = parent_dir->child_lba;

  /* Search for file by name */
  while (current_lba != 0)
  {
    if (disk_read(current_lba, &child, 1) != 0)
    {
      return -2;
    }

    if (strncmp(child.name, name, DUCKFS_NAME_LEN) == 0)
    {
      break;
    }

    prev_lba = current_lba;
    current_lba = child.next_file_lba;
  }

  if (current_lba == 0)
  {
    return -3; /* not found */
  }

  /* Remove from parent's linked list */
  if (prev_lba == 0) 
  {
    /* It's the first child */
    parent_dir->child_lba = child.next_file_lba;
    if (disk_write(parent_dir->first_block_lba, parent_dir, 1) != 0)
    {
      return -4;
    }
  }
  else
  {
    duckfs_file_header_t prev;
    if (disk_read(prev_lba, &prev, 1) != 0)
    {
      return -5;
    }
    
    prev.next_file_lba = child.next_file_lba;
    if (disk_write(prev_lba, &prev, 1) != 0)
    {
      return -6;
    }
  }

  /* Free file data blocks */
  uint64_t current = child.first_block_lba;
  while (current != 0)
  {
    uint64_t next;
    if (disk_read(current, &next, 1) != 0)
    {
      return -7;
    }

    /* Insert this block into the freelist */
    if (disk_write(current, &super.free_list_lba, 1) != 0)
    {
      return -8;
    }
    
    super.free_list_lba = current;
    current = next;
  }

  /* Free file header block */
  if (disk_write(current_lba, &super.free_list_lba, 1) != 0)
  {
    return -9;
  }
  
  super.free_list_lba = current_lba;

  /* Save updated superblock */
  if (disk_write(8, &super, 1) != 0)
  {
    return -10;
  }

  return 0;
}

int32_t duckfs_delete_dir(const char* name, duckfs_file_header_t* parent_dir)
{
  if (!parent_dir || parent_dir->type != DUCKFS_TYPE_DIR)
  {
    return -1;
  }

  duckfs_file_header_t dir;
  uint64_t prev_lba = 0;
  uint64_t current_lba = parent_dir->child_lba;

  /* Find directory */
  while (current_lba != 0)
  {
    if (disk_read(current_lba, &dir, 1) != 0)
    {
      return -2;
    }
    
    if (strncmp(dir.name, name, DUCKFS_NAME_LEN) == 0 && dir.type == DUCKFS_TYPE_DIR)
    {
      break;
    }
    
    prev_lba = current_lba;
    current_lba = dir.next_file_lba;
  }

  if (current_lba == 0)
  {
    return -3; /* not found */
  }

  /* Delete contents recursively */
  uint64_t child_lba = dir.child_lba;
  while (child_lba != 0)
  {
    duckfs_file_header_t child;
    if (disk_read(child_lba, &child, 1) != 0)
    {
      return -4;
    }

    uint64_t next_lba = child.next_file_lba;

    if (child.type == DUCKFS_TYPE_FILE)
    {
      if (duckfs_delete_file(child.name, &dir) != 0)
      {
        return -5;
      }
    }
    else if (child.type == DUCKFS_TYPE_DIR)
    {
      if (duckfs_delete_dir(child.name, &dir) != 0)
      {
        return -6;
      }
    }

    child_lba = next_lba;
  }

  /* Remove dir from parent */
  if (prev_lba == 0)
  {
    parent_dir->child_lba = dir.next_file_lba;
    if (disk_write(parent_dir->first_block_lba, parent_dir, 1) != 0)
    {
      return -7;
    }
  }
  else
  {
    duckfs_file_header_t prev;
    if (disk_read(prev_lba, &prev, 1) != 0)
    {
      return -8;
    }
    
    prev.next_file_lba = dir.next_file_lba;
    if (disk_write(prev_lba, &prev, 1) != 0)
    {
      return -9;
    }
  }

  /* Free directory header block */
  if (disk_write(current_lba, &super.free_list_lba, 1) != 0)
  {
    return -10;
  }
  
  super.free_list_lba = current_lba;

  /* Save updated superblock */
  return disk_write(8, &super, 1);
}