// Copyright (c) 2026 Cesanta Software Limited
// All rights reserved

#include "hardware/pio.h"
#include "mongoose.h"

#define WIFI_SSID "YOUR_WIFI_NETWORK_NAME"  // SET THIS!
#define WIFI_PASS "YOUR_WIFI_PASSWORD"      // SET THIS!

#define SPIPIO pio0
#define SPISM 0
#define DATAPIN 2
#define CLKPIN 3
#define CSPIN 1
#define PWRPIN 0

#define piospi_wrap_target 0
#define piospi_wrap 6
#define piospi_pio_version 0
#define piospi_offset_rx_ 7u

static const uint16_t piospi_program_instructions[] = {
    //     .wrap_target
    0x80a0,  //  0: pull   block           side 0
    0xa027,  //  1: mov    x, osr          side 0
    0x6068,  //  2: out    null, 8         side 0
    0x6001,  //  3: out    pins, 1         side 0
    0x1043,  //  4: jmp    x--, 3          side 1
    0xe000,  //  5: set    pins, 0         side 0
    0xc020,  //  6: irq    wait 0          side 0
             //     .wrap
    0xe080,  //  7: set    pindirs, 0      side 0
    0x6020,  //  8: out    x, 32           side 0
    0x4001,  //  9: in     pins, 1         side 0
    0x1049,  // 10: jmp    x--, 9          side 1
    0xe000,  // 11: set    pins, 0         side 0
    0xe081,  // 12: set    pindirs, 1      side 0
    0xc020,  // 13: irq    wait 0          side 0
    0x0000,  // 14: jmp    0               side 0
};

static const struct pio_program piospi_program = {
    .instructions = piospi_program_instructions,
    .length = 15,
    .origin = -1,
    .pio_version = 0,
#if PICO_PIO_VERSION > 0
    .used_gpio_ranges = 0x0
#endif
};

static pio_sm_config piospi_program_get_default_config(uint offset) {
  pio_sm_config c = pio_get_default_sm_config();
  sm_config_set_wrap(&c, offset + piospi_wrap_target, offset + piospi_wrap);
  sm_config_set_sideset(&c, 1, false, false);
  return c;
}

static pio_sm_config piospi_init(PIO pio, uint sm, uint addr, uint data,
                                 uint clk) {
  pio_gpio_init(pio, data);
  pio_gpio_init(pio, clk);
  pio_sm_set_pins_with_mask(pio, sm, 0, (1 << data) | (1 << clk));
  pio_sm_config c = piospi_program_get_default_config(addr);
  sm_config_set_out_pins(&c, data, 1);
  sm_config_set_in_pins(&c, data);
  sm_config_set_set_pins(&c, data, 1);
  sm_config_set_sideset_pins(&c, clk);
  pio_sm_set_consecutive_pindirs(pio, sm, data, 1, true);
  pio_sm_set_consecutive_pindirs(pio, sm, clk, 1, true);
  pio_sm_set_pindirs_with_mask(pio, sm, (1 << data) | (1 << clk),
                               (1 << data) | (1 << clk));
  sm_config_set_in_shift(&c, false, true, 8);  // push bytes, MSB first, auto
#if PICO_RP2040
  sm_config_set_clkdiv(&c, 18);  // Run at 133/1 = <50MHz (2x data rate)
#else
  sm_config_set_clkdiv(&c, 20);  // Run at 150/1 = <50MHz (2x data rate)
#endif
  return c;
}
static void piospi_tx(PIO pio, uint sm, pio_sm_config *c, uint addr,
                      uint data) {
  sm_config_set_fifo_join(c, PIO_FIFO_JOIN_TX);
  sm_config_set_out_shift(c, false, true, 8);  // pull bytes, MSB first, auto
  pio_sm_init(pio, sm, addr, c);
}
static void piospi_rx(PIO pio, uint sm, pio_sm_config *c, uint addr,
                      uint data) {
  sm_config_set_fifo_join(c, PIO_FIFO_JOIN_NONE);
  sm_config_set_out_shift(c, false, true, 32);  // pull words, MSB first, auto
  pio_sm_init(pio, sm, addr + piospi_offset_rx_, c);
}
static void piospi_done(PIO pio, uint sm) {
  while (!pio_interrupt_get(pio, 0));
  pio_sm_set_enabled(pio, sm, false);
  pio_interrupt_clear(pio, 0);
}

static pio_sm_config s_piospi_sm_config;
static uint s_piospi_sm_addr;

