// SPDX-License-Identifier: GPL-2.0-or-later
//
// GMMOS's USB keyboards on xHCI: kernel/drivers/bus/xhci/xhci.c with
// everything but finding a keyboard and reading it taken out.
//
// It does not take over the live driver's state. The kernel may have died
// inside that driver, mid-command, with its rings half-written; so each
// controller is reset and enumerated from nothing, and of the live driver only
// two things are used: where the registers are mapped, and the scratchpad
// pages the controller was already lent (DCBAA entry 0 survives in memory).
//
// What is left, in order:
//
//   - halt, reset, a command ring, an event ring, run
//   - every connected root port: reset, Enable Slot, Address Device, the
//     device descriptor, the configuration
//   - a USB 2 hub, up to two deep: power its ports, reset each connected one,
//     and enumerate what is behind it
//   - a boot-protocol keyboard interface: Set Configuration, Set Protocol
//     (boot), Set Idle, Configure Endpoint for its interrupt IN endpoint, and
//     a Normal TRB queued on it
//
// and then the event ring is polled and every completed report re-queued. No
// interrupts, no hot-plug, no mice, no storage, no USB 3 hubs (a keyboard is
// never behind one's SuperSpeed half). Memory is a static pool: the DMA
// allocator belongs to the dead kernel.

#include "core/types.h"
#include "core/defs.h"
#include "arch/memlayout.h"
#include "dev/usb.h"
#include "gmmos/gmmos.h"
#include "../drivers/bus/xhci/xhci.h"

#define NHC     4
#define NKBD    4
#define NTRB    64                      // per ring, link TRB included
#define MAXSLOT 32

// --- memory ----------------------------------------------------------------

static uchar pool[256 * 1024] __attribute__((aligned(4096)));
static uint64 used;

static void *
palloc(uint64 size, uint64 align)
{
  used = (used + align - 1) & ~(align - 1);
  if (used + size > sizeof(pool))
    return 0;
  void *p = pool + used;
  used += size;
  memset(p, 0, size);
  return p;
}

static addr_t
phys(volatile void *p)
{
  return V2P((addr_t)p);
}

// --- one controller --------------------------------------------------------

struct ring {
  volatile struct trb *trb;
  int enq, cycle;
};

struct hc {
  volatile uchar *cap, *op, *rt;
  volatile uint32 *db;
  int ctxsize, maxslots, maxports;
  uint64 *dcbaa;
  struct ring cmd;
  volatile struct trb *ev;
  int evdeq, evcycle;
  struct ring ep0[MAXSLOT + 1];
};

struct kbd {
  struct hc *hc;
  int slot, dci;
  struct ring ring;
  uchar *buf;
  int len;
};

static struct hc hcs[NHC];
static struct kbd kbds[NKBD];
static int nkbds;
static struct gmmos_hidkeys keys;

static inline uint32 rd32(volatile uchar *p, int o)  { return *(volatile uint32 *)(p + o); }
static inline void   wr32(volatile uchar *p, int o, uint32 v) { *(volatile uint32 *)(p + o) = v; }
static inline void   wr64(volatile uchar *p, int o, uint64 v)
{
  *(volatile uint32 *)(p + o) = (uint32)v;
  *(volatile uint32 *)(p + o + 4) = (uint32)(v >> 32);
}

static int
ring_init(struct ring *r)
{
  r->trb = palloc(NTRB * 16, 4096);
  if (r->trb == 0)
    return -1;
  r->enq = 0;
  r->cycle = 1;
  r->trb[NTRB - 1].param = phys(r->trb);
  r->trb[NTRB - 1].control = TRB_TYPE(TRB_LINK) | TRB_TC;
  return 0;
}

static void
push(struct ring *r, uint64 param, uint32 status, uint32 control)
{
  volatile struct trb *t = &r->trb[r->enq];
  t->param = param;
  t->status = status;
  __sync_synchronize();
  t->control = (control & ~TRB_C) | (r->cycle ? TRB_C : 0);
  __sync_synchronize();
  if (++r->enq == NTRB - 1) {
    r->trb[NTRB - 1].control = TRB_TYPE(TRB_LINK) | TRB_TC | (r->cycle ? TRB_C : 0);
    r->enq = 0;
    r->cycle ^= 1;
  }
}

