#include "src/kernel/usb.h"
#include "src/kernel/log.h"
#include "src/kernel/panic.h"
#include "src/kernel/stream.h"
#include "src/kernel/xhci.h"
#include "src/lib/layout.h"
#include "src/lib/memory.h"
#include "src/lib/util.h"

typedef struct __attribute__((packed)) {
  uint8_t bm_request_type;
  uint8_t b_request;
  uint16_t w_value;
  uint16_t w_index;
  uint16_t w_length;
} usb_setup_t;

#define USB_SETUP_DIRECTION_R(s) BITS_R((s)->bm_request_type, 7, 7)
#define USB_SETUP_DIRECTION_W(s, v) BITS_W((s)->bm_request_type, 7, 7, v)
#define USB_SETUP_TYPE_R(s) BITS_R((s)->bm_request_type, 6, 5)
#define USB_SETUP_TYPE_W(s, v) BITS_W((s)->bm_request_type, 6, 5, v)
#define USB_SETUP_RECIPIENT_R(s) BITS_R((s)->bm_request_type, 4, 0)
#define USB_SETUP_RECIPIENT_W(s, v) BITS_W((s)->bm_request_type, 4, 0, v)

#define USB_SETUP_DIRECTION_H2D 0
#define USB_SETUP_DIRECTION_D2H 1

#define USB_SETUP_TYPE_STANDARD 0
#define USB_SETUP_TYPE_CLASS 1
#define USB_SETUP_TYPE_VENDOR 2

#define USB_SETUP_RECIPIENT_DEVICE 0
#define USB_SETUP_RECIPIENT_INTERFACE 1
#define USB_SETUP_RECIPIENT_ENDPOINT 2
#define USB_SETUP_RECIPIENT_OTHER 3

#define USB_SETUP_GET_DESCRIPTOR 0x06
#define USB_SETUP_SET_CONFIGURATION 0x09
#define USB_SETUP_SET_IDLE 0x0A
#define USB_SETUP_SET_PROTOCOL 0x0B

#define USB_SETUP_DESCRIPTOR_TYPE_R(s) BITS_R((s)->w_value, 15, 8)
#define USB_SETUP_DESCRIPTOR_TYPE_W(s, v) BITS_W((s)->w_value, 15, 8, v)

#define USB_SETUP_DESCRIPTOR_INDEX_R(s) BITS_R((s)->w_value, 7, 0)
#define USB_SETUP_DESCRIPTOR_INDEX_W(s, v) BITS_W((s)->w_value, 7, 0, v)

typedef struct __attribute__((packed)) {
  uint8_t length;
  uint8_t descriptor_type;
} usb_descriptor_base_t;

#define USB_DESCRIPTOR_TYPE_DEVICE 0x01
#define USB_DESCRIPTOR_TYPE_CONFIGURATION 0x02
#define USB_DESCRIPTOR_TYPE_INTERFACE 0x04
#define USB_DESCRIPTOR_TYPE_ENDPOINT 0x05

typedef struct __attribute__((packed)) {
  usb_descriptor_base_t base;
  uint16_t usb_version;
  uint8_t device_class;
  uint8_t device_subclass;
  uint8_t device_protocol;
  uint8_t max_packet_size;
  uint16_t vendor_id;
  uint16_t product_id;
  uint16_t device_version;
  uint8_t manufacturer_str;
  uint8_t product_str;
  uint8_t serial_str;
  uint8_t num_configurations;
} usb_device_descriptor_t;

typedef struct __attribute__((packed)) {
  usb_descriptor_base_t base;
  uint16_t total_length;
  uint8_t num_interfaces;
  uint8_t configuration_value;
  uint8_t configuration_str;
  uint8_t attributes;
  uint8_t max_power;
} usb_configuration_descriptor_t;

typedef struct __attribute__((packed)) {
  usb_descriptor_base_t base;
  uint8_t interface_number;
  uint8_t alternate_setting;
  uint8_t num_endpoints;
  uint8_t interface_class;
  uint8_t interface_subclass;
  uint8_t interface_protocol;
  uint8_t interface_str;
} usb_interface_descriptor_t;

typedef struct __attribute__((packed)) {
  usb_descriptor_base_t base;
  uint8_t endpoint_address;
  uint8_t attributes;
  uint16_t max_packet_size;
  uint8_t interval;
} usb_endpoint_descriptor_t;

typedef struct __attribute__((packed)) {
  usb_descriptor_base_t base;
  uint16_t hid_version;
  uint8_t country_code;
  uint8_t num_descriptors;
  uint8_t report_descriptor_type;
  uint16_t report_descriptor_length;
} usb_hid_descriptor_t;