void piospi_write(uint8_t *data, size_t len) {
  size_t initial_len = len;
  piospi_tx(SPIPIO, SPISM, &s_piospi_sm_config, s_piospi_sm_addr, DATAPIN);
  __compiler_memory_barrier();
  pio_sm_put(SPIPIO, SPISM, (len * 8) - 1);  // bits to transmit
  pio_sm_set_enabled(SPIPIO, SPISM, true);
  while (len--) pio_sm_put_blocking(SPIPIO, SPISM, (*data++) << 24);
  piospi_done(SPIPIO, SPISM);
}

// Read data block from SPI interface
void piospi_read(uint8_t *data, size_t len) {
  piospi_rx(SPIPIO, SPISM, &s_piospi_sm_config, s_piospi_sm_addr, DATAPIN);
  __compiler_memory_barrier();
  pio_sm_put(SPIPIO, SPISM, (len * 8) - 1);  // bits to receive
  pio_sm_set_enabled(SPIPIO, SPISM, true);
  while (len--) *data++ = pio_sm_get_blocking(SPIPIO, SPISM);
  piospi_done(SPIPIO, SPISM);
}

void hwspecific_spi_init(void) {
  gpio_init(PWRPIN);
  gpio_set_dir(PWRPIN, GPIO_OUT);
  gpio_put(PWRPIN, 0);
  gpio_init(CSPIN);
  gpio_set_dir(CSPIN, GPIO_OUT);
  gpio_put(CSPIN, 0);
  // init SPI pins so they are idle and stable during power up
  gpio_init(DATAPIN);
  gpio_set_dir(DATAPIN, GPIO_OUT);
  gpio_put(DATAPIN, 0);
  gpio_init(CLKPIN);
  gpio_set_dir(CLKPIN, GPIO_OUT);
  gpio_put(CLKPIN, 0);
  mg_delayms(100);
  gpio_put(CSPIN, 1);
  gpio_put(PWRPIN, 1);
  s_piospi_sm_addr = pio_add_program(SPIPIO, &piospi_program);
  s_piospi_sm_config =
      piospi_init(SPIPIO, SPISM, s_piospi_sm_addr, DATAPIN, CLKPIN);
  mg_delayms(50);
}

void hwspecific_spi_begin(void *arg) {
  gpio_put(CSPIN, 0);
  (void) arg;
}

// either write or read, not both
void hwspecific_spi_txn(void *arg, uint8_t *txdata, uint8_t *rxdata,
                        size_t len) {
  if (txdata != NULL) piospi_write(txdata, len);
  if (rxdata != NULL) piospi_read(rxdata, len);
  (void) arg;
}

void hwspecific_spi_end(void *arg) {
  gpio_put(CSPIN, 1);
  (void) arg;
}

static const struct mg_tcpip_spi spi = {NULL, hwspecific_spi_begin,
                                        hwspecific_spi_end, hwspecific_spi_txn};

#ifndef CYW43_RESOURCE_ATTRIBUTE
#define CYW43_RESOURCE_ATTRIBUTE
#endif
#include "pico-sdk/lib/cyw43-driver/firmware/w43439A0_7_95_49_00_combined.h"
#include "pico-sdk/lib/cyw43-driver/firmware/wifi_nvram_43439.h"
static const struct mg_tcpip_driver_cyw_firmware fw = {
    (const uint8_t *) w43439A0_7_95_49_00_combined,
    (size_t) CYW43_WIFI_FW_LEN,
    (const uint8_t *) wifi_nvram_4343,
    (size_t) sizeof(wifi_nvram_4343),
    (const uint8_t *) (w43439A0_7_95_49_00_combined +
                       sizeof(w43439A0_7_95_49_00_combined) - CYW43_CLM_LEN),
    (size_t) CYW43_CLM_LEN};

// mif user states
enum { AP, SCANNING, STOPPING_AP, CONNECTING, READY };
static unsigned int state;

