#ifndef _KERNEL_DEVICE_XHCI_H
#define _KERNEL_DEVICE_XHCI_H

#include <aux/compiler.h>
#include <device/def_device_op.h>
#include <device/device.h>
#include <device/usb/usb.h>
#include <libk/list.h>
#include <libk/std.h>
#include <proc/proc.h>
#include <proc/reschedule.h>
#include <proc/suspension_q.h>

/* REF:
 * https://www.intel.com/content/dam/www/public/us/en/documents/technical-specifications/extensible-host-controler-interface-usb-xhci.pdf
 */

/* clang-format off */

/* capability registers */
#define XHCI_CAPLENGTH  0x00
#define XHCI_RSVD       0x01
#define XHCI_HCIVERSION 0x02
#define XHCI_HCSPARAMS1 0x04
#define XHCI_HCSPARAMS2 0x08
#define XHCI_HCSPARAMS3 0x0C
#define XHCI_HCCPARAMS1 0x10
#define XHCI_DBOFF      0x14
#define XHCI_RTSOFF     0x18
#define XHCI_HCCPARAMS2 0x1C

/* operational registers */
#define XHCI_USBCMD     0x00
#define XHCI_USBSTS     0x04
#define XHCI_PAGESIZE   0x08
#define XHCI_DNCTRL     0x14
#define XHCI_CRCR       0x18
#define XHCI_DCBAAP     0x30
#define XHCI_CONFIG     0x38

/* port registers */
#define XHCI_PORTSC     0x00
#define XHCI_PORTPMSC   0x04
#define XHCI_PORTLI     0x08

/* runtime registers */
#define XHCI_MFINDEX    0x00
/* + IRQ sets (0x20) */
#define XHCI_IMAN       0x00
#define XHCI_IMOD       0x04
#define XHCI_ERSTSZ     0x08
#define XHCI_ERSTBA     0x10
#define XHCI_ERDP       0x18

/* HCCPARAMS1 */
#define XHCI_HCCPARAMS1_AC64          0
#define XHCI_HCCPARAMS1_BNC           1
#define XHCI_HCCPARAMS1_CSZ           2
#define XHCI_HCCPARAMS1_PPC           3
#define XHCI_HCCPARAMS1_PIND          4
#define XHCI_HCCPARAMS1_LHRC          5
#define XHCI_HCCPARAMS1_LTC           6
#define XHCI_HCCPARAMS1_NSS           7
#define XHCI_HCCPARAMS1_PAE           8
#define XHCI_HCCPARAMS1_SPC           9
#define XHCI_HCCPARAMS1_SEC           10
#define XHCI_HCCPARAMS1_CFC           11
#define XHCI_HCCPARAMS1_MAX_PSA_SIZE  12
#define XHCI_HCCPARAMS1_XECP          16

/* XHCI extended capabilities */
#define XHCI_XECP_CAP_ID     0
#define XHCI_XECP_NEXT_PTR   8

/* USB legacy support capability */
#define XHCI_USBLEGSUP_CAP_ID     0
#define XHCI_USBLEGSUP_NEXT_PTR   8
#define XHCI_USBLEGSUP_BIOS_SEMA  16
#define XHCI_USBLEGSUP_OS_SEMA    24

/* Supported protocol capability */
#define XHCI_SUPPROTO_DW0_CAP_ID      0
#define XHCI_SUPPROTO_DW0_NEXT_PTR    8
#define XHCI_SUPPROTO_DW0_MINOR_REV   16
#define XHCI_SUPPROTO_DW0_MAJOR_REV   24
#define XHCI_SUPPROTO_DW1_NAME_STRING 0
#define XHCI_SUPPROTO_DW2_PORT_OFF    0
#define XHCI_SUPPROTO_DW2_PORT_COUNT  8
#define XHCI_SUPPROTO_DW2_PROT_DEF    16
#define XHCI_SUPPROTO_DW2_PSIC        28
#define XHCI_SUPPROTO_DW3_SLOT_TYPE   0

/* HCSPARAMS1 */
#define XHCI_HCSPARAMS1_MAX_DEV_SLOTS 0
#define XHCI_HCSPARAMS1_MAX_INTRS     8
#define XHCI_HCSPARAMS1_MAX_PORTS     24

/* HCSPARAMS2 */
#define XHCI_HCSPARAMS2_IST           0
#define XHCI_HCSPARAMS2_ERST_MAX      4
#define XHCI_HCSPARAMS2_MAX_SCRTCH_HI 21
#define XHCI_HCSPARAMS2_SPR           26
#define XHCI_HCSPARAMS2_MAX_SCRTCH_LO 27

