#include <asm/io.h>
#include <kernel/con.h>
#include <kernel/kernel.h>
#include <stdint.h>

#define ST_ERR 1
#define ST_DRQ 8
#define ST_SRV 16
#define ST_DF  32
#define ST_RDY 64
#define ST_BSY 128

/* sector buffer */
static char sbuf[512];

void hd_init(void) {
  int i;
  uint8_t s;
  uint16_t *buf = (uint16_t*) sbuf;

  /* select the master drive */
  outb(0x1F6, 0xA0);
  /* send the IDENTIFY command */
  outb(0x1F7, 0xEC);

  while(1) {
    s = inb(0x1F7);
    if(!s) {
      printk("[hd] Master drive not detected on primary bus!\n");
      panic();
    }
    if(s & ST_ERR) {
      if(!inb(0x1F4) && !inb(0x1F5))
        continue;
      printk("[hd] IDENTIFY command error!\n");
      panic();
    }
    if(s & ST_DRQ && !(s & ST_ERR))
      break;
  }

  printk("[hd] Master drive detected on primary bus! (status: 0x%02x)\n", s);

  for(i = 0; i < 256; i++)
    buf[i] = inw(0x1F0);
}

int hd_read(void *buf, uint32_t lba, uint8_t c) {
  int n;
  uint8_t s;
  uint16_t *p = buf;

  if(!c)
    return 0;

  /* select the master disk on the primary bus */
  outb(0x1F6, 0xE0 | ((lba >> 24) & 0xF));
  /* set the number of sectors to be read */
  outb(0x1F2, c);
  /* set the LBA address for the read operation */
  outb(0x1F3, lba);
  outb(0x1F4, lba >> 8);
  outb(0x1F5, lba >> 16);
  /* send the READ command */
  outb(0x1F7, 0x20);

  while(c--) {
    /* wait for the drive */
    while(1) {
      /* poll the drive's status */
      s = inb(0x1F7);
      if(s & ST_DF) {
        printk("[hd] Drive failure!\n");
        panic();
      }
      if(s & ST_ERR)
        return -1;
      if(!(s & ST_BSY) && s & ST_DRQ)
        break;
    }

    /* buffer is full, copy it's data */
    n = 256;
    while(n--)
      *(p++) = inw(0x1F0);
  }

  return 0;
}

int hd_write(void *buf, uint32_t lba, uint8_t c) {
  int n;
  uint8_t s;
  uint16_t *p = buf;

  if(!c)
    return 0;

  /* select the master disk on the primary bus */
  outb(0x1F6, 0xE0 | ((lba >> 24) & 0xF));
  /* set the number of sectors to be written */
  outb(0x1F2, c);
  /* set the LBA address for the write operation */
  outb(0x1F3, lba);
  outb(0x1F4, lba >> 8);
  outb(0x1F5, lba >> 16);
  /* send the WRITE command */
  outb(0x1F7, 0x30);

  while(c--) {
    /* wait for the drive */
    while(1) {
      /* poll the drive's status */
      s = inb(0x1F7);
      if(s & ST_DF) {
        printk("[hd] Drive failure!\n");
        panic();
      }
      if(s & ST_ERR)
        return -1;
      if(!(s & ST_BSY) && s & ST_DRQ)
        break;
    }

    /* refill the drive's buffer */
    n = 256;
    while(n--)
      outw(0x1F0, *(p++));
  }

  return 0;
}