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

#define SECTS(n) ((n) * (BLOCK_SIZE / 512))
#define BLOCKS(n) ((n) / (BLOCK_SIZE / 512))

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

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

static uint32_t nblocks;

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

  if(lba > nblocks - 1)
    return -1;

  /* select the master disk on the primary bus */
  outb(0x1F6, 0xE0 | ((lba >> 24) & 0xF));
  /* set the number of sectors to be read */
  outb(0x1F2, 1);
  /* 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);

  /* wait for the drive */
  for(;;) {
    /* 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;
}

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

  if(lba > nblocks - 1)
    return -1;

  /* select the master disk on the primary bus */
  outb(0x1F6, 0xE0 | ((lba >> 24) & 0xF));
  /* set the number of sectors to be written */
  outb(0x1F2, 1);
  /* 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);

  /* wait for the drive */
  for(;;) {
    /* 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;
}

int hd_read_block(void *buf, uint16_t block) {
  int ret;
  size_t i;
  uint32_t s = SECTS(block);
  uint32_t e = SECTS(block + 1);

  if(nblocks <= s || nblocks <= e)
    return -1;

  for(i = 0; i < (e - s); i++) {
    ret = read_sect(buf, i + s);

    if(ret < 0)
      return -1;

    buf = ((char*) buf + 512);
  }

  return 0;
}

int hd_write_block(void *buf, uint16_t block) {
  int ret;
  size_t i;
  uint32_t s = SECTS(block);
  uint32_t e = SECTS(block + 1);

  if(nblocks <= s || nblocks <= e)
    return -1;

  for(i = 0; i < (e - s); i++) {
    ret = write_sect(buf, i + s);

    if(ret < 0)
      return -1;

    buf = ((char*) buf + 512);
  }

  return 0;
}

int 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);

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

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

  nblocks = *((uint32_t*) &buf[60]);
  if(!nblocks) {
    printk("[hd] Drive does not support LBA28 addressing!\n");
    return -1;
  }

  printk("[hd] Primary bus master drive initialized! (status: 0x%02x, nsects: 0x%01x)\n", s, nblocks);
  return 0;
}