/* event types */
#define XHCI_TRB_NORMAL         1
#define XHCI_TRB_SETUP_STAGE    2
#define XHCI_TRB_DATA_STAGE     3
#define XHCI_TRB_STATUS_STAGE   4
#define XHCI_TRB_ISOCH          5
#define XHCI_TRB_LINK           6
#define XHCI_TRB_EVENT_DATA     7
#define XHCI_TRB_NOOP           8
#define XHCI_TRB_SLOT_ENAB_CMD  9
#define XHCI_TRB_SLOT_DISB_CMD  10
#define XHCI_TRB_ADDR_DEV_CMD   11
#define XHCI_TRB_CFG_ENDP_CMD   12
#define XHCI_TRB_EVAL_CTX_CMD   13
#define XHCI_TRB_RESET_ENDP_CMD 14
#define XHCI_TRB_STOP_ENDP_CMD  15
#define XHCI_TRB_SET_DRDQP_CMD  16
#define XHCI_TRB_RESET_DEV_CMD  17
#define XHCI_TRB_FORCE_EVT_CMD  18
#define XHCI_TRB_NEGO_BANDW_CMD 19
#define XHCI_TRB_SET_LTV_CMD    20
#define XHCI_TRB_PORT_BANDW_CMD 21
#define XHCI_TRB_FORCE_HEADER   22
#define XHCI_TRB_NOOP_CMD       23
#define XHCI_TRB_GET_EXTPRP_CMD 24
#define XHCI_TRB_SET_EXTPRP_CMD 25
#define XHCI_TRB_TRANSFER_EVENT 32
#define XHCI_TRB_CMD_CMPL_EVENT 33
#define XHCI_TRB_PORT_STS_CHNG  34
#define XHCI_TRB_BANDW_RQ_EVENT 35
#define XHCI_TRB_DOORBELL_EVENT 36
#define XHCI_TRB_HOST_CTRL_EVNT 37
#define XHCI_TRB_DEV_NOTIF_EVNT 38
#define XHCI_TRB_MFINDEX_WRAP   39

/* generic TRB bits */
#define XHCI_GTRB_TRB_TYPE  10
#define XHCI_GTRB_CYCLE_BIT 0

/* command completion event TRB */
#define XHCI_CCETRB_CTRL_SLOT_ID  24
#define XHCI_CCETRB_CTRL_VFID     16
#define XHCI_CCETRB_CTRL_TRB_TYPE 10
#define XHCI_CCETRB_CTRL_CYCLEBIT 0
#define XHCI_CCETRB_STS_CMPL_CODE 24

/* transfer event TRB */
#define XHCI_TETRB_CTRL_SLOT_ID   24
#define XHCI_TETRB_CTRL_ENDPOINT  16
#define XHCI_TETRB_CTRL_TRB_TYPE  10
#define XHCI_TETRB_CTRL_EVT_DATA  2
#define XHCI_TETRB_CTRL_CYCLEBIT  0
#define XHCI_TETRB_STS_CMPL_CODE  24

/* port status change event TRB */
#define XHCI_PSCETRB_CTRL_TRB_TYPE 10
#define XHCI_PSCETRB_CTRL_CYCLEBIT 0
#define XHCI_PSCETRB_STS_CMPL_CODE 24
#define XHCI_PSCETRB_PARAM_PORT_ID 24

/* link TRB */
#define XHCI_LNKTRB_CTRL_TRB_TYPE  10
#define XHCI_LNKTRB_CTRL_IOC       5
#define XHCI_LNKTRB_CTRL_CH        4
#define XHCI_LNKTRB_CTRL_TC        1
#define XHCI_LNKTRB_CTRL_CYCLEBIT  0
#define XHCI_LNKTRB_STS_INTR_TRGT  22

/* setup stage TRB */
#define XHCI_SSTRB_CTRL_TRT        16
#define XHCI_SSTRB_CTRL_TRB_TYPE   10
#define XHCI_SSTRB_CTRL_IDT        6
#define XHCI_SSTRB_CTRL_IOC        5
#define XHCI_SSTRB_CTRL_CYCLEBIT   0
#define XHCI_SSTRB_STS_INTR_TRGT   22
#define XHCI_SSTRB_PARAM_WLENGTH        48
#define XHCI_SSTRB_PARAM_WINDEX         32
#define XHCI_SSTRB_PARAM_WVALUE         16
#define XHCI_SSTRB_PARAM_BREQUEST       8
#define XHCI_SSTRB_PARAM_BMREQUEST_TYPE 0

/* setup stage TRB transfer types */
#define XHCI_SSTRB_TRT_NO_DATA_STAGE  0
#define XHCI_SSTRB_TRT_OUT_DATA_STAGE 2
#define XHCI_SSTRB_TRT_IN_DATA_STAGE  3

