// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

#include <assert.h>
#include <stdint.h>
#include <string.h>

#include "usb_utils.h"
#include "usbdpi.h"

// Seed numbers for the LFSR generators in each transfer direction for
// the given stream number
#define USBTST_LFSR_SEED(s) (uint8_t)(0x10U + (s)*7U)
#define USBDPI_LFSR_SEED(s) (uint8_t)(0x9BU - (s)*7U)
// Seed number of packet retrying
#define RETRY_LFSR_SEED(s) (uint8_t)(0x24U + (s)*7U)

// Simple LFSR for 8-bit sequences
#define LFSR_ADVANCE(lfsr) \
  (((lfsr) << 1) ^         \
   ((((lfsr) >> 1) ^ ((lfsr) >> 2) ^ ((lfsr) >> 3) ^ ((lfsr) >> 7)) & 1u))

// Stream signature words
#define STREAM_SIGNATURE_HEAD 0x579EA01AU
#define STREAM_SIGNATURE_TAIL 0x160AE975U

// Verbose logging/diagnostic reporting
// TODO - consider folding this into the existing log level
static const bool verbose = false;

static const bool expect_sig = true;

// Determine the next stream for which IN data packets shall be requested
static inline unsigned in_stream_next(usbdpi_ctx_t *ctx);

// Determine the next stream for which OUT data shall be sent
static inline unsigned out_stream_next(usbdpi_ctx_t *ctx);

// Check a data packet received from the test software (usbdev_stream_test)
static bool stream_data_check(usbdpi_ctx_t *ctx, usbdpi_stream_t *s,
                              const usbdpi_transfer_t *rx, bool accept);

// Generate a data packet as if it had been received from the device
static usbdpi_transfer_t *stream_data_gen(usbdpi_ctx_t *ctx, usbdpi_stream_t *s,
                                          unsigned len);

// Check a data packet received from the test software (usbdev_stream_test)
// and collect the data, combined with our LFSR-generated random stream,
// for later transmission back to the device
static usbdpi_transfer_t *stream_data_process(usbdpi_ctx_t *ctx,
                                              usbdpi_stream_t *s,
                                              usbdpi_transfer_t *rx);

// Check the stream signature
static bool stream_sig_check(usbdpi_ctx_t *ctx, usbdpi_stream_t *s,
                             usbdpi_transfer_t *rx);

// Determine the next stream for which IN data packets shall be requested
inline unsigned in_stream_next(usbdpi_ctx_t *ctx) {
  uint8_t id = ctx->stream_in;
  if (++id >= ctx->nstreams) {
    id = 0U;
  }
  ctx->stream_in = id;
  return id;
}

// Determine the next stream for which OUT data shall be sent
inline unsigned out_stream_next(usbdpi_ctx_t *ctx) {
  uint8_t id = ctx->stream_out;
  if (++id >= ctx->nstreams) {
    id = 0U;
  }
  ctx->stream_out = id;
  return id;
}

// Initialize streaming state for the given number of streams
bool streams_init(usbdpi_ctx_t *ctx, unsigned nstreams, bool retrieve,
                  bool checking, bool retrying, bool send) {
  assert(ctx);

  // Can we support the requested number of streams?
  if (!nstreams || nstreams > USBDPI_MAX_STREAMS) {
    return false;
  }

  if (verbose) {
    printf("[usbdpi] Stream test running with %u streams(s)\n", nstreams);
    printf("[usbdpi] - retrieve %c checking %c retrying %c send %c\n",
           retrieve ? 'Y' : 'N', checking ? 'Y' : 'N', retrying ? 'Y' : 'N',
           send ? 'Y' : 'N');
  }

  // Remember the number of streams and initialize the arbitration of
  // IN and OUT traffic
  ctx->nstreams = nstreams;
  ctx->stream_in = 0U;
  ctx->stream_out = nstreams - 1U;

  for (unsigned id = 0U; id < nstreams; id++) {
    // Poll device for IN packets in streaming test?
    ctx->stream[id].retrieve = retrieve;
    // Attempt to sent OUT packets to device in streaming test?
    ctx->stream[id].send = send;
    // Checking of received data against expected LFSR output
    ctx->stream[id].checking = checking;
    // Request retrying of IN packets, feigning error
    ctx->stream[id].retrying = retrying;
    // Endpoints to be used by this stream
    ctx->stream[id].ep_in = 1U + id;
    ctx->stream[id].ep_out = 1U + id;
    // LFSR state for this byte stream
    ctx->stream[id].tst_lfsr = USBTST_LFSR_SEED(id);
    ctx->stream[id].dpi_lfsr = USBDPI_LFSR_SEED(id);
    // LFSR-controlled packet retrying state
    ctx->stream[id].retry_lfsr = RETRY_LFSR_SEED(id);
    ctx->stream[id].nretries = 0U;
    // No received packets
    ctx->stream[id].received = NULL;
  }
  return true;
}

