/*
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "sw/device/lib/usbdev/usbdev.h"

#include "hw/top_matcha/sw/autogen/top_matcha.h"
#include "sw/device/lib/arch/device.h"
#include "sw/device/lib/dif/dif_pinmux.h"
#include "sw/device/lib/testing/pinmux_testutils.h"
#include "sw/device/lib/testing/test_framework/check.h"
#include "sw/device/lib/testing/usb_testutils_controlep.h"

const unsigned kDataEp = 1;
const unsigned kControlEp = 2;

typedef struct {
  uint32_t size;
} ControlPkt;

static dif_pinmux_t pinmux;

static const uint8_t kUsbdevProfileBulkDescriptors[] = {
    USB_CFG_DSCR_HEAD(USB_CFG_DSCR_LEN + (2 * (USB_INTERFACE_DSCR_LEN +
                                               2 * (USB_EP_DSCR_LEN))),
                      2),
    VEND_INTERFACE_DSCR(0, 2, 0x51, 1),
    USB_BULK_EP_DSCR(0, kDataEp, USBDEV_MAX_PACKET_SIZE, 0),
    USB_BULK_EP_DSCR(1, kDataEp, USBDEV_MAX_PACKET_SIZE, 0),

    VEND_INTERFACE_DSCR(1, 2, 0x51, 1),
    USB_BULK_EP_DSCR(0, kControlEp, USBDEV_MAX_PACKET_SIZE, 0),
    USB_BULK_EP_DSCR(1, kControlEp, USBDEV_MAX_PACKET_SIZE, 0),
};

static void rx_data(void* ctx, dif_usbdev_rx_packet_info_t packet_info,
                    dif_usbdev_buffer_t buf) {
  UsbdevContext* uctx = (UsbdevContext*)ctx;
  size_t bytes_written;
  CHECK_DIF_OK(dif_usbdev_buffer_read(
      uctx->usbdev.dev, uctx->usbdev.buffer_pool, &buf,
      uctx->dst + uctx->bytes_received, /*dst_len=*/1024, &bytes_written));
  uctx->bytes_received += packet_info.length;
}

static void rx_control(void* ctx, dif_usbdev_rx_packet_info_t packet_info,
                       dif_usbdev_buffer_t buf) {
  ControlPkt control;
  size_t bytes_written;
  UsbdevContext* uctx = (UsbdevContext*)ctx;
  CHECK(packet_info.length == sizeof(control));
  CHECK_DIF_OK(dif_usbdev_buffer_read(
      uctx->usbdev.dev, uctx->usbdev.buffer_pool, &buf, (uint8_t*)&control,
      sizeof(control), &bytes_written));
  CHECK(bytes_written == sizeof(control));
  uctx->bytes_to_receive = control.size;
  if (uctx->verbose) {
    LOG_INFO("The host will send %d bytes of data.", uctx->bytes_to_receive);
  }
}

dif_result_t UsbdevInit(UsbdevContext* ctx, UsbdevProfile profile) {
  if (ctx == NULL) {
    return kDifBadArg;
  }
  if (profile != kUsbdevProfileBulk) {
    return kDifBadArg;
  }

  memset(&ctx->usbdev, 0, sizeof(ctx->usbdev));
  memset(&ctx->usbdev_control, 0, sizeof(ctx->usbdev_control));
  ctx->finished = false;
  ctx->bytes_received = 0;
  ctx->bytes_to_receive = ~0U;

  CHECK_DIF_OK(dif_pinmux_init(
      mmio_region_from_addr(TOP_MATCHA_PINMUX_AON_BASE_ADDR), &pinmux));
  pinmux_testutils_init(&pinmux);

  CHECK_DIF_OK(dif_pinmux_input_select(&pinmux,
                                       kTopMatchaPinmuxPeripheralInUsbdevSense,
                                       kTopMatchaPinmuxInselIoc7));
  // VBUS SENSE on Nexus is active-low. USBDEV wants active high, so invert the
  // signal via pinmux.
  if (kDeviceType == kDeviceFpgaNexus) {
    dif_pinmux_pad_attr_t attrs;
    CHECK_DIF_OK(dif_pinmux_pad_get_attrs(&pinmux, kTopMatchaMuxedPadsIoc7,
                                          kDifPinmuxPadKindMio, &attrs));
    attrs.flags |= kDifPinmuxPadAttrInvertLevel;
    CHECK_DIF_OK(dif_pinmux_pad_write_attrs(
        &pinmux, kTopMatchaMuxedPadsIoc7, kDifPinmuxPadKindMio, attrs, &attrs));
  }

  bool en_diff_rcvr = kDeviceType == kDeviceFpgaNexus ? true : false;
  usb_testutils_init(&ctx->usbdev, false, en_diff_rcvr, false);
  usb_testutils_controlep_init(&ctx->usbdev_control, &ctx->usbdev, /*ep=*/0,
                               kUsbdevProfileBulkDescriptors,
                               sizeof(kUsbdevProfileBulkDescriptors), NULL, 0);

  if (ctx->verbose) {
    LOG_INFO("Waiting for USB to finish configuring...");
  }
  while (ctx->usbdev_control.device_state != kUsbTestutilsDeviceConfigured) {
    usb_testutils_poll(&ctx->usbdev);
  }
  if (ctx->verbose) {
    LOG_INFO("USB finished configuration!");
  }

  usb_testutils_endpoint_setup(&ctx->usbdev, kDataEp, kUsbdevOutStream, ctx,
                               /*tx_done=*/NULL, rx_data, /*flush=*/NULL,
                               /*reset=*/NULL);
  usb_testutils_endpoint_setup(&ctx->usbdev, kControlEp, kUsbdevOutStream, ctx,
                               /*tx_done=*/NULL, rx_control, /*flush=*/NULL,
                               /*reset=*/NULL);

  return kDifOk;
}

dif_result_t UsbdevPoll(UsbdevContext* ctx) {
  usb_testutils_poll(&ctx->usbdev);
  if (ctx->bytes_received >= ctx->bytes_to_receive) {
    ctx->finished = true;
  }
  return kDifOk;
}