static int
pop(struct hc *h, struct trb *out)
{
  volatile struct trb *e = &h->ev[h->evdeq];
  uint32 c = e->control;
  if (((c & TRB_C) ? 1 : 0) != h->evcycle)
    return 0;
  __sync_synchronize();
  out->param = e->param;
  out->status = e->status;
  out->control = c;
  if (++h->evdeq == NTRB) {
    h->evdeq = 0;
    h->evcycle ^= 1;
  }
  wr64(h->rt, XHCI_IR0 + IR_ERDP, (phys(h->ev) + h->evdeq * 16) | ERDP_EHB);
  return 1;
}

static void report(struct hc *h, struct trb *ev);

// The next event of this type, within ms -- for a transfer event, the one
// for EP0 of this slot. Keyboard reports that arrive meanwhile, a key pressed
// while a second keyboard is enumerated, are handled rather than lost.
static int
wait_event(struct hc *h, int type, int slot, struct trb *out, int ms)
{
  for (int i = 0; i < ms * 100; i++) {
    while (pop(h, out)) {
      int t = TRB_GET_TYPE(out->control);
      if (t == TRB_XFER_EVENT && (type != t || TRB_GET_SLOT(out->control) != slot ||
                                  TRB_GET_EPID(out->control) != DCI_EP0)) {
        report(h, out);
        continue;
      }
      if (t == type)
        return 0;
    }
    microdelay(10);
  }
  return -1;
}

static int
command(struct hc *h, uint64 param, uint32 control, struct trb *ev)
{
  struct trb e;
  push(&h->cmd, param, 0, control);
  h->db[0] = 0;
  if (wait_event(h, TRB_CMD_COMPLETE, 0, &e, 500) < 0)
    return -1;
  if (ev)
    *ev = e;
  return TRB_CC(e.status);
}

static uint32 *
ctx(struct hc *h, void *base, int i)
{
  return (uint32 *)((uchar *)base + i * h->ctxsize);
}

// --- devices ---------------------------------------------------------------

struct dev {
  int slot, port, speed, route, tier;
  int tt_slot, tt_port;
  void *devctx;
};

static uchar dbuf[512] __attribute__((aligned(64)));

static int
control(struct hc *h, int slot, int type, int req, int value, int index,
        void *buf, int len)
{
  struct usb_setup s = { type, req, value, index, len };
  struct ring *r = &h->ep0[slot];
  int in = (type & 0x80) != 0;
  struct trb ev;

  push(r, *(uint64 *)&s, 8, TRB_TYPE(TRB_SETUP) | TRB_IDT |
       ((uint32)(len ? (in ? TRT_IN : TRT_OUT) : TRT_NONE) << 16));
  if (len)
    push(r, phys(buf), len, TRB_TYPE(TRB_DATA) | (in ? DIR_IN : 0));
  push(r, 0, 0, TRB_TYPE(TRB_STATUS) | TRB_IOC | ((len == 0 || !in) ? DIR_IN : 0));
  h->db[slot] = DCI_EP0;
  if (wait_event(h, TRB_XFER_EVENT, slot, &ev, 500) < 0)
    return -1;
  int cc = TRB_CC(ev.status);
  return cc == CC_SUCCESS || cc == CC_SHORT_PACKET ? 0 : -1;
}

static int
mps0(int speed)
{
  return speed == XSPEED_HIGH ? 64 : speed == XSPEED_SUPER ? 512 : 8;
}

