#include <asm/interrupt.h>
#include <asm/io.h>
#include <errno.h>
#include <kernel/con.h>
#include <kernel/fs.h>
#include <kernel/kernel.h>
#include <kernel/sched.h>
#include <limits.h>
#include <stdint.h>
#include <string.h>
#include <sys/types.h>

#define FDBASE 0x3F0

#define FDSRA   FDBASE
#define FDSRB   (FDBASE + 1)
#define FDDOR   (FDBASE + 2)
#define FDTDR   (FDBASE + 3)
#define FDMSR   (FDBASE + 4)
#define FDDRS   (FDBASE + 4)
#define FDDATA  (FDBASE + 5)
#define FDDIR   (FDBASE + 7)
#define FDCCR   (FDBASE + 7)

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

#define NTYPES (sizeof(ftypes) / sizeof(struct ftype))
#define NFDS (sizeof(fd) / sizeof(struct fdrive))

/* helper macros for converting LBA into CHS */
#define C(lba, ft) \
  (lba / (ft->h * ft->s))
#define H(lba, ft) \
  ((lba / ft->s) % ft->h)
#define S(lba, ft) \
  ((lba % ft->s) + 1)

#define PRESENT(fd) ((fd.type) && (fd.type != ftypes))

/* TODO: replace any delay loops with sleep_on() */
/* TODO: implement a head scheduling algorithm */
/* TODO: implement write functionality */
/* TODO: implement access to multiple drives */
/* TODO: seek before reading */
/* TODO: check drive ready state in floppy driver (ensure the drive door is closed before any operation) */
/* TODO: cut back on printk() usage */

/* TODO: implement a timeout to prevent system hangs */
#define WAITIRQ(x) ({ \
    irq_fired = 0; \
    (x); \
    while(!irq_fired); \
    irq_fired = 0; \
    })

extern void fd_isr(void);

struct sense_result {
  uint8_t st0;
  uint8_t cyl;
} __attribute__((packed));

struct ftype {
  uint16_t c;
  uint8_t h;
  uint8_t s;
  uint8_t drate; /* 0: 500Kb/s, 1: 300Kb/s, 2: 250Kb/s, 3: 1Mb/s */
  int delay;     /* time to spin up (ms) */
  char *name;
};

struct fdrive {
  struct ftype *type;
  struct task_struct *seek_wait;
  uint8_t cyl;
  uint8_t st0;
};

/* possible floppy disk/drive configurations */
static struct ftype ftypes[] = {
  {  0, 0,  0, 0,   0, "none" },
  { 40, 2,  9, 1, 500, "360KB 5 1/4\"" },
  { 80, 2, 15, 0, 500, "1.2MB 5 1/4\"" },
  { 80, 2,  9, 2, 300, "720KB 3 1/2\"" },
  { 80, 2, 18, 0, 300, "1.44MB 3 1/2\"" },
  { 80, 2, 36, 3, 300, "2.88MB 3 1/2\"" }
};

static volatile int irq_fired;

static struct fdrive fd[4];

static char buffer[512] __attribute__((aligned (0x10000)));

static struct task_struct *read_wait;

static ssize_t read_data(void *buf, size_t len) {
  ssize_t c = 0;
  uint8_t *p = buf;

  if(!buf)
    return 0;

  if(len > SSIZE_MAX)
    len = SSIZE_MAX;

  /*
   * TODO: fix this.
   * This entire routine really is garbage.
   * It doesn't use a timeout (potentially
   * hanging the entire system) and doesn't
   * handle the DIO flag properly.
   */
  while(len--) {
    while(!(inb(FDMSR) & 0x80));

    if(!(inb(FDMSR) & 0x40))
      break;

    *p++ = inb(FDDATA);
    c++;

    /* delay */
    inb(0x80);
  }

  /* empty out any remaining bytes the FDC may have for us */
  while((inb(FDMSR) & 0xC0) == 0xC0)
    inb(FDDATA);

  if(!c)
    return -1;
  return c;
}

static void write_data(void *buf, size_t len) {
  uint8_t *p = buf;

  if(!buf)
    return;

  /* TODO: timeout waiting for RQM flag */
  while(len--) {
    while(!(inb(FDMSR) & 0x80));

    if((inb(FDMSR) & 0x40)) {
      printk("[fd] Expected DIO=0. Drive not ready for data!\n");
      panic();
    }

    outb(FDDATA, *p++);

    /* delay */
    inb(0x80);
  }
}

static int seek(uint8_t drive, int cyl) {
  drive &= 3;

  char cmd_seek[] = {
    0x0F,
    drive,
    cyl
  };

  if(!PRESENT(fd[drive]))
    return -EINVAL;

  WAITIRQ(write_data(cmd_seek, sizeof(cmd_seek)));

  /* wait for the seek command to finish */
  while((inb(FDMSR) & (1 << drive)));

  return 0;
}

static int sense_intr(struct sense_result *sr) {
  ssize_t in;
  char cmd_sense[] = { 0x08 };

  /* send the sense interrupt command */
  write_data(cmd_sense, sizeof(cmd_sense));
  in = read_data(sr, sizeof(struct sense_result));
  if(in < sizeof(struct sense_result))
    return -1;

  printk("[fd] Sense command success (Drive: %c:, CYL: 0x%02x, ST0: 0x%02x)\n", 'A' + (sr->st0 & 3), sr->cyl, sr->st0);
  return 0;
}

