#include <device/def_device_op.h>
#include <device/device.h>
#include <device/storage/partitions.h>
#include <device/storage/usbdrv.h>
#include <device/usb/xhci.h>
#include <devices.h>
#include <libk/align.h>
#include <libk/endianess.h>
#include <libk/string.h>
#include <limine/requests.h>
#include <mm/malloc.h>
#include <mm/pmm.h>
#include <proc/proc.h>
#include <proc/reschedule.h>
#include <status.h>
#include <sys/debug.h>
#include <sys/stall.h>

/* REF:
 * https://en.wikipedia.org/wiki/SCSI
 * https://www.seagate.com/files/staticfiles/support/docs/manual/Interface%20manuals/100293068j.pdf
 */

#define SCSI_TEST_UNIT_READY 0x00
#define SCSI_READ_CAPACITY10 0x25
#define SCSI_READ10          0x28
#define SCSI_WRITE10         0x2A

static int usb_ms_send_cbw(struct xhci* xhci, struct xhci_usb_device* usb_device,
                           uint8_t bulk_out_endpoint, uintptr_t cbw_phys) {
  return xhci_bulk_transfer(xhci, usb_device, bulk_out_endpoint, cbw_phys,
                            sizeof(struct usb_ms_cbw));
}

static int usb_ms_read_data(struct xhci* xhci, struct xhci_usb_device* usb_device,
                            uint8_t bulk_in_endpoint, uintptr_t buffer_phys, size_t buffer_size) {
  return xhci_bulk_transfer(xhci, usb_device, bulk_in_endpoint, buffer_phys, buffer_size);
}

static int usb_ms_write_data(struct xhci* xhci, struct xhci_usb_device* usb_device,
                             uint8_t bulk_out_endpoint, uintptr_t buffer_phys, size_t buffer_size) {
  return xhci_bulk_transfer(xhci, usb_device, bulk_out_endpoint, buffer_phys, buffer_size);
}

static int usb_ms_read_csw(struct xhci* xhci, struct xhci_usb_device* usb_device,
                           uint8_t bulk_in_endpoint, uintptr_t csw_phys) {
  return xhci_bulk_transfer(xhci, usb_device, bulk_in_endpoint, csw_phys,
                            sizeof(struct usb_ms_csw));
}

static int usb_ms_scsi_read_capacity(struct xhci* xhci, struct xhci_usb_device* usb_device,
                                     uint8_t bulk_in_endpoint, uint8_t bulk_out_endpoint,
                                     size_t* sector_size, size_t* sector_count) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  int ret;

  uintptr_t cbw_phys = pmm_alloc(1);
  struct usb_ms_cbw* cbw = (struct usb_ms_cbw*)(cbw_phys + (uintptr_t)hhdm->offset);
  memset(cbw, 0, PAGE_SIZE);

  uintptr_t data_phys = pmm_alloc(1);
  uint8_t* data = (uint8_t*)(data_phys + (uintptr_t)hhdm->offset);
  memset(data, 0, PAGE_SIZE);

  uintptr_t csw_phys = pmm_alloc(1);
  struct usb_ms_csw* csw = (struct usb_ms_csw*)(csw_phys + (uintptr_t)hhdm->offset);
  memset(csw, 0, PAGE_SIZE);

  uint8_t cdb[10] = {
      SCSI_READ_CAPACITY10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  };

  cbw->signature = USB_CBW_SIGNATURE;
  cbw->tag = 0x12345678;
  cbw->dir = 0x80;
  cbw->lun = 0;
  cbw->cmd_len = sizeof(cdb);
  cbw->length = 8;
  memcpy(cbw->data, cdb, sizeof(cdb));

  if ((ret = usb_ms_send_cbw(xhci, usb_device, bulk_out_endpoint, cbw_phys)) < 0) {
    goto done;
  }

  if ((ret = usb_ms_read_data(xhci, usb_device, bulk_in_endpoint, data_phys, 8)) < 0) {
    goto done;
  }

  if ((ret = usb_ms_read_csw(xhci, usb_device, bulk_in_endpoint, csw_phys)) < 0) {
    goto done;
  }

  if (csw->status != 0 || csw->signature != USB_CSW_SIGNATURE) {
    DEBUG("Failed to read capacity/sect size. CSW status=%u residue=%u\n", csw->status,
          csw->residue);
    ret = -ST_USB_CTRL_ERROR;
    goto done;
  }

  uint32_t last_lba = BSWAP32(*(uint32_t*)&data[0]);
  uint32_t blk_len = BSWAP32(*(uint32_t*)&data[4]);

  *sector_size = blk_len;
  *sector_count = (size_t)last_lba + 1;

  ret = ST_OK;