// Enable Slot and Address Device: the device answers on its own address with
// EP0 working at the default packet size.
static int
address(struct hc *h, struct dev *d)
{
  struct trb ev;
  if (command(h, 0, TRB_TYPE(TRB_ENABLE_SLOT), &ev) != CC_SUCCESS)
    return -1;
  d->slot = TRB_GET_SLOT(ev.control);
  if (d->slot <= 0 || d->slot > MAXSLOT || d->slot > h->maxslots)
    return -1;

  d->devctx = palloc(h->ctxsize * 32, 64);
  void *in = palloc(h->ctxsize * 33, 64);
  if (!d->devctx || !in || ring_init(&h->ep0[d->slot]) < 0)
    return -1;
  h->dcbaa[d->slot] = phys(d->devctx);

  uint32 *c = ctx(h, in, 0), *sc = ctx(h, in, 1), *ep0 = ctx(h, in, 2);
  c[1] = 0x3;
  sc[0] = (d->route & 0xFFFFF) | ((uint32)d->speed << 20) | (1u << 27);
  sc[1] = (uint32)d->port << 16;
  sc[2] = d->tt_slot ? ((uint32)d->tt_slot | ((uint32)d->tt_port << 8)) : 0;
  ep0[1] = (EPTYPE_CONTROL << 3) | (3 << 1) | ((uint32)mps0(d->speed) << 16);
  *(uint64 *)&ep0[2] = phys(h->ep0[d->slot].trb) | 1;
  ep0[4] = 8;
  if (command(h, phys(in), TRB_TYPE(TRB_ADDRESS_DEV) | TRB_SLOT(d->slot), 0) != CC_SUCCESS)
    return -1;

  // A full-speed device's real EP0 packet size is in byte 7 of its device
  // descriptor, and the first eight bytes can be read at any size.
  if (control(h, d->slot, 0x80, 6, 0x0100, 0, dbuf, 8) < 0)
    return -1;
  if (d->speed == XSPEED_FULL && dbuf[7] && dbuf[7] != 8) {
    memset(in, 0, h->ctxsize * 3);
    c[1] = 0x2;
    ep0[1] = (EPTYPE_CONTROL << 3) | (3 << 1) | ((uint32)dbuf[7] << 16);
    *(uint64 *)&ep0[2] = phys(h->ep0[d->slot].trb) | h->ep0[d->slot].cycle;
    ep0[4] = 8;
    command(h, phys(in), TRB_TYPE(TRB_EVAL_CTX) | TRB_SLOT(d->slot), 0);
  }
  return 0;
}

// Configure Endpoint with only the slot context: how a slot is told it is a
// hub, and how many ports it has. Before anything behind it is addressed.
static int
make_hub(struct hc *h, struct dev *d, int nports, int mtt)
{
  void *in = palloc(h->ctxsize * 3, 64);
  if (!in)
    return -1;
  uint32 *c = ctx(h, in, 0), *sc = ctx(h, in, 1);
  c[1] = 1;
  sc[0] = (d->route & 0xFFFFF) | ((uint32)d->speed << 20) | (1u << 27) |
          (1u << 26) | (mtt ? (1u << 25) : 0);
  sc[1] = ((uint32)d->port << 16) | ((uint32)nports << 24);
  sc[2] = d->tt_slot ? ((uint32)d->tt_slot | ((uint32)d->tt_port << 8)) : 0;
  return command(h, phys(in), TRB_TYPE(TRB_CONFIG_EP) | TRB_SLOT(d->slot), 0)
         == CC_SUCCESS ? 0 : -1;
}

static int
interval(int speed, int bInterval)
{
  int i;
  if (speed == XSPEED_HIGH || speed == XSPEED_SUPER)
    i = bInterval - 1;
  else {
    int f = bInterval ? bInterval : 1;
    for (i = 0; (1 << (i + 1)) <= f; i++)
      ;
    i += 3;
  }
  return i < 0 ? 0 : i > 15 ? 15 : i;
}

static void
arm(struct kbd *k)
{
  push(&k->ring, phys(k->buf), k->len, TRB_TYPE(TRB_NORMAL) | TRB_IOC | TRB_ISP);
  k->hc->db[k->slot] = k->dci;
}

