#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/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 */

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

DEFINE_DEVICE_OP (usbdrv_write) {}

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

DEFINE_DEVICE_OP (usbdrv_get_size) {}

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;

  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)
    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)
    return false;

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

  return true;
}

DEFINE_DEVICE_FINI (usbdrv_fini) {}