static void reset(void) {
  int i;
  int wait = 0;
  ssize_t in;
  struct sense_result sr;

  /* TODO: set timings depending on the drive */
  char cmd_specify[] = {
    0x03,
    0xDF, /* SRT: 3ms, HUT: 240ms */
    0x02  /* HLT: 16ms, NDMA: 0 */
  };

  char cmd_recalibrate[] = { 0x07, 0x00 };

  /* disable interrupts and put the FDC into reset */
  /* TODO: store FDDOR as static variable as the register is write-only */
  printk("[fd] Resetting FDC...\n");
  outb(FDDOR, inb(FDDOR) & 0xF3);

  /* program the data-rate */
  /* TODO: set this depending on floppy type */
  outb(FDDRS, 0);
  outb(FDCCR, 0);

  /* bring the FDC out of reset and re-enable interrupts */
  outb(FDDOR, inb(FDDOR) | 0x0C);

  /* the IRQ handler will issue a sense-interrupt for us */
  WAITIRQ(NULL);

  /* send the specify command */
  write_data(cmd_specify, sizeof(cmd_specify));

  for(i = 0; i < NFDS; i++) {
    if(!PRESENT(fd[i]))
      continue;
    printk("[fd] Recalibrating Drive %c:...\n", 'A' + i);
    cmd_recalibrate[1] = i & 3;
    write_data(cmd_recalibrate, sizeof(cmd_recalibrate));
    wait |= 1 << i;
  }

  /* wait for all drives to finish recalibrating */
  while((inb(FDMSR) & wait));
}

int fd_read_sect(int sect) {
  ssize_t in;
  struct {
    uint8_t st0;
    uint8_t st1;
    uint8_t st2;
    uint8_t c;
    uint8_t h;
    uint8_t r;
    uint8_t n;
  } __attribute__((packed)) res;

  char cmd[] = {
    0x46,
    H(sect, fd[0].type) << 2,
    C(sect, fd[0].type),
    H(sect, fd[0].type),
    S(sect, fd[0].type),
    0x02,
    S(sect, fd[0].type),
    0x1B,
    0xFF
  };

  cli();
  if(read_wait)
    sleep_on(&read_wait);
  sti();

  printk("[fd] Reading sector 0x%01x (C: 0x%01x, H: 0x%01x, S: 0x%01x)\n", sect, cmd[2], cmd[3], cmd[4]);

  WAITIRQ(write_data(cmd, sizeof(cmd)));

  in = read_data(&res, sizeof(res));
  if(in == sizeof(res) && !(res.st0 & 0xC0)) {
    printk("[fd] Successfully read a sector (ST0: 0x%02x, N: 0x%02x)!\n", res.st0, res.n);
    printk("  DATA: %s\n", buffer);
    return 0;
  }

  printk("[fd] Read command failed!\n");
  return -1;
}

void fd_interrupt(void) {
  struct sense_result sr;
  irq_fired = 1;

  wake_up(&read_wait);

  if((inb(FDMSR) & 0xC0) != 0x80)
    return;
  if(sense_intr(&sr) < sizeof(sr))
    return;
  /* deposit result into current drive struct */
  fd[sr.cyl & 3].cyl = sr.cyl;
  fd[sr.st0 & 3].st0 = sr.st0;
}

void fd_block_read(struct buffer *b) {
  int i;
  void *p = b->b_data;

  for(i = SECTS(b->b_block); i < SECTS(b->b_block + 1); i++) {
    if(fd_read_sect(i)) {
      b->b_device = 0;
      return;
    }
    memcpy(p, buffer, sizeof(buffer));
    p += 512;
  }

  b->b_present = 1;
}

void fd_init(void) {
  int i;
  int delay = 0;
  uint8_t tmp;

  /* put the controller into reset */
  outb(FDDOR, 0);

  outb(0x70, 0x10);
  tmp = inb(0x71);

  fd[0].type = &ftypes[(tmp >> 4) > NTYPES ? 0 : (tmp >> 4)];
  fd[1].type = &ftypes[(tmp & 0xF) > NTYPES ? 0 : (tmp & 0xF)];
  fd[2].type = ftypes;
  fd[3].type = ftypes;

  /* setup the drive (motors A and B on) */
  /* TODO: motor timeout */
  outb(FDDOR, 0x3C);

  /* wait for all drives to spin up */
  for(i = 0; i < NFDS; i++)
    if(PRESENT(fd[i]))
      delay = fd[i].type->delay > delay ? fd[i].type->delay : delay;
  delay = (delay / 10) + ticks;
  while(ticks < delay);

  /* initialize the DMA controller */
  outb(0x0F, 0x0F);
  outb(0x0C, 0xFF);
  outb(0x04, (uint32_t) buffer);
  outb(0x04, (uint32_t) buffer >> 8);
  outb(0x0C, 0xFF);
  outb(0x05, sizeof(buffer));
  outb(0x05, sizeof(buffer) >> 8);
  outb(0x81, (uint32_t) buffer >> 16);
  outb(0x0B, 0x56);
  outb(0x0A, 0x02);

  register_isr(0x26, 0, &fd_isr);
  irq_enable(0x40);

  reset();

  printk("[fd] Floppy subsystem initialized (A: %s, B: %s)\n", fd[0].type->name, fd[1].type->name);
}