static int
keyboard(struct hc *h, struct dev *d, int config, int iface,
         const struct usb_endpoint_desc *ep)
{
  if (nkbds == NKBD)
    return -1;
  struct kbd *k = &kbds[nkbds];

  if (control(h, d->slot, 0x00, 9, config, 0, 0, 0) < 0)          // SET_CONFIGURATION
    return -1;
  control(h, d->slot, 0x21, 0x0B, 0, iface, 0, 0);                // SET_PROTOCOL boot
  control(h, d->slot, 0x21, 0x0A, 0, iface, 0, 0);                // SET_IDLE: changes only

  int mps = ep->wMaxPacketSize & 0x7FF;
  k->hc = h;
  k->slot = d->slot;
  k->dci = (ep->bEndpointAddress & 0xF) * 2 + 1;
  k->len = mps < 64 ? mps : 64;
  k->buf = palloc(64, 64);
  void *in = palloc(h->ctxsize * 33, 64);
  if (!k->buf || !in || ring_init(&k->ring) < 0)
    return -1;

  uint32 *c = ctx(h, in, 0), *sc = ctx(h, in, 1), *ec = ctx(h, in, k->dci + 1);
  uint32 *cur = ctx(h, d->devctx, 0);
  c[1] = 1 | (1u << k->dci);
  sc[0] = (cur[0] & ~(0x1Fu << 27)) | ((uint32)k->dci << 27);
  sc[1] = cur[1];
  sc[2] = cur[2];
  ec[0] = (uint32)interval(d->speed, ep->bInterval) << 16;
  ec[1] = (EPTYPE_INT_IN << 3) | (3 << 1) | ((uint32)mps << 16);
  *(uint64 *)&ec[2] = phys(k->ring.trb) | 1;
  ec[4] = (uint32)mps | ((uint32)mps << 16);
  if (command(h, phys(in), TRB_TYPE(TRB_CONFIG_EP) | TRB_SLOT(d->slot), 0) != CC_SUCCESS)
    return -1;

  nkbds++;
  arm(k);
  return 0;
}

static void enumerate(struct hc *h, struct dev *d);

// A USB 2 hub: power every port, then for each that has something on it,
// reset it and enumerate what comes up.
static void
hub(struct hc *h, struct dev *d, int config, int mtt)
{
  if (d->tier >= 2)
    return;
  if (control(h, d->slot, 0x00, 9, config, 0, 0, 0) < 0)
    return;
  if (control(h, d->slot, 0xA0, 6, 0x2900, 0, dbuf, 9) < 0)
    return;
  int nports = dbuf[2], pgood = dbuf[5] * 2;
  if (nports > 15)
    nports = 15;
  if (make_hub(h, d, nports, mtt) < 0)
    return;

  for (int p = 1; p <= nports; p++)
    control(h, d->slot, 0x23, 3, 8, p, 0, 0);                     // PORT_POWER
  microdelay((pgood < 100 ? 100 : pgood) * 1000);

  for (int p = 1; p <= nports; p++) {
    uchar st[4];
    if (control(h, d->slot, 0xA3, 0, 0, p, st, 4) < 0 || !(st[0] & 1))
      continue;
    control(h, d->slot, 0x23, 3, 4, p, 0, 0);                     // PORT_RESET
    int ok = 0;
    for (int i = 0; i < 50 && !ok; i++) {
      microdelay(10000);
      if (control(h, d->slot, 0xA3, 0, 0, p, st, 4) == 0 && (st[2] & 0x10))
        ok = 1;
    }
    control(h, d->slot, 0x23, 1, 20, p, 0, 0);                    // C_PORT_RESET
    control(h, d->slot, 0x23, 1, 16, p, 0, 0);                    // C_PORT_CONNECTION
    if (!ok || !(st[0] & 2))
      continue;
    microdelay(10000);

    struct dev c;
    memset(&c, 0, sizeof(c));
    c.port = d->port;
    c.tier = d->tier + 1;
    c.route = d->route | (p << (4 * d->tier));
    c.speed = (st[1] & 0x02) ? XSPEED_LOW : (st[1] & 0x04) ? XSPEED_HIGH : XSPEED_FULL;
    // A slow device behind a high-speed hub goes through that hub's
    // transaction translator; behind a full-speed hub, through whichever one
    // the hub itself is behind.
    if (c.speed != XSPEED_HIGH && d->speed == XSPEED_HIGH) {
      c.tt_slot = d->slot;
      c.tt_port = p;
    } else {
      c.tt_slot = d->tt_slot;
      c.tt_port = d->tt_port;
    }
    enumerate(h, &c);
  }
}