// Check a data packet received from the test software (usbdev_stream_test)
bool stream_data_check(usbdpi_ctx_t *ctx, usbdpi_stream_t *s,
                       const usbdpi_transfer_t *rx, bool accept) {
  assert(rx);

  // The byte count _includes_ the DATAx PID and the two CRC bytes
  //
  // Note: we are expecting a LFSR-generated byte stream, but we do not
  //       make assumptions about the number or size of the data packets
  //       that make up the stream

  unsigned num_bytes = transfer_length(rx);
  unsigned idx = rx->data_start + 1u;  // Skip past the DATAx PID
  bool ok = false;

  // Validate the received packet - data length valid and checksum present
  if (num_bytes >= sizeof(rx->data) || idx + 2u > num_bytes) {
    printf("[usbdpi] Unexpected/malformed data packet (0x%x 0x%x)\n", idx,
           num_bytes);
  } else {
    // Data field within received packet
    const uint8_t *sp = &rx->data[idx];
    num_bytes -= 3u;

    // Check that the CRC16 checksum of the data field is as expected
    uint16_t rx_crc = sp[num_bytes] | (sp[num_bytes + 1u] << 8);
    uint16_t crc = CRC16(sp, num_bytes);
    if (rx_crc != crc) {
      printf("[usbdpi] Mismatched CRC16 0x%04x received, expected 0x%04x\n",
             rx_crc, crc);
    } else {
      // Data toggle synchronization
      unsigned pid = ctx->ep_in[s->ep_in].next_data;
      if (rx->data[0] == pid) {
        // If we've decided to reject this packet then we still check its
        // content but we do not advance the data toggle because we're
        // pretending that we didn't receive it successfully
        if (accept) {
          ctx->ep_in[s->ep_in].next_data = DATA_TOGGLE_ADVANCE(pid);
        }

        // Tentatively acceptable, but we still have to check and report any and
        // all mismatched bytes
        ok = accept;
      }

      // Iff running a performance investigation, checking may be undesired
      // because it causes us to reject and retry the transmission
      if (s->checking) {
        // Note: use a local copy of the LFSR so that we can check the data
        //       field even on those packets that we choose to reject
        uint8_t tst_lfsr = s->tst_lfsr;
        while (num_bytes-- > 0U) {
          uint8_t recvd = *sp++;
          if (recvd != tst_lfsr) {
            printf(
                "[usbdpi] Mismatched data from device 0x%02x, "
                "expected 0x%02x\n",
                recvd, tst_lfsr);
            ok = false;
          }
          // Advance our local LFSR
          tst_lfsr = LFSR_ADVANCE(tst_lfsr);
        }

        // Update the LFSR only if we've accepted valid data and will not
        // be receiving this data again
        if (accept && ok) {
          s->tst_lfsr = tst_lfsr;
        }
      } else {
        printf("[usbdpi] Warning: Stream data checking disabled\n");
      }
    }
  }

  return ok;
}

// Generate a data packet as if it had been received from the device
usbdpi_transfer_t *stream_data_gen(usbdpi_ctx_t *ctx, usbdpi_stream_t *s,
                                   unsigned len) {
  usbdpi_transfer_t *tr = transfer_alloc(ctx);
  if (tr) {
    // Pretend that we have successfully received the packet with the correct
    // data toggling...
    uint8_t data = ctx->ep_in[s->ep_in].next_data;
    ctx->ep_in[s->ep_in].next_data = DATA_TOGGLE_ADVANCE(data);
    // ...and that the data is as expected
    uint8_t *dp = transfer_data_start(tr, data, len);
    for (unsigned idx = 0U; idx < len; idx++) {
      dp[idx] = s->tst_lfsr;
      s->tst_lfsr = LFSR_ADVANCE(s->tst_lfsr);
    }
    transfer_data_end(tr, dp + len);
  }
  return tr;
}

