#include <stdio.h>
#include <drivers/pit.h>
#include <port_io.h>

#include <drivers/ide.h>


static int32_t ide_wait(uint16_t io, int32_t check_drq)
{
  for (int32_t i = 0; i < 5000; i++)
  {
    uint8_t status = inb(io + ATA_REG_STATUS);
    if (!(status & ATA_SR_BSY))
    {
      if (!check_drq || (status & ATA_SR_DRQ))
      {
        return 0;
      }
    }
    pit_sleep(2);
  }
  return 1;
}

static int32_t ide_wait_ready(uint16_t io)
{
  for (int32_t i = 0; i < 100000; i++)
  {
    uint8_t status = inb(io + ATA_REG_STATUS);
    if (!(status & ATA_SR_BSY))
    {
      return 0;
    }
  }
  return 1;
}

int32_t ide_identify(uint8_t drive, uint16_t* buffer)
{
  uint16_t io = ATA_PRIMARY_IO;
  uint8_t slave_bit = (drive & 1) << 4;

  ide_wait(io, 0);
  outb(io + ATA_REG_HDDEVSEL, 0xA0 | slave_bit);
  outb(io + ATA_REG_SECCOUNT0, 0);
  outb(io + ATA_REG_LBA0, 0);
  outb(io + ATA_REG_LBA1, 0);
  outb(io + ATA_REG_LBA2, 0);
  outb(io + ATA_REG_COMMAND, 0xEC); /* IDENTIFY DEVICE command */

  uint8_t status = inb(io + ATA_REG_STATUS);
  if (status == 0)
  {
    return -1; /* No device */
  }

  if (ide_wait(io, 1))
  {
    return -2;
  }

  insw(io + ATA_REG_DATA, buffer, 256); /* Read 512 bytes */
  return 0;
}

void ide_initialize(void) {
  outb(ATA_PRIMARY_CTRL, 0x02); /* Disable IRQs from IDE disk controllers TODO: should probably use IRQs soon */
  
  uint16_t identify_buf[256];
  if (ide_identify(0, identify_buf) == 0)
  {
    char model[41];
    for (int i = 0; i < 20; ++i)
    {
      model[i * 2 + 0] = identify_buf[27 + i] >> 8;
      model[i * 2 + 1] = identify_buf[27 + i] & 0xFF;
    }
    model[40] = 0;
    printf("Disk model: %s\n", model);
  }
}

int32_t ide_read48(uint8_t drive, uint64_t lba, uint8_t sector_count, void* buffer) {
  if (sector_count == 0)
  {
    return -1;
  }

  uint16_t io = ATA_PRIMARY_IO;
  uint8_t slave_bit = (drive & 1) << 4;

  ide_wait(io, 0); /* Wait for BSY to clear */

  outb(io + ATA_REG_HDDEVSEL, 0xE0 | slave_bit); /* Select drive */

  
  outb(io + ATA_REG_SECCOUNT0, (sector_count >> 8) & 0xFF);  /* high count */
  outb(io + ATA_REG_LBA0,      (lba >> 40) & 0xFF);           /* LBA5 */
  outb(io + ATA_REG_LBA1,      (lba >> 32) & 0xFF);           /* LBA4 */
  outb(io + ATA_REG_LBA2,      (lba >> 24) & 0xFF);           /* LBA3 */

  outb(io + ATA_REG_SECCOUNT0, sector_count & 0xFF);          /* low count */
  outb(io + ATA_REG_LBA0,      (lba >> 0) & 0xFF);            /* LBA0 */
  outb(io + ATA_REG_LBA1,      (lba >> 8) & 0xFF);            /* LBA1 */
  outb(io + ATA_REG_LBA2,      (lba >> 16) & 0xFF);           /* LBA2 */



  outb(io + ATA_REG_COMMAND, ATA_CMD_READ_SECTORS_EXT);
  

  uint16_t* ptr = (uint16_t*)buffer;
  for (int i = 0; i < sector_count; i++)
  {
    if (ide_wait(io, 1))
    {
      uint8_t status = inb(io + ATA_REG_STATUS);
      printf("IDE DRQ wait failed (read), status = 0x%02X\n", status);
      return -2;
    }

    uint8_t status = inb(io + ATA_REG_STATUS);
    if (status & ATA_SR_ERR)
    {
      return -3;
    }

    insw(io + ATA_REG_DATA, ptr, 256);
    ptr += 256;
  }

  return 0;
}


int32_t ide_write48(uint8_t drive, uint64_t lba, uint8_t sector_count, const void* buffer)
{
  if (sector_count == 0)
  {
    return -1;
  }

  uint16_t io = ATA_PRIMARY_IO;
  uint8_t slave_bit = (drive & 1) << 4;

  ide_wait(io, 0); /* Wait for BSY to clear */

  outb(io + ATA_REG_HDDEVSEL, 0xE0 | slave_bit); /* Select drive */

  /* Send high bytes of sector count and LBA first */
  outb(io + ATA_REG_SECCOUNT0, 0);                     // Sector count high byte
  outb(io + ATA_REG_LBA0, (lba >> 40) & 0xFF);         // LBA[47:40]
  outb(io + ATA_REG_LBA1, (lba >> 32) & 0xFF);         // LBA[39:32]
  outb(io + ATA_REG_LBA2, (lba >> 24) & 0xFF);         // LBA[31:24]

  /* Send low bytes */
  outb(io + ATA_REG_SECCOUNT0, sector_count);          // Sector count low byte
  outb(io + ATA_REG_LBA0, (lba >> 0) & 0xFF);           // LBA[7:0]
  outb(io + ATA_REG_LBA1, (lba >> 8) & 0xFF);           // LBA[15:8]
  outb(io + ATA_REG_LBA2, (lba >> 16) & 0xFF);          // LBA[23:16]

  /* Send write command */
  outb(io + ATA_REG_COMMAND, ATA_CMD_WRITE_SECTORS_EXT); /* 0x34 */

  const uint16_t* ptr = (const uint16_t*)buffer;
  for (int i = 0; i < sector_count; i++)
  {
    if (ide_wait(io, 1))
    {
      uint8_t status = inb(io + ATA_REG_STATUS);
      printf("IDE DRQ wait failed (write), status = 0x%02X\n", status);
      return -2;
    }

    /* Check for drive error */
    uint8_t status = inb(io + ATA_REG_STATUS);
    if (status & ATA_SR_ERR)
    {
      return -3;
    }

    outsw(io + ATA_REG_DATA, ptr, 256); /* 256 words = 512 bytes */
    inb(io + ATA_REG_STATUS); /* Dummy read */
    ptr += 256;
  }

  /* Wait for BSY=0 before issuing FLUSH CACHE */
  if (ide_wait_ready(io))
  {
    return -5;
  }
  
  outb(io + ATA_REG_COMMAND, 0xE7);
  
  if (ide_wait(io, 0)) {
    printf("FLUSH CACHE timeout, status=0x%02X\n", inb(io + ATA_REG_STATUS));
    return -4;
  }

  return 0;
}

int32_t read_sector(uint64_t lba, void* buffer)
{
  return ide_read48(0, lba, 1, buffer);
}

int32_t write_sector(uint64_t lba, const void* buffer)
{
  return ide_write48(0, lba, 1, buffer);
}