/* data stage TRB */
#define XHCI_DSTRB_CTRL_DIR       16
#define XHCI_DSTRB_CTRL_TRB_TYPE  10
#define XHCI_DSTRB_CTRL_IDT       6
#define XHCI_DSTRB_CTRL_IOC       5
#define XHCI_DSTRB_CTRL_CH        4
#define XHCI_DSTRB_CTRL_NS        3
#define XHCI_DSTRB_CTRL_ISP       2
#define XHCI_DSTRB_CTRL_ENT       1
#define XHCI_DSTRB_CTRL_CYCLEBIT  0
#define XHCI_DSTRB_STS_INTR_TRGT  22
#define XHCI_DSTRB_STS_TD_SIZE    17

/* status stage TRB */
#define XHCI_STSTRB_CTRL_DIR      16
#define XHCI_STSTRB_CTRL_TRB_TYPE 10
#define XHCI_STSTRB_CTRL_IOC      5
#define XHCI_STSTRB_CTRL_CH       4
#define XHCI_STSTRB_CTRL_ENT      1
#define XHCI_STSTRB_CTRL_CYCLEBIT 0
#define XHCI_STSTRB_STS_INTR_TRGT 22

/* extended capabilities */
#define XHCI_EXTCAP_USB_LEGACY_SUPPORT 1
#define XHCI_EXTCAP_SUPPORTED_PROTOCOL 2
#define XHCI_EXTCAP_EXT_POWER_MNGMNT   3
#define XHCI_EXTCAP_IO_VIRT            4
#define XHCI_EXTCAP_MSG_INTR           5
#define XHCI_EXTCAP_LOCAL_MEMORY       6
#define XHCI_EXTCAP_USB_DEBUG          7
#define XHCI_EXTCAP_EXT_MSG_INTR       8
#define XHCI_EXTCAP_VENDOR_DEFINED     9

/* PORTSC bits */
#define XHCI_PORTSC_CCS       0
#define XHCI_PORTSC_PED       1
#define XHCI_PORTSC_OCA       3
#define XHCI_PORTSC_PR        4
#define XHCI_PORTSC_PLS       5
#define XHCI_PORTSC_PP        9
#define XHCI_PORTSC_PORTSPEED 10
#define XHCI_PORTSC_PIC       14
#define XHCI_PORTSC_LWS       16
#define XHCI_PORTSC_CSC       17
#define XHCI_PORTSC_PEC       18
#define XHCI_PORTSC_WRC       19
#define XHCI_PORTSC_OCC       20
#define XHCI_PORTSC_PRC       21
#define XHCI_PORTSC_PLC       22
#define XHCI_PORTSC_CEC       23
#define XHCI_PORTSC_CAS       24
#define XHCI_PORTSC_WCE       25
#define XCHI_PORTSC_WDE       26
#define XHCI_PORTSC_WOE       27
#define XHCI_PORTSC_DR        30
#define XHCI_PORTSC_WPR       31

/* endpoint context */
#define XHCI_EPCTX_EP_STATE     0
#define XHCI_EPCTX_MULT         8
#define XHCI_EPCTX_MAXPSTRAMS   10
#define XHCI_EPCTX_LSA          15
#define XHCI_EPCTX_INTERVAL     16
#define XHCI_EPCTX_MAX_ESIT_HI  24
#define XHCI_EPCTX_ERR_COUNT    1
#define XHCI_EPCTX_EP_TYPE      3
#define XHCI_EPCTX_HID          7
#define XHCI_EPCTX_MAX_BURST_SZ 8
#define XHCI_EPCTX_MAX_PKT_SZ   16
#define XHCI_EPCTX_DCS          0
#define XHCI_EPCTX_TR_DQ_PTR    4
#define XHCI_EPCTX_AVG_TRB_LEN  0
#define XHCI_EPCTX_MAX_ESIT_LOW 16

/* endpoint types */
#define XHCI_EP_INVALID   0
#define XHCI_EP_ISOCH_OUT 1
#define XHCI_EP_BULK_OUT  2
#define XHCI_EP_INTR_OUT  3
#define XHCI_EP_CTRL_BI   4
#define XHCI_EP_ISOCH_IN  5
#define XHCI_EP_BULK_IN   6
#define XHCI_EP_INTR_IN   7