// Process a received data packet to produce a corresponding reply packet
// by XORing our LFSR output with the received data
//
// Note: for now we do this even if the received data mismatches because
//       only the CPU software has the capacity to decide upon and report
//       test status
usbdpi_transfer_t *stream_data_process(usbdpi_ctx_t *ctx, usbdpi_stream_t *s,
                                       usbdpi_transfer_t *rx) {
  // Note: checkStreamData has already been called on this packet
  assert(rx);

  // The byte count _includes_ the DATAx PID and the two CRC bytes
  unsigned num_bytes = rx->num_bytes;
  unsigned idx = rx->data_start + 1u;  // Skip past the DATAx PID

  // Data field within received packet
  const uint8_t *sp = &rx->data[idx];
  num_bytes -= 3u;

  // Allocate a new buffer for the reply
  usbdpi_transfer_t *reply = transfer_alloc(ctx);
  assert(reply);

  // Construct OUT token packet to the target endpoint, using the
  // appropriate DATAx PID
  const uint8_t ep_out = s->ep_out;
  transfer_token(reply, USB_PID_OUT, ctx->dev_address, ep_out);
  uint8_t *dp =
      transfer_data_start(reply, ctx->ep_out[ep_out].next_data, num_bytes);
  assert(dp);

  while (num_bytes-- > 0U) {
    uint8_t recvd = *sp++;

    // Simply XOR the two LFSR-generated streams together
    *dp++ = recvd ^ s->dpi_lfsr;
    if (verbose) {
      printf("[usbdpi] 0x%02x <- 0x%02x ^ 0x%02x\n", *(dp - 1), recvd,
             s->dpi_lfsr);
    }

    // Advance our LFSR
    //
    // TODO - decide whether we want to do this here; if the device
    // responds with a NAK, requiring us to retry, or we decide to
    // resend the packet, then we don't want to advance again
    s->dpi_lfsr = LFSR_ADVANCE(s->dpi_lfsr);
  }

  transfer_data_end(reply, dp);

  return reply;
}

// Check the stream signature
bool stream_sig_check(usbdpi_ctx_t *ctx, usbdpi_stream_t *s,
                      usbdpi_transfer_t *rx) {
  // Packet should be PID, data field and CRC16
  if (transfer_length(rx) == 3U + 0x10U) {
    const uint8_t *sig = transfer_data_field(rx);
    if (sig) {
      // Signature format:
      //   Bits Description
      //   32   head signature
      //   8    initial vaue of LFSR
      //   8    stream number
      //   16   reserved, SBZ
      //   32   number of bytes to be transferred
      //   32   tail signature
      // Note: all 32-bit quantities are in little endian order

      uint32_t num_bytes = get_le32(&sig[8]);
      if (verbose) {
        printf("[usbdpi] Stream signature at %p head 0x%x tail 0x%x\n", sig,
               get_le32(&sig[0]), get_le32(&sig[12]));
      }

      // Basic validation check; words are transmitted in little endian order
      if (get_le32(&sig[0]) == STREAM_SIGNATURE_HEAD &&
          get_le32(&sig[12]) == STREAM_SIGNATURE_TAIL &&
          // sanity check on transfer length, though we rely upon the CPU
          // software to send, receive and count the number of bytes
          num_bytes > 0U && num_bytes < 0x10000000U && !sig[6] && !sig[7]) {
        // Signature includes the initial value of the device-side LFSR
        s->tst_lfsr = sig[4];
        // Update data toggle
        uint8_t pid = transfer_data_pid(rx);
        ctx->ep_in[s->ep_in].next_data = DATA_TOGGLE_ADVANCE(pid);
        return true;
      }
    }
  }
  return false;
}