// Address a device, and keep it only if it is a hub or a keyboard.
static void
enumerate(struct hc *h, struct dev *d)
{
  if (address(h, d) < 0)
    return;
  int devclass = dbuf[4], devproto = dbuf[6];

  if (control(h, d->slot, 0x80, 6, 0x0200, 0, dbuf, 9) < 0)
    return;
  int total = dbuf[2] | (dbuf[3] << 8);
  if (total > (int)sizeof(dbuf))
    total = sizeof(dbuf);
  if (control(h, d->slot, 0x80, 6, 0x0200, 0, dbuf, total) < 0)
    return;
  int config = dbuf[5];

  if (devclass == 9) {
    if (d->speed != XSPEED_SUPER)
      hub(h, d, config, devproto == 2);
    return;
  }

  // The first boot keyboard interface, and its interrupt IN endpoint.
  int iface = -1;
  for (int o = 0; o + 2 <= total && dbuf[o] >= 2; o += dbuf[o]) {
    if (dbuf[o + 1] == 4 && o + 9 <= total) {
      const struct usb_interface_desc *id = (const void *)&dbuf[o];
      iface = (id->bInterfaceClass == 3 && id->bInterfaceSubClass == 1 &&
               id->bInterfaceProtocol == 1) ? id->bInterfaceNumber : -1;
    } else if (dbuf[o + 1] == 5 && o + 7 <= total && iface >= 0) {
      const struct usb_endpoint_desc *ep = (const void *)&dbuf[o];
      if ((ep->bEndpointAddress & 0x80) && USB_EP_XFERTYPE(ep->bmAttributes) == 3) {
        struct usb_endpoint_desc copy = *ep;
        keyboard(h, d, config, iface, &copy);
        return;
      }
    }
  }
}