static uint8_t descriptor_buffer[PAGE_SIZE] __attribute__((aligned(64)));

static void scan_device(xhci_slot_t slot, usb_device_descriptor_t *device, usb_configuration_descriptor_t *configuration,
                        usb_interface_descriptor_t *interface, usb_endpoint_descriptor_t *endpoint) {
  memory_set(0, sizeof(usb_device_descriptor_t), device);
  memory_set(0, sizeof(usb_configuration_descriptor_t), configuration);
  memory_set(0, sizeof(usb_interface_descriptor_t), interface);
  memory_set(0, sizeof(usb_endpoint_descriptor_t), endpoint);

  usb_setup_t setup;
  memory_set(0, sizeof(setup), &setup);
  USB_SETUP_DIRECTION_W(&setup, USB_SETUP_DIRECTION_D2H);
  setup.b_request = USB_SETUP_GET_DESCRIPTOR;
  USB_SETUP_DESCRIPTOR_TYPE_W(&setup, USB_DESCRIPTOR_TYPE_DEVICE);
  setup.w_length = sizeof(usb_device_descriptor_t);
  memory_set(0, sizeof(descriptor_buffer), descriptor_buffer);
  xhci_control_transfer(slot, *(uint64_t *)&setup, &descriptor_buffer, setup.w_length);

  memory_copy(descriptor_buffer, sizeof(usb_device_descriptor_t), device);

  memory_set(0, sizeof(setup), &setup);
  USB_SETUP_DIRECTION_W(&setup, USB_SETUP_DIRECTION_D2H);
  setup.b_request = USB_SETUP_GET_DESCRIPTOR;
  USB_SETUP_DESCRIPTOR_TYPE_W(&setup, USB_DESCRIPTOR_TYPE_CONFIGURATION);
  setup.w_length = sizeof(usb_configuration_descriptor_t);
  memory_set(0, sizeof(descriptor_buffer), descriptor_buffer);
  xhci_control_transfer(slot, *(uint64_t *)&setup, &descriptor_buffer, setup.w_length);

  memory_set(0, sizeof(setup), &setup);
  USB_SETUP_DIRECTION_W(&setup, USB_SETUP_DIRECTION_D2H);
  setup.b_request = USB_SETUP_GET_DESCRIPTOR;
  USB_SETUP_DESCRIPTOR_TYPE_W(&setup, USB_DESCRIPTOR_TYPE_CONFIGURATION);
  setup.w_length = ((usb_configuration_descriptor_t *)&descriptor_buffer)->total_length;
  memory_set(0, sizeof(descriptor_buffer), descriptor_buffer);
  xhci_control_transfer(slot, *(uint64_t *)&setup, &descriptor_buffer, setup.w_length);

  usb_descriptor_base_t *descriptor = (usb_descriptor_base_t *)&descriptor_buffer;
  while (descriptor->descriptor_type) { // Assuming `descriptor_buffer` longer than any potential descriptor set.
    if (descriptor->descriptor_type == USB_DESCRIPTOR_TYPE_CONFIGURATION && !configuration->base.descriptor_type) {
      memory_copy(descriptor, sizeof(usb_configuration_descriptor_t), configuration);
    }
    if (descriptor->descriptor_type == USB_DESCRIPTOR_TYPE_INTERFACE && !interface->base.descriptor_type) {
      memory_copy(descriptor, sizeof(usb_interface_descriptor_t), interface);
    }
    if (descriptor->descriptor_type == USB_DESCRIPTOR_TYPE_ENDPOINT && !endpoint->base.descriptor_type) {
      memory_copy(descriptor, sizeof(usb_endpoint_descriptor_t), endpoint);
    }
    descriptor = (usb_descriptor_base_t *)((uint8_t *)descriptor + descriptor->length);
  }
}

void usb_init() {
  PRINT_LN(kernel_log, "Enumerating devices...");
  usb_enumerate(kernel_log);
  PRINT_LN(kernel_log, "Done.");
}

void usb_enumerate(stream_t *out) {
  for (xhci_port_t port = 0; port < xhci_port_count(); port++) {
    if (!xhci_port_connected(port)) {
      continue;
    }

    xhci_slot_t slot = xhci_attach(port);
    if (slot == (xhci_slot_t)-1) {
      continue;
    }

    usb_device_descriptor_t device;
    usb_configuration_descriptor_t configuration;
    usb_interface_descriptor_t interface;
    usb_endpoint_descriptor_t endpoint;
    scan_device(slot, &device, &configuration, &interface, &endpoint);

    xhci_detach(slot);

    PRINT_LN(
        out,
        "%hx:%hx usb_version=%hx device_class=%hhx device_subclass=%hhx device_protocol=%hhx interface_class=%hhx interface_subclass=%hhx "
        "interface_protocol=%hhx endpoint_address=%hhx max_packet_size=%hx interval=%hhx",
        device.vendor_id, device.product_id, device.usb_version, device.device_class, device.device_subclass, device.device_protocol,
        interface.interface_class, interface.interface_subclass, interface.interface_protocol, endpoint.endpoint_address,
        endpoint.max_packet_size, endpoint.interval);
  }
}

xhci_port_t usb_find_by_interface(uint8_t class, uint8_t subclass, uint8_t protocol) {
  for (xhci_port_t port = 0; port < xhci_port_count(); port++) {
    if (!xhci_port_connected(port)) {
      continue;
    }

    xhci_slot_t slot = xhci_attach(port);
    if (slot == (xhci_slot_t)-1) {
      continue;
    }

    usb_device_descriptor_t device;
    usb_configuration_descriptor_t configuration;
    usb_interface_descriptor_t interface;
    usb_endpoint_descriptor_t endpoint;
    scan_device(slot, &device, &configuration, &interface, &endpoint);

    xhci_detach(slot);

    if ((class == NUL || interface.interface_class == class) && (subclass == NUL || interface.interface_subclass == subclass) &&
        (protocol == NUL || interface.interface_protocol == protocol)) {
      return port;
    }
  }
  return (xhci_port_t)-1;
}

xhci_slot_t usb_attach(xhci_port_t port) {
  return xhci_attach(port);
}

xhci_endpoint_t usb_open_endpoint(xhci_slot_t slot, uint8_t type, uint16_t max_packet_size) {
  usb_device_descriptor_t device;
  usb_configuration_descriptor_t configuration;
  usb_interface_descriptor_t interface;
  usb_endpoint_descriptor_t endpoint;
  scan_device(slot, &device, &configuration, &interface, &endpoint);

  usb_setup_t setup;
  memory_set(0, sizeof(setup), &setup);
  USB_SETUP_DIRECTION_W(&setup, USB_SETUP_DIRECTION_H2D);
  USB_SETUP_TYPE_W(&setup, USB_SETUP_TYPE_STANDARD);
  USB_SETUP_RECIPIENT_W(&setup, USB_SETUP_RECIPIENT_DEVICE);
  setup.b_request = USB_SETUP_SET_CONFIGURATION;
  setup.w_value = configuration.configuration_value;
  xhci_control_transfer(slot, *(uint64_t *)&setup, NUL, 0);

  memory_set(0, sizeof(setup), &setup);
  USB_SETUP_DIRECTION_W(&setup, USB_SETUP_DIRECTION_H2D);
  USB_SETUP_TYPE_W(&setup, USB_SETUP_TYPE_CLASS);
  USB_SETUP_RECIPIENT_W(&setup, USB_SETUP_RECIPIENT_INTERFACE);
  setup.b_request = USB_SETUP_SET_PROTOCOL;
  xhci_control_transfer(slot, *(uint64_t *)&setup, NUL, 0);

  memory_set(0, sizeof(setup), &setup);
  USB_SETUP_DIRECTION_W(&setup, USB_SETUP_DIRECTION_H2D);
  USB_SETUP_TYPE_W(&setup, USB_SETUP_TYPE_CLASS);
  USB_SETUP_RECIPIENT_W(&setup, USB_SETUP_RECIPIENT_INTERFACE);
  setup.b_request = USB_SETUP_SET_IDLE;
  xhci_control_transfer(slot, *(uint64_t *)&setup, NUL, 0);

  ASSERT(endpoint.max_packet_size == max_packet_size, "usb_open_endpoint: Unexpected max packet size.");

  return xhci_open_endpoint(slot, endpoint.endpoint_address, type, endpoint.max_packet_size, 7);
}

uint16_t usb_read_endpoint(xhci_slot_t slot, xhci_endpoint_t endpoint, uint16_t max, void *to) {
  return xhci_read_endpoint(slot, endpoint, max, to);
}

void usb_close_endpoint(xhci_slot_t slot, xhci_endpoint_t endpoint) {
  return xhci_close_endpoint(slot, endpoint);
}

void usb_detach(xhci_slot_t slot) {
  return xhci_detach(slot);
}