static void mif_fn(struct mg_tcpip_if *ifp, int ev, void *ev_data) {
  // TODO(): should we include this inside ifp ? add an fn_data ?
  if (ev == MG_TCPIP_EV_STATE_CHANGE) {
    MG_INFO(("State change: %u", *(uint8_t *) ev_data));
  }
  switch (state) {
    case AP:  // we are in AP mode, wait for a user connection to trigger a scan
              // or a connection to a network
      if (ev == MG_TCPIP_EV_STATE_CHANGE &&
          *(uint8_t *) ev_data == MG_TCPIP_STATE_READY) {
        MG_INFO(("Access Point READY !"));

        // simulate user request to scan for networks
        bool res = mg_wifi_scan();
        MG_INFO(("Starting scan: %s", res ? "OK" : "FAIL"));
        if (res) state = SCANNING;
      }
      break;
    case SCANNING:
      if (ev == MG_TCPIP_EV_WIFI_SCAN_RESULT) {
        struct mg_wifi_scan_bss_data *bss =
            (struct mg_wifi_scan_bss_data *) ev_data;
        MG_INFO(("BSS: %.*s (%u) (%M) %d dBm %u", bss->SSID.len, bss->SSID.buf,
                 bss->channel, mg_print_mac, bss->BSSID, (int) bss->RSSI,
                 bss->security));
      } else if (ev == MG_TCPIP_EV_WIFI_SCAN_END) {
        // struct mg_tcpip_driver_cyw_data *d = (struct mg_tcpip_driver_cyw_data
        // *) ifp->driver_data;
        MG_INFO(("Wi-Fi scan finished"));

        // simulate user selection of a network (1/2: stop AP)
        bool res = mg_wifi_ap_stop();
        MG_INFO(("Manually stopping AP: %s", res ? "OK" : "FAIL"));
        if (res) state = STOPPING_AP;
        // else we have a hw/fw problem
      }
      break;
    case STOPPING_AP:
      if (ev == MG_TCPIP_EV_STATE_CHANGE &&
          *(uint8_t *) ev_data == MG_TCPIP_STATE_DOWN) {
        struct mg_wifi_data *wifi =
            &((struct mg_tcpip_driver_cyw_data *) ifp->driver_data)->wifi;
        wifi->apmode = false;

        // simulate user selection of a network (2/2: actual connect)
        bool res = mg_wifi_connect(wifi);
        MG_INFO(("Manually connecting: %s", res ? "OK" : "FAIL"));
        if (res) {
          state = CONNECTING;
        }  // else manually start AP as below
      }
      break;
    case CONNECTING:
      if (ev == MG_TCPIP_EV_STATE_CHANGE &&
          *(uint8_t *) ev_data == MG_TCPIP_STATE_READY) {
        MG_INFO(("READY!"));
        state = READY;

        // simulate user code disconnection and go back to AP mode (1/2:
        // disconnect)
        bool res = mg_wifi_disconnect();
        MG_INFO(("Manually disconnecting: %s", res ? "OK" : "FAIL"));
      } else if (ev == MG_TCPIP_EV_WIFI_CONNECT_ERR) {
        MG_ERROR(("Wi-Fi connect failed"));
        // manually start AP as below
      }
      break;
    case READY:
      // go back to AP mode after a disconnection (simulation 2/2), you could
      // retry
      if (ev == MG_TCPIP_EV_STATE_CHANGE &&
          *(uint8_t *) ev_data == MG_TCPIP_STATE_DOWN) {
        struct mg_wifi_data *wifi =
            &((struct mg_tcpip_driver_cyw_data *) ifp->driver_data)->wifi;
        bool res = mg_wifi_ap_start(wifi);
        MG_INFO(("Disconnected"));
        MG_INFO(("Manually starting AP: %s", res ? "OK" : "FAIL"));
        if (res) {
          state = AP;
          wifi->apmode = true;
        }
      }
      break;
  }
}

static struct mg_tcpip_driver_cyw_data d = {
    {WIFI_SSID, WIFI_PASS, "mongoose", "mongoose", 0, 0, 0, 0, 10, true},
    (struct mg_tcpip_spi *) &spi,
    (struct mg_tcpip_driver_cyw_firmware *) &fw,
    false,
};

void driver_init(struct mg_mgr *mgr) {
  hwspecific_spi_init();

  d.wifi.apip = MG_IPV4(192, 168, 169, 1),
  d.wifi.apmask = MG_IPV4(255, 255, 255, 0),
  state = d.wifi.apmode ? AP : CONNECTING;

  // Initialise Mongoose network stack
  // Either set use_dhcp or enter a static config.
  // For static configuration, specify IP/mask/GW in network byte order
  struct mg_tcpip_if mif = {
      .ip = 0,
      .driver = (struct mg_tcpip_driver *) &mg_tcpip_driver_cyw,
      .driver_data = (struct mg_tcpip_driver_cyw_data *) &d,
      .fn = mif_fn,
      //      .recv_queue.size = 8192
  };
  mg_tcpip_init(mgr, &mif);
}