static int
start(struct hc *h, volatile uchar *cap)
{
  h->cap = cap;
  h->op = cap + (rd32(cap, XHCI_CAPLENGTH) & 0xFF);
  h->rt = cap + (rd32(cap, XHCI_RTSOFF) & ~0x1Fu);
  h->db = (volatile uint32 *)(cap + (rd32(cap, XHCI_DBOFF) & ~0x3u));
  uint32 hcs1 = rd32(cap, XHCI_HCSPARAMS1), hcs2 = rd32(cap, XHCI_HCSPARAMS2);
  h->maxslots = HCS1_MAXSLOTS(hcs1) < MAXSLOT ? HCS1_MAXSLOTS(hcs1) : MAXSLOT;
  h->maxports = HCS1_MAXPORTS(hcs1);
  h->ctxsize = (rd32(cap, XHCI_HCCPARAMS1) & HCC1_CSZ) ? 64 : 32;

  // The scratchpad array the live driver lent the controller, from the DCBAA
  // it gave it. Read before the reset, which forgets where that was.
  uint64 scratch = 0;
  uint64 old = rd32(h->op, XHCI_DCBAAP) | ((uint64)rd32(h->op, XHCI_DCBAAP + 4) << 32);
  if (old)
    scratch = *(uint64 *)P2V(old & ~0x3Full);
  int spb = (HCS2_SPB_HI(hcs2) << 5) | HCS2_SPB_LO(hcs2);
  if (spb && !scratch)
    return -1;

  wr32(h->op, XHCI_USBCMD, rd32(h->op, XHCI_USBCMD) & ~(CMD_RS | CMD_INTE));
  for (int i = 0; i < 100 && !(rd32(h->op, XHCI_USBSTS) & STS_HCH); i++)
    microdelay(1000);
  wr32(h->op, XHCI_USBCMD, CMD_HCRST);
  int i;
  for (i = 0; i < 1000; i++) {
    microdelay(1000);
    if (!(rd32(h->op, XHCI_USBCMD) & CMD_HCRST) && !(rd32(h->op, XHCI_USBSTS) & STS_CNR))
      break;
  }
  if (i == 1000)
    return -1;

  h->dcbaa = palloc((MAXSLOT + 1) * 8, 4096);
  uint64 *erst = palloc(64, 64);
  h->ev = palloc(NTRB * 16, 4096);
  if (!h->dcbaa || !erst || !h->ev || ring_init(&h->cmd) < 0)
    return -1;
  h->dcbaa[0] = scratch;
  h->evdeq = 0;
  h->evcycle = 1;
  erst[0] = phys(h->ev);
  erst[1] = NTRB;

  wr32(h->op, XHCI_CONFIG, h->maxslots);
  wr64(h->op, XHCI_DCBAAP, phys(h->dcbaa));
  wr64(h->op, XHCI_CRCR, phys(h->cmd.trb) | CRCR_RCS);
  wr32(h->rt, XHCI_IR0 + IR_IMAN, IMAN_IP);           // pending cleared, IE off
  wr32(h->rt, XHCI_IR0 + IR_ERSTSZ, 1);
  wr64(h->rt, XHCI_IR0 + IR_ERDP, phys(h->ev));
  wr64(h->rt, XHCI_IR0 + IR_ERSTBA, phys(erst));
  wr32(h->op, XHCI_USBCMD, CMD_RS);
  for (i = 0; i < 100 && (rd32(h->op, XHCI_USBSTS) & STS_HCH); i++)
    microdelay(1000);
  if (rd32(h->op, XHCI_USBSTS) & STS_HCH)
    return -1;

  // USB 3 ports train on their own after a reset, USB 2 ports need one; a
  // device that has not reconnected yet gets a moment.
  microdelay(300000);
  for (int p = 1; p <= h->maxports; p++) {
    uint32 sc = rd32(h->op, XHCI_PORTSC(p));
    if (!(sc & PORT_CCS))
      continue;
    if (!(sc & PORT_PED)) {
      wr32(h->op, XHCI_PORTSC(p), (sc & ~PORT_RW1C) | PORT_PR);
      for (i = 0; i < 200; i++) {
        microdelay(1000);
        sc = rd32(h->op, XHCI_PORTSC(p));
        if (sc & PORT_PRC)
          break;
      }
      wr32(h->op, XHCI_PORTSC(p), (sc & ~PORT_RW1C) | PORT_PRC | PORT_CSC);
      sc = rd32(h->op, XHCI_PORTSC(p));
    }
    if (!(sc & PORT_PED))
      continue;
    struct dev d;
    memset(&d, 0, sizeof(d));
    d.port = p;
    d.speed = PORT_SPEED(sc);
    enumerate(h, &d);
  }
  return 0;
}

static void
report(struct hc *h, struct trb *ev)
{
  int slot = TRB_GET_SLOT(ev->control), dci = TRB_GET_EPID(ev->control);
  for (int i = 0; i < nkbds; i++) {
    struct kbd *k = &kbds[i];
    if (k->hc != h || k->slot != slot || k->dci != dci)
      continue;
    int cc = TRB_CC(ev->status);
    if ((cc == CC_SUCCESS || cc == CC_SHORT_PACKET) &&
        k->len - (int)(ev->status & 0xFFFFFF) >= 8)
      gmmos_hid_report(&keys, k->buf);
    arm(k);
    return;
  }
}

static int
probe(void)
{
  volatile uchar *caps[NHC];
  int n = xhci_gmmos_controllers(caps, NHC);

  for (int i = 0; i < n; i++)
    start(&hcs[i], caps[i]);
  return nkbds > 0;
}

static int
poll(void)
{
  struct trb ev;
  int k = gmmos_hid_next(&keys);
  if (k)
    return k;
  for (int i = 0; i < NHC; i++) {
    if (!hcs[i].ev)
      continue;
    while (pop(&hcs[i], &ev))
      if (TRB_GET_TYPE(ev.control) == TRB_XFER_EVENT)
        report(&hcs[i], &ev);
  }
  return gmmos_hid_next(&keys);
}

const struct gmmos_kbd gmmos_xhci = {
  .name  = "USB (xHCI)",
  .probe = probe,
  .poll  = poll,
};