// Service streaming data (usbdev_stream_test)
void streams_service(usbdpi_ctx_t *ctx) {
  if (verbose) {
    //    printf("[usbdpi] streams_service hostSt %u in %u out %u\n",
    //    ctx->hostSt,
    //           ctx->stream_in, ctx->stream_out);
  }

  // Maximum time for transmission of a packet ought to be circa 80 bytes of
  // data, 640 bits. Allowing for bitstuffing this means we need to leave ~800
  const unsigned min_time_left = 800U;

  switch (ctx->hostSt) {
    // --------------------------------------------
    // Try to transmit a data packet to the device
    // --------------------------------------------
    case HS_STARTFRAME:
    case HS_STREAMOUT: {
      // Decide whether we have enough time within this frame to attempt
      // another transmission
      uint32_t next_frame = ctx->frame_start + FRAME_INTERVAL;
      if ((next_frame - ctx->tick_bits) > min_time_left) {  // HACK
        unsigned id = out_stream_next(ctx);
        usbdpi_stream_t *s = &ctx->stream[id];
        if (s->send) {
          // Start by trying to transmit a data packet that we've received, if
          // any
          if (s->received) {
            if (ctx->sending) {
              transfer_release(ctx, ctx->sending);
            }

            // Scramble the oldest received packet with our LFSR-generated byte
            // stream and send it to the device
            usbdpi_transfer_t *reply = stream_data_process(ctx, s, s->received);
            transfer_send(ctx, reply);

            uint32_t max_bits =
                transfer_length(ctx->sending) * 10 + 160;  // HACK
            ctx->wait = USBDPI_TIMEOUT(ctx, max_bits);
            ctx->bus_state = kUsbBulkOut;
            ctx->lastrxpid = 0;
            ctx->hostSt = HS_WAITACK;
          } else {
            // Nothing to send, try receiving...
            ctx->hostSt = HS_STREAMIN;
          }
        } else {
          // We're not sending anything - discard any received data
          while (s->received) {
            usbdpi_transfer_t *tr = s->received;
            s->received = tr->next;
            transfer_release(ctx, tr);
          }
          ctx->hostSt = HS_STREAMIN;
        }
      } else {
        // Wait until the next bus frame
        ctx->hostSt = HS_NEXTFRAME;
      }
    } break;

    // Await acknowledgement of the packet that we just transmitted
    case HS_WAITACK:
      if (ctx->sending) {
        // Forget reference to the buffer we just tried to send; the received
        // packet remains in the list of received buffers to try again later
        transfer_release(ctx, ctx->sending);
        ctx->sending = NULL;
      }

      if (ctx->bus_state == kUsbBulkOutAck) {
        bool proceed = false;

        if (verbose) {
          printf("[usbdpi] OUT - response is PID 0x%02x from device (%s)\n",
                 ctx->lastrxpid, decode_pid(ctx->lastrxpid));
        }

        switch (ctx->lastrxpid) {
          case USB_PID_ACK: {
            // Transmitted packet was accepted, so we can retire it...
            usbdpi_stream_t *s = &ctx->stream[ctx->stream_out];
            usbdpi_transfer_t *rx = s->received;
            assert(rx);
            s->received = rx->next;
            transfer_release(ctx, rx);

            uint8_t ep_out = s->ep_out;
            ctx->ep_out[ep_out].next_data =
                DATA_TOGGLE_ADVANCE(ctx->ep_out[ep_out].next_data);
            proceed = true;
          } break;

          // We may receive a NAK from the device if it is unable to receive the
          // packet right now
          case USB_PID_NAK:
            printf(
                "[usbdpi] frame 0x%x tick_bits 0x%x NAK received from device\n",
                ctx->frame, ctx->tick_bits);
            proceed = true;
            break;

          default:
            printf("[usbdpi] Unexpected PID 0x%02x from device (%s)\n",
                   ctx->lastrxpid, decode_pid(ctx->lastrxpid));
            ctx->hostSt = HS_NEXTFRAME;
            break;
        }

        ctx->hostSt = proceed ? HS_STREAMIN : HS_NEXTFRAME;
      } else if (ctx->tick_bits >= ctx->wait) {
        printf("[usbdpi] Timed out waiting for OUT response\n");
        ctx->hostSt = HS_NEXTFRAME;
      }
      break;

    // ---------------------------------------------
    // Try to collect a data packet from the device
    // ---------------------------------------------
    case HS_STREAMIN: {
      // Decide whether we have enough time within this frame to attempt
      // another fetch
      //
      // TODO - find out what the host behaviour should be at this point;
      //        the device must be required to respond within a certain
      //        time interval, and then the bus transmission speed
      //        determines the maximum delay
      uint32_t next_frame = ctx->frame_start + FRAME_INTERVAL;
      if ((next_frame - ctx->tick_bits) > min_time_left) {  // HACK
        unsigned id = in_stream_next(ctx);
        usbdpi_stream_t *s = &ctx->stream[id];
        if (s->retrieve) {
          // Ensure that a buffer is available for constructing a transfer
          usbdpi_transfer_t *tr = ctx->sending;
          if (!tr) {
            tr = transfer_alloc(ctx);
            assert(tr);

            ctx->sending = tr;
          }

          transfer_token(tr, USB_PID_IN, ctx->dev_address, s->ep_in);

          transfer_send(ctx, tr);
          ctx->bus_state = kUsbBulkInToken;
          ctx->hostSt = HS_WAIT_PKT;
          ctx->lastrxpid = 0;
        } else {
          // We're not required to poll for IN data, but if we're sending we
          //   must still fake the reception of valid packet data because
          //   the sw test will be expecting valid data
          if (s->send && !s->received) {
            // For simplicity we just create max length packets
            const unsigned len = USBDEV_MAX_PACKET_SIZE;
            s->received = stream_data_gen(ctx, s, len);
          }
          ctx->hostSt = HS_STREAMOUT;
        }
      } else {
        // Wait until the next bus frame
        ctx->hostSt = HS_NEXTFRAME;
      }
    } break;

    case HS_WAIT_PKT:
      // Wait max time for a response + packet
      ctx->wait = ctx->tick_bits + 18 + 8 + 8 + 64 * 8 + 16;
      ctx->hostSt = HS_ACKIFDATA;
      break;
    case HS_ACKIFDATA:
      if (ctx->bus_state == kUsbBulkInData && ctx->recving) {
        // We have a response from the device
        switch (ctx->lastrxpid) {
          case USB_PID_DATA0:
          case USB_PID_DATA1: {
            // Steal the received packet; it belongs to the stream
            usbdpi_transfer_t *rx = ctx->recving;
            ctx->recving = NULL;
            // Decide whether we want to ACK or NAK this packet
            unsigned id = ctx->stream_in;
            usbdpi_stream_t *s = &ctx->stream[id];
            bool accept = false;
            if (s->retrying && s->nretries) {
              s->nretries--;
            } else {
              // Decide the number of retries for the next data packet
              // Note: by randomizing the number of retries, rather than
              // independently deciding each accept/reject, we guarantee an
              // upper bound on the run time
              switch (s->retry_lfsr & 7U) {
                case 7U:
                  s->nretries = 3U;
                  break;
                case 6U:
                case 5U:
                  s->nretries = 2U;
                  break;
                case 4U:
                  s->nretries = 1U;
                  break;
                default:
                  s->nretries = 0U;
                  break;
              }
              s->retry_lfsr = LFSR_ADVANCE(s->retry_lfsr);
              accept = true;
            }
            if (!accept) {
              printf("[usbdpi] Requesting resend of data\n");
              usb_monitor_log(ctx->mon, "[usbdpi] Requesting resend of data\n");
            }

            if (expect_sig && !s->sig_recvd) {
              // Note: the stream signature is primarily of use on a physical
              // USB connection to a host since the endpoint to port mapping is
              // variable. With t-l sim we can rely upon the first packet being
              // the signature and nothing else
              accept = stream_sig_check(ctx, s, rx);
              transfer_release(ctx, rx);
              // TODO - run time error, signal test failure to the software
              assert(accept);
              if (accept) {
                s->sig_recvd = true;
              }
            } else {
              if (stream_data_check(ctx, s, rx, accept)) {
                // Collect the received packets in preparation for later
                // transmission with modification back to the device
                usbdpi_transfer_t *tr = s->received;
                if (tr) {
                  while (tr->next)
                    tr = tr->next;
                  tr->next = rx;
                } else {
                  s->received = rx;
                }
              } else {
                transfer_release(ctx, rx);
                accept = false;
              }
            }

            usbdpi_transfer_t *tr = ctx->sending;
            assert(tr);

            transfer_status(ctx, tr, accept ? USB_PID_ACK : USB_PID_NAK);

            ctx->hostSt = HS_STREAMOUT;
          } break;

          case USB_PID_NAK:
            // No data available
            ctx->hostSt = HS_STREAMOUT;
            break;

          default:
            printf("[usbdpi] Unexpected PID 0x%02x from device (%s)\n",
                   ctx->lastrxpid, decode_pid(ctx->lastrxpid));
            ctx->hostSt = HS_NEXTFRAME;
            break;
        }
      } else if (ctx->tick_bits >= ctx->wait) {
        printf("[usbdpi] Timed out waiting for IN response\n");
        ctx->hostSt = HS_NEXTFRAME;
      }
      break;

    default:
      break;
  }
}