#include <device/def_device_op.h>
#include <device/device.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
 */

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

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,
                             uint64_t* lockflags) {
  return xhci_bulk_transfer (xhci, usb_device, bulk_in_endpoint, buffer_phys, buffer_size,
                             lockflags);
}

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

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,
                                      uint64_t* lockflags) {
  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] = {
      0x25, 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, lockflags)) < 0) {
    goto done;
  }

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

  if ((ret = usb_ms_read_csw (xhci, usb_device, bulk_in_endpoint, csw_phys, lockflags)) < 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 = BE2LE32 (*(uint32_t*)&data[0]);
  uint32_t blk_len = BE2LE32 (*(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,
                                        uint64_t* lockflags) {
  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] = {
      0x00, 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, lockflags)) < 0) {
    goto done;
  }

  if ((ret = usb_ms_read_csw (xhci, usb_device, bulk_in_endpoint, csw_phys, lockflags)) < 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, uint64_t* lockflags) {
  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] = {
      0x28,
      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, lockflags)) < 0) {
    goto done;
  }

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

  if ((ret = usb_ms_read_csw (xhci, usb_device, bulk_in_endpoint, csw_phys, lockflags)) < 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 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) {
  uint64_t fd;

  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;

  spin_unlock (&device->lock, *lockflags);
  spin_lock (&usbdrv->xhci->device->lock, &fd);
  spin_lock (&device->lock, lockflags);

  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, &fd);

  spin_unlock (&device->lock, *lockflags);
  spin_unlock (&usbdrv->xhci->device->lock, fd);
  spin_lock (&device->lock, lockflags);

  return ret;
}

DEFINE_DEVICE_OP (usbdrv_write) { return -ST_XDRV_WRITE_ERROR; }

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_INIT (usbdrv_init) {
  uint64_t fd;
  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;

  spin_lock (&usbdrv->xhci->device->lock, &fd);

  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, init->lockflags);

  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, init->lockflags);
  }

  if (ret < 0) {
    free (usbdrv);
    spin_unlock (&usbdrv->xhci->device->lock, fd);
    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, init->lockflags);

  if (ret < 0) {
    free (usbdrv);
    spin_unlock (&usbdrv->xhci->device->lock, fd);
    return false;
  }

  spin_unlock (&usbdrv->xhci->device->lock, fd);

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