#include "wokwi-api.h"
#include <stdlib.h>
#include <stdio.h>

typedef struct {
  pin_t cs, dc, res;
  spi_dev_t spi;
  buffer_t fb;
  uint32_t width, height;
  uint8_t spi_buffer[1];
  uint8_t page;
  uint8_t col;
} chip_state_t;

// Forward declarations
static void on_cs_change(void *user_data, pin_t pin, uint32_t value);
static void chip_spi_done(void *user_data, uint8_t *buffer, uint32_t count);

void chip_init(void) {
  printf("Initing CH1115 chip\n");
  fflush(stdout);

  chip_state_t *chip = malloc(sizeof(chip_state_t));
  chip->fb = framebuffer_init(&chip->width, &chip->height);
  
  chip->cs = pin_init("CS", INPUT_PULLUP);
  chip->dc = pin_init("DC", INPUT);
  chip->res = pin_init("RES", INPUT_PULLUP);

  const spi_config_t spi_cfg = {
    .sck = pin_init("SCL", INPUT),
    .mosi = pin_init("SDA", INPUT),
    .miso = NO_PIN,
    .mode = 0,
    .done = chip_spi_done,
    .user_data = chip,
  };
  chip->spi = spi_init(&spi_cfg);

  const pin_watch_config_t cs_watch = {
    .edge = BOTH,
    .pin_change = on_cs_change,
    .user_data = chip,
  };
  pin_watch(chip->cs, &cs_watch);
}

static void on_cs_change(void *user_data, pin_t pin, uint32_t value) {
  chip_state_t *chip = (chip_state_t*)user_data;
  if (value == LOW) {
    // Start receiving 1 byte at a time
    spi_start(chip->spi, chip->spi_buffer, 1);
  } else {
    spi_stop(chip->spi);
  }
}

static void chip_spi_done(void *user_data, uint8_t *buffer, uint32_t count) {
  chip_state_t *chip = (chip_state_t*)user_data;
  if (!count) return;

  uint8_t val = buffer[0];
  bool is_data = pin_read(chip->dc) == HIGH;

  if (is_data) {
    // Write 8 vertical pixels to the Wokwi RGBA Framebuffer
    for (int i = 0; i < 8; i++) {
      // 0xFFFFFFFF = White, 0x000000FF = Black
      uint32_t color = (val & (1 << i)) ? 0xFFFFFFFF : 0x000000FF; 
      uint32_t y = (chip->page * 8) + i;
      uint32_t offset = (y * chip->width + chip->col) * 4;
      buffer_write(chip->fb, offset, &color, 4);
    }
    chip->col++; // Auto-increment column
  } else {
    // CH1115 Core Commands
    if ((val & 0xB0) == 0xB0) chip->page = val & 0x0F; // Set Page Start
    else if ((val & 0xF0) == 0x00) chip->col = (chip->col & 0xF0) | (val & 0x0F); // Lower Col
    else if ((val & 0xF0) == 0x10) chip->col = (chip->col & 0x0F) | ((val & 0x0F) << 4); // Upper Col
  }

  // Continue receiving if CS is still active
  if (pin_read(chip->cs) == LOW) {
    spi_start(chip->spi, chip->spi_buffer, 1);
  }
}