/* slot context */
#define XHIC_SLCTX_ROUTE_STR    0
#define XHCI_SLCTX_SPEED        20
#define XHCI_SLCTX_MTT          25
#define XHCI_SLCTX_HUB          26
#define XHCI_SLCTX_CTX_ENTRIES  27
#define XHCI_SLCTX_MAXEXITLTNCY 0
#define XHCI_SLCTX_ROOTHUBPRNUM 16
#define XHCI_SLCTX_PORT_COUNT   24
#define XHCI_SLCTX_PRNT_HUB_SLT 0
#define XHCI_SLCTX_PARENT_PORT  8
#define XHCI_SLCTX_TT_THINKTIME 16
#define XHCI_SLCTX_INTR_TARGET  22
#define XHCI_SLCTX_USB_DEV_ADDR 0
#define XHCI_SLCTX_SLOT_STATE   27

/* clang-format on */

/* event ring segment table entry */
struct xhci_erst_entry {
  volatile uint64_t ptr;
  volatile uint32_t size;
  volatile uint32_t _rsvd;
} PACKED;

/* transfer request block */
struct xhci_trb {
  volatile uint64_t param;
  volatile uint32_t status;
  volatile uint32_t ctrl;
} PACKED;

struct xhci_ctx32 {
  volatile uint32_t dw[8];
} PACKED;

struct xhci_ctx64 {
  volatile uint32_t dw[16];
} PACKED;

struct xhci_input_ctx32 {
  struct xhci_ctx32 ctrl;
  struct xhci_ctx32 slot;
  struct xhci_ctx32 endpoints[31];
};

struct xhci_input_ctx64 {
  struct xhci_ctx64 ctrl;
  struct xhci_ctx64 slot;
  struct xhci_ctx64 endpoints[31];
};

struct xhci_init {
  uintptr_t xhci_mmio_base;
  bool irqs_support;
  uint8_t irq;
};

#define XHCI_PORT_USB2 0
#define XHCI_PORT_USB3 1

struct xhci_port {
  struct list_node_link ports_link;
  uint8_t port_value;
  int type;
};

struct xhci_ring {
  struct xhci_trb* trbs;
  uintptr_t phys;
  uint32_t idx;
  uint32_t size;
  uint8_t cycle_bit;

  atomic_bool pending;
  uint8_t last_slot_id;
  uint8_t last_cmpl_code;
};

struct xhci_usb_device_config {
  struct usb_config_desc desc;
};

struct xhci_usb_device_if {
  struct usb_if_desc desc;
};

struct xhci_usb_device_endpoint {
  struct usb_endpoint_desc desc;
  struct xhci_ring transfer_ring;
};

#define XHCI_USB_DEVICE_CONFIGS_MAX   16
#define XHCI_USB_DEVICE_IFS_MAX       16
#define XHCI_USB_DEVICE_ENDPOINTS_MAX 16

struct xhci_usb_device {
  struct list_node_link usb_devices_link;
  struct xhci_port* xhci_port;
  int slot_id;
  struct device* device;

  struct xhci_ring endpoint0_ring;

  struct usb_device_desc device_desc;
  struct usb_config_desc config_desc;

  struct xhci_usb_device_config configs[XHCI_USB_DEVICE_CONFIGS_MAX];
  size_t configs_count;

  struct xhci_usb_device_if ifs[XHCI_USB_DEVICE_IFS_MAX];
  size_t ifs_count;

  struct xhci_usb_device_endpoint endpoints[XHCI_USB_DEVICE_ENDPOINTS_MAX];
  size_t endpoints_count;
};

struct xhci_port_status_change {
  struct list_node_link port_changes_link;
  uint8_t port;
  uint32_t portsc;
};

struct xhci {
  struct device* device;

  bool irqs_support;
  uint8_t irq;

  uintptr_t xhci_mmio_base;
  uintptr_t xhci_oper_base;
  uintptr_t xhci_runtime_base;
  uintptr_t xhci_doorbell_base;

  uintptr_t* xhci_dcbaa;
  uintptr_t xhci_dcbaa_phys;

  uint32_t max_scratchpad;
  uintptr_t* scratchpads;
  uintptr_t scratchpads_phys;

  uint32_t max_slots;

  struct xhci_ring cmd_ring;
  struct xhci_ring event_ring;

  struct xhci_erst_entry* erst;
  uintptr_t erst_phys;

  size_t xhci_ctx_size;

  struct list_node_link* xhci_ports;
  struct list_node_link* xhci_usb_devices;

  struct list_node_link* port_changes;
};

int xhci_bulk_transfer (struct xhci* xhci, struct xhci_usb_device* usb_device,
                        uint8_t endpoint_addr, uintptr_t buffer_phys, size_t buffer_size,
                        uint64_t* lockflags);

DEFINE_DEVICE_INIT (xhci_init);

DEFINE_DEVICE_FINI (xhci_fini);

DEFINE_DEVICE_OP (xhci_poll_driver);

#endif // _KERNEL_DEVICE_XHCI_H