done:
  pmm_free(cbw_phys, 1);
  pmm_free(data_phys, 1);
  pmm_free(csw_phys, 1);
  return ret;
}

static int usb_ms_scsi_test_unit_ready(struct xhci* xhci, struct xhci_usb_device* usb_device,
                                       uint8_t bulk_in_endpoint, uint8_t bulk_out_endpoint) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  int ret;

  uintptr_t cbw_phys = pmm_alloc(1);
  struct usb_ms_cbw* cbw = (struct usb_ms_cbw*)(cbw_phys + (uintptr_t)hhdm->offset);
  memset(cbw, 0, PAGE_SIZE);

  uintptr_t csw_phys = pmm_alloc(1);
  struct usb_ms_csw* csw = (struct usb_ms_csw*)(csw_phys + (uintptr_t)hhdm->offset);
  memset(csw, 0, PAGE_SIZE);

  uint8_t cdb[10] = {
      SCSI_TEST_UNIT_READY, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  };

  cbw->signature = USB_CBW_SIGNATURE;
  cbw->tag = 0x12345678;
  cbw->dir = 0x00;
  cbw->lun = 0;
  cbw->cmd_len = sizeof(cdb);
  cbw->length = 0;
  memcpy(cbw->data, cdb, sizeof(cdb));

  if ((ret = usb_ms_send_cbw(xhci, usb_device, bulk_out_endpoint, cbw_phys)) < 0) {
    goto done;
  }

  if ((ret = usb_ms_read_csw(xhci, usb_device, bulk_in_endpoint, csw_phys)) < 0) {
    goto done;
  }

  if (csw->status != 0 || csw->signature != USB_CSW_SIGNATURE) {
    ret = -ST_USB_CTRL_ERROR;
    goto done;
  }

  ret = ST_OK;

done:
  pmm_free(cbw_phys, 1);
  pmm_free(csw_phys, 1);
  return ret;
}

static int usb_ms_scsi_read(struct xhci* xhci, struct xhci_usb_device* usb_device,
                            uint8_t bulk_in_endpoint, uint8_t bulk_out_endpoint,
                            uint8_t* out_buffer, uint32_t lba, uint16_t sector_count,
                            size_t sector_size) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  int ret;

  uintptr_t cbw_phys = pmm_alloc(1);
  struct usb_ms_cbw* cbw = (struct usb_ms_cbw*)(cbw_phys + (uintptr_t)hhdm->offset);
  memset(cbw, 0, PAGE_SIZE);

  size_t data_pages = div_align_up(sector_count * sector_size, PAGE_SIZE);
  uintptr_t data_phys = pmm_alloc(data_pages);
  uint8_t* data = (uint8_t*)(data_phys + (uintptr_t)hhdm->offset);
  memset(data, 0, PAGE_SIZE * data_pages);

  uintptr_t csw_phys = pmm_alloc(1);
  struct usb_ms_csw* csw = (struct usb_ms_csw*)(csw_phys + (uintptr_t)hhdm->offset);
  memset(csw, 0, PAGE_SIZE);

  uint8_t cdb[10] = {
      SCSI_READ10,
      0x00,
      /* LBA */
      (uint8_t)((lba & 0xFF000000) >> 24),
      (uint8_t)((lba & 0x00FF0000) >> 16),
      (uint8_t)((lba & 0x0000FF00) >> 8),
      (uint8_t)((lba & 0xFF)),
      0x00,
      /* Transfer Length */
      (uint8_t)((sector_count & 0xFF00) >> 8),
      (uint8_t)((sector_count & 0x00FF)),
      0x00,
  };

  cbw->signature = USB_CBW_SIGNATURE;
  cbw->tag = 0x12345678;
  cbw->dir = 0x80;
  cbw->lun = 0;
  cbw->cmd_len = sizeof(cdb);
  cbw->length = sector_count * sector_size;
  memcpy(cbw->data, cdb, sizeof(cdb));

  if ((ret = usb_ms_send_cbw(xhci, usb_device, bulk_out_endpoint, cbw_phys)) < 0) {
    goto done;
  }

  if ((ret = usb_ms_read_data(xhci, usb_device, bulk_in_endpoint, data_phys,
                              sector_count * sector_size)) < 0) {
    goto done;
  }

  if ((ret = usb_ms_read_csw(xhci, usb_device, bulk_in_endpoint, csw_phys)) < 0) {
    goto done;
  }

  if (csw->status != 0 || csw->signature != USB_CSW_SIGNATURE) {
    DEBUG("Failed to read sectors. CSW status=%u residue=%u\n", csw->status, csw->residue);
    ret = -ST_USB_CTRL_ERROR;
    goto done;
  }

  memcpy(out_buffer, data, sector_count * sector_size);

  ret = ST_OK;

done:
  pmm_free(cbw_phys, 1);
  pmm_free(data_phys, data_pages);
  pmm_free(csw_phys, 1);
  return ret;
}

static int usb_ms_scsi_write(struct xhci* xhci, struct xhci_usb_device* usb_device,
                             uint8_t bulk_in_endpoint, uint8_t bulk_out_endpoint,
                             uint8_t* in_buffer, uint32_t lba, uint16_t sector_count,
                             size_t sector_size) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;
  int ret;

  uintptr_t cbw_phys = pmm_alloc(1);
  struct usb_ms_cbw* cbw = (struct usb_ms_cbw*)(cbw_phys + (uintptr_t)hhdm->offset);
  memset(cbw, 0, PAGE_SIZE);

  uintptr_t buffer_phys = (uintptr_t)in_buffer - (uintptr_t)hhdm->offset;

  uintptr_t csw_phys = pmm_alloc(1);
  struct usb_ms_csw* csw = (struct usb_ms_csw*)(csw_phys + (uintptr_t)hhdm->offset);
  memset(csw, 0, PAGE_SIZE);

  uint8_t cdb[10] = {
      SCSI_WRITE10,
      0x00,
      /* LBA */
      (uint8_t)((lba & 0xFF000000) >> 24),
      (uint8_t)((lba & 0x00FF0000) >> 16),
      (uint8_t)((lba & 0x0000FF00) >> 8),
      (uint8_t)((lba & 0xFF)),
      0x00,
      /* Transfer Length */
      (uint8_t)((sector_count & 0xFF00) >> 8),
      (uint8_t)((sector_count & 0x00FF)),
      0x00,
  };

  cbw->signature = USB_CBW_SIGNATURE;
  cbw->tag = 0x12345678;
  cbw->dir = 0x00;
  cbw->lun = 0;
  cbw->cmd_len = sizeof(cdb);
  cbw->length = sector_count * sector_size;
  memcpy(cbw->data, cdb, sizeof(cdb));

  if ((ret = usb_ms_send_cbw(xhci, usb_device, bulk_out_endpoint, cbw_phys)) < 0) {
    goto done;
  }

  if ((ret = usb_ms_write_data(xhci, usb_device, bulk_out_endpoint, buffer_phys,
                               sector_count * sector_size)) < 0) {
    goto done;
  }

  if ((ret = usb_ms_read_csw(xhci, usb_device, bulk_in_endpoint, csw_phys)) < 0) {
    goto done;
  }

  if (csw->status != 0 || csw->signature != USB_CSW_SIGNATURE) {
    DEBUG("Failed to read sectors. CSW status=%u residue=%u\n", csw->status, csw->residue);
    ret = -ST_USB_CTRL_ERROR;
    goto done;
  }

  ret = ST_OK;

done:
  pmm_free(cbw_phys, 1);
  pmm_free(csw_phys, 1);
  return ret;
}

static void usbdrv_setup_endpoints(struct usbdrv* usbdrv) {
  for (size_t ep = 0; ep < usbdrv->usb_device->endpoints_count; ep++) {
    struct xhci_usb_device_endpoint* endpoint = &usbdrv->usb_device->endpoints[ep];

    uint8_t type = endpoint->desc.attrs & 0x03;
    bool in = (endpoint->desc.endpoint_addr & 0x80) != 0;

    uint32_t xhci_type = 0;

    if (type == 2) {
      if (in)
        xhci_type = XHCI_EP_BULK_IN;
      else
        xhci_type = XHCI_EP_BULK_OUT;
    } else if (type == 3) {
      if (in)
        xhci_type = XHCI_EP_INTR_IN;
      else
        xhci_type = XHCI_EP_INTR_OUT;
    }

    if (xhci_type == XHCI_EP_BULK_IN)
      usbdrv->bulk_in = endpoint->desc;
    else if (xhci_type == XHCI_EP_BULK_OUT)
      usbdrv->bulk_out = endpoint->desc;
  }
}

DEFINE_DEVICE_OP(usbdrv_read) {
  if (a1 == NULL || a2 == NULL || a3 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t sector = *(size_t*)a1;
  size_t sector_count = *(size_t*)a2;
  uint8_t* buffer = a3;

  struct usbdrv* usbdrv = device->udata;

  int ret = usb_ms_scsi_read(usbdrv->xhci, usbdrv->usb_device, usbdrv->bulk_in.endpoint_addr,
                             usbdrv->bulk_out.endpoint_addr, buffer, sector, sector_count,
                             usbdrv->sector_size);
  return ret;
}

DEFINE_DEVICE_OP(usbdrv_write) {
  if (a1 == NULL || a2 == NULL || a3 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t sector = *(size_t*)a1;
  size_t sector_count = *(size_t*)a2;
  uint8_t* buffer = a3;

  struct usbdrv* usbdrv = device->udata;

  int ret = usb_ms_scsi_write(usbdrv->xhci, usbdrv->usb_device, usbdrv->bulk_in.endpoint_addr,
                              usbdrv->bulk_out.endpoint_addr, buffer, sector, sector_count,
                              usbdrv->sector_size);
  return ret;
}

DEFINE_DEVICE_OP(usbdrv_get_device_type) {
  if (a1 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  int* device_type = (int*)a1;

  *device_type = XDRV_TYPE_USBDRV;

  return ST_OK;
}

DEFINE_DEVICE_OP(usbdrv_get_sector_size) {
  if (a1 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t* secsize = (size_t*)a1;

  struct usbdrv* usbdrv = device->udata;

  *secsize = usbdrv->sector_size;

  return ST_OK;
}

DEFINE_DEVICE_OP(usbdrv_get_size) {
  if (a1 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t* size = (size_t*)a1;

  struct usbdrv* usbdrv = device->udata;

  *size = (usbdrv->sector_size * usbdrv->sector_count);

  return ST_OK;
}

DEFINE_DEVICE_OP(usbdrv_partition_rescan) {
  struct list_node_link *subdevice_link, *tmp_subdevice_link;
  list_foreach(device->subdevices, subdevice_link, tmp_subdevice_link) {
    struct device* subdevice = list_entry(subdevice_link, struct device, subdevices_link);

    list_remove(device->subdevices, &subdevice->subdevices_link);

    device_delete(subdevice->key, proc, rctx);
  }

  int r = device_probe_partitions(proc, rctx, device);

  return r;
}

DEFINE_DEVICE_INIT(usbdrv_init) {
  int ret;
  struct usbdrv_init* init = arg;

  struct usbdrv* usbdrv = malloc(sizeof(*usbdrv));

  if (usbdrv == NULL)
    return false;

  memset(usbdrv, 0, sizeof(*usbdrv));
  usbdrv->xhci = init->xhci;
  usbdrv->usb_device = init->usb_device;

  device->udata = usbdrv;

  usbdrv_setup_endpoints(usbdrv);

  ret = usb_ms_scsi_test_unit_ready(usbdrv->xhci, usbdrv->usb_device, usbdrv->bulk_in.endpoint_addr,
                                    usbdrv->bulk_out.endpoint_addr);

  if (ret < 0) {
    stall_ms(100);

    ret =
        usb_ms_scsi_test_unit_ready(usbdrv->xhci, usbdrv->usb_device, usbdrv->bulk_in.endpoint_addr,
                                    usbdrv->bulk_out.endpoint_addr);
  }

  if (ret < 0) {
    free(usbdrv);
    return false;
  }

  ret = usb_ms_scsi_read_capacity(usbdrv->xhci, usbdrv->usb_device, usbdrv->bulk_in.endpoint_addr,
                                  usbdrv->bulk_out.endpoint_addr, &usbdrv->sector_size,
                                  &usbdrv->sector_count);

  if (ret < 0) {
    free(usbdrv);
    return false;
  }

  DEBUG("sector_count = %zu, sector_size = %u\n", usbdrv->sector_count, usbdrv->sector_size);

  return true;
}

DEFINE_DEVICE_FINI(usbdrv_fini) {
  struct usbdrv* usbdrv = device->udata;

  free(usbdrv);
}