#include "led.h"
#include "stm32_dma.h"
#include <string.h>

extern uint8_t __heap_base__[];
extern uint8_t __heap_end__[];

led_ctx_t led_drivers[NUM_DRIVERS];
ws2812_signal_t signals[C_PALETTE_SIZE] = {0};
dma_ctx_t transfer_dma[2] = {0};
uint8_t * const framebuffer = __heap_base__;
event_listener_t event_listeners[NUM_DRIVERS];

led_t palette[C_PALETTE_SIZE] = {
  {0, 0, 0},       // C_OFF
  {10, 0, 0},      // C_DARK_GREEN
  {0, 10, 0},      // C_DARK_RED
  {0, 0, 10},      // C_DARK_BLUE
  {0, 6, 6},       // C_DARK_PURPLE
  {6, 0, 6},       // C_DARK_CYAN
  {3, 3, 3},       // C_DARK_WHITE
  {100, 0, 0},     // C_GREEN
  {0, 100, 0},     // C_RED
  {0, 0, 100},     // C_BLUE
  {255, 255, 255}, // C_WHITE
  {2, 2, 2},       // C_GRAY
};

void dma_copy_led(uint8_t palette_index, uint8_t * buf);

void fill_dmabuf(led_ctx_t * driver) {
  uint32_t curr = 0;
  for(uint32_t i = 0; i < LEDS_PER_IRQ; i++) {
    curr = driver->buf_cnt + i;
    if(curr < RESET_PRE){
      memset(driver->dma_buf_ptr + (i*BITS_PER_LED), 0x00, BITS_PER_LED);
    } else if(curr < (driver->tx_led_cnt + RESET_PRE)) {
      uint32_t led_index = curr - RESET_PRE;
      uint8_t palette_idx = driver->framebuffer[led_index];
      uint8_t * buf = driver->dma_buf_ptr + (i*BITS_PER_LED);
      uint32_t cndtr0, cndtr1;
      size_t dma_idx;
      // Loop until a DMA is ready
      while(true) {
        // Check if transfer_dma 0 is ready
        cndtr0 = transfer_dma[0].stream->channel->CNDTR;
        if(cndtr0 == 0) {
          dma_idx = 0;
          dmaStreamDisable(transfer_dma[0].stream);
          break;
        }
        // Check if transfer_dma 1 is ready
        cndtr1 = transfer_dma[1].stream->channel->CNDTR;
        if(cndtr1 == 0) {
          dma_idx = 1;
          dmaStreamDisable(transfer_dma[1].stream);
          break;
        }
      }
      dmaStartMemCopy(transfer_dma[dma_idx].stream, STM32_DMA_CR_PSIZE_WORD | STM32_DMA_CR_MSIZE_WORD | STM32_DMA_CR_PL(2), signals[palette_idx].bit_sigs, buf, BITS_PER_LED/4);
    } else {
      uint8_t * ptr = driver->dma_buf_ptr + (i*BITS_PER_LED);
      memset(ptr, 0x00, DMA_BUF_HALF - (i*BITS_PER_LED));
      break;
    }
  }
  driver->buf_cnt += LEDS_PER_IRQ;
}

void dma_irq(void * args, uint32_t flags) {
  chSysLockFromISR();
  led_ctx_t * driver = args;

  driver->tx_cnt += LEDS_PER_IRQ;

  if(flags & STM32_DMA_ISR_HTIF) { // during HT the first half is done, so overwrite it
    driver->dma_buf_ptr = &driver->dma_buf[0];
  } else if(flags & STM32_DMA_ISR_TCIF) { // during TC the second half is done, so overwrite it
    driver->dma_buf_ptr = &driver->dma_buf[DMA_BUF_HALF];
  } else if(flags & STM32_DMA_ISR_TEIF) { // Error flag
    chSysHalt("dma error interrupt flag");
  } else {
    chSysHalt("dma isr called without flags");
  }

  // If the interrupt was called after the number of LEDs needed has been written,
  // stop the transfer and signal the thread
  if(driver->tx_cnt >= (driver->tx_led_cnt + RESET_PRE)) {
    driver->dma_stream->channel->CCR &= ~(STM32_DMA_CR_EN);
    // Disable timer DMA request and pull CCR low so the WS2812B sees a reset pulse
    switch(driver->tim_channel) {
      case 0: driver->pwm_driver->tim->DIER &= ~STM32_TIM_DIER_CC1DE; break;
      case 1: driver->pwm_driver->tim->DIER &= ~STM32_TIM_DIER_CC2DE; break;
      case 2: driver->pwm_driver->tim->DIER &= ~STM32_TIM_DIER_CC3DE; break;
      case 3: driver->pwm_driver->tim->DIER &= ~STM32_TIM_DIER_CC4DE; break;
    }
    driver->pwm_driver->tim->CCR[driver->tim_channel] = 0;
    driver->upd_lck = 0;
    driver->done = 0;
    driver->tx_cnt = 0;
    driver->tx_led_cnt = 0;
    chEvtBroadcastFlagsI(&driver->event_source, LED_EVENT_DONE);
  } else { // otherwise fill the dma buffer with new leds
    chEvtBroadcastFlagsI(&driver->event_source, LED_EVENT_BUFRDY);
  }

  chSysUnlockFromISR();
}

int setup_led_transfer(led_ctx_t * driver) {
  if(driver->upd_lck) {
    return 1;
  }

  // Setup the transfer buffer, and initialize the current count
  driver->upd_lck = 1;
  driver->done = 0;
  driver->tx_cnt = 0;
  driver->buf_cnt = 0;
  driver->tx_led_cnt = driver->num_leds;
  // Setup the first LEDS_PER_IRQ leds
  driver->dma_buf_ptr = &driver->dma_buf[0];
  fill_dmabuf(driver);
  driver->dma_buf_ptr = &driver->dma_buf[DMA_BUF_HALF];
  fill_dmabuf(driver);

  // Wait for all pending transfer_dma copies to land before the PWM DMA starts reading
  while (transfer_dma[0].stream->channel->CNDTR != 0 ||
         transfer_dma[1].stream->channel->CNDTR != 0) {}

  // Set DMA circular
  dmaStreamSetTransactionSize(driver->dma_stream, DMA_BUF_LEN);
  dmaStreamSetMemory0(driver->dma_stream, driver->dma_buf);
  dmaStreamClearInterrupt(driver->dma_stream);
  return 0;
}

int start_led_transfer(led_ctx_t * driver) {
  driver->dma_stream->channel->CCR |= STM32_DMA_CR_CIRC | STM32_DMA_CR_HTIE | STM32_DMA_CR_TCIE | STM32_DMA_CR_EN;
  switch(driver->tim_channel) {
  case 0:
    driver->pwm_driver->tim->DIER |= STM32_TIM_DIER_CC1DE;
    break;
  case 1:
    driver->pwm_driver->tim->DIER |= STM32_TIM_DIER_CC2DE;
    break;
  case 2:
    driver->pwm_driver->tim->DIER |= STM32_TIM_DIER_CC3DE;
    break;
  case 3:
    driver->pwm_driver->tim->DIER |= STM32_TIM_DIER_CC4DE;
    break;
  }
  driver->pwm_driver->tim->CCR[driver->tim_channel] = 0;

  return 0;
}

PWMConfig pwmcfg = {
                     PWM_TIM_FREQ,
                     PWM_PERIOD, // PWM Period of 800 kHz
                     NULL,
                     {
                       {PWM_OUTPUT_ACTIVE_HIGH, NULL},
                       {PWM_OUTPUT_DISABLED, NULL},
                       {PWM_OUTPUT_DISABLED, NULL},
                       {PWM_OUTPUT_DISABLED, NULL}
                     },
                     0,
                     0,
                     0
                   };

extern led_cfg_t led_cfgs[NUM_DRIVERS];

void driver_init(led_cfg_t * cfg, thread_t * draw_thread) {
  if(cfg->num >= NUM_DRIVERS){
    chSysHalt("Tried to initialize a driver outside the max driver range");
  }

  led_ctx_t * driver = &led_drivers[cfg->num];
  driver->num = cfg->num;
  driver->num_leds = cfg->num_leds;

  uint32_t offset = 0;
  for (uint8_t i = 0; i < cfg->num; i++) {
    offset += led_cfgs[i].num_leds;
  }
  driver->framebuffer = framebuffer + offset;
  driver->pwm_driver = cfg->pwm_driver;
  driver->tim_channel = cfg->tim_channel;
  driver->draw_thread = draw_thread;
  chEvtObjectInit(&driver->event_source);
  driver->dma_stream = dmaStreamAlloc(cfg->dma_stream_id, 0, dma_irq, driver);
  if(driver->dma_stream == NULL) {
    chSysHalt("Failed to allocate DMA stream for driver");
  }

  pwmStart(cfg->pwm_driver, &pwmcfg);
  cfg->pwm_driver->tim->CR1 |= STM32_TIM_CR1_ARPE;
  cfg->pwm_driver->tim->CR2 |= STM32_TIM_CR2_CCDS;

  dmaStreamSetPeripheral(driver->dma_stream, &(cfg->pwm_driver->tim->CCR[cfg->tim_channel]));
  dmaStreamSetMode(driver->dma_stream, STM32_DMA_CR_DIR_M2P | STM32_DMA_CR_PL(1) | STM32_DMA_CR_MINC | STM32_DMA_CR_PSIZE_HWORD | STM32_DMA_CR_MSIZE_BYTE);

  driver->port = cfg->port;
  driver->pin = cfg->pin;
  palSetPadMode(driver->port, driver->pin, PAL_MODE_ALTERNATE(cfg->af));
  pwmEnableChannel(driver->pwm_driver, driver->tim_channel, 0);
}

led_cfg_t led_cfgs[NUM_DRIVERS] = {
  {STM32_DMA_STREAM_ID(1, 4), &PWMD3,  0, GPIOB, 4, 1, 0, 10},
  {STM32_DMA_STREAM_ID(1, 5), &PWMD15, 0, GPIOA, 2, 0, 1, 10},
  {STM32_DMA_STREAM_ID(1, 3), &PWMD16, 0, GPIOA, 6, 5, 2, 10},
  {STM32_DMA_STREAM_ID(1, 1), &PWMD17, 0, GPIOA, 7, 5, 3, 10},
  {STM32_DMA_STREAM_ID(1, 2), &PWMD1,  0, GPIOA, 8, 2, 4, 10},
};

ws2812_signal_t convert_palette(led_t * colour) {
  ws2812_signal_t ret = {0};
  memset(ret.bit_sigs, SIG_LOW, BITS_PER_LED);
  for(uint32_t i = 0; i < 8; i++) {
    if(colour->g & (1 << i)) {
      ret.bit_sigs[7 - i] = SIG_HIGH;
    }

    if(colour->r & (1 << i)) {
      ret.bit_sigs[15 - i] = SIG_HIGH;
    }

    if(colour->b & (1 << i)) {
      ret.bit_sigs[23 - i] = SIG_HIGH;
    }
  }

  return ret;
}

void led_init(thread_t * draw_thread) {
  size_t total_leds = 0;
  for (size_t i = 0; i < NUM_DRIVERS; i++) {
    total_leds += led_cfgs[i].num_leds;
  }
  if (total_leds > (size_t)((uint8_t *)__heap_end__ - (uint8_t *)__heap_base__)) {
    chSysHalt(NULL);
  }

  for(size_t i = 0; i < C_PALETTE_SIZE; i++) {
    signals[i] = convert_palette(&palette[i]);
  }

  for(size_t i = 0; i < NUM_DRIVERS; i++) {
    driver_init(&led_cfgs[i], draw_thread);
  }

  transfer_dma[0].id = 0;
  transfer_dma[0].stream = dmaStreamAlloc(STM32_DMA_STREAM_ID(1, 6), 3, NULL, NULL);

  transfer_dma[1].id = 1;
  transfer_dma[1].stream = dmaStreamAlloc(STM32_DMA_STREAM_ID(1, 7), 3, NULL, NULL);

  if(transfer_dma[0].stream == NULL || transfer_dma[1].stream == NULL) {
    chSysHalt("Failed to allocate DMA streams for led transfers");
  }

  memset(framebuffer, 0x00, total_leds);
}

void led_draw(event_listener_t* event_listeners) {
  uint32_t draws_done = 0;
  eventmask_t evt = {0};

  for(uint32_t driver_index = 0; driver_index < NUM_DRIVERS; driver_index++) {
    setup_led_transfer(&led_drivers[driver_index]);
  }
  for(uint32_t driver_index = 0; driver_index < NUM_DRIVERS; driver_index++) {
    start_led_transfer(&led_drivers[driver_index]);
    for(uint32_t i = 0; i < 1000; i++) {
      __asm volatile("NOP");
    }
  }

  // wait for all the drivers to finish drawing
  while(draws_done < NUM_DRIVERS) {
    evt = chEvtWaitAny(ALL_EVENTS);
    for(uint8_t index = 0; index < NUM_DRIVERS; index++) {
      if(!(evt & EVENT_MASK(index))) {
        continue;
      }
      eventflags_t flags = chEvtGetAndClearFlags(&event_listeners[index]);
      if(flags & LED_EVENT_DONE) {
        draws_done += 1;
      } else if(flags & LED_EVENT_BUFRDY) {
        fill_dmabuf(&led_drivers[index]);
      }
    }
  }
}

size_t led_fb_used(void) {
  size_t total = 0;
  for (uint8_t i = 0; i < NUM_DRIVERS; i++) {
    total += led_drivers[i].num_leds;
  }
  return total;
}

void led_resize(uint8_t driver_num, uint8_t new_num_leds) {
  if (driver_num >= NUM_DRIVERS) return;

  size_t new_total = 0;
  size_t old_total = led_drivers[driver_num].num_leds;
  for (uint8_t i = 0; i < NUM_DRIVERS; i++) {
    new_total += (i == driver_num) ? new_num_leds : led_drivers[i].num_leds;
  }
  if (new_total > (size_t)((uint8_t *)__heap_end__ - (uint8_t *)__heap_base__)) return;

  uint32_t new_offsets[NUM_DRIVERS];
  uint32_t off = 0;
  for (uint8_t i = 0; i < NUM_DRIVERS; i++) {
    new_offsets[i] = off;
    off += (i == driver_num) ? new_num_leds : led_drivers[i].num_leds;
  }

  if (new_num_leds >= led_drivers[driver_num].num_leds) {
    for (int i = (int)NUM_DRIVERS - 1; i > (int)driver_num; i--) {
      memmove(framebuffer + new_offsets[i], led_drivers[i].framebuffer, led_drivers[i].num_leds);
    }
  } else {
    for (uint8_t i = driver_num + 1; i < NUM_DRIVERS; i++) {
      memmove(framebuffer + new_offsets[i], led_drivers[i].framebuffer, led_drivers[i].num_leds);
    }
  }

  led_drivers[driver_num].num_leds = new_num_leds;
  for (uint8_t i = driver_num; i < NUM_DRIVERS; i++) {
    led_drivers[i].framebuffer = framebuffer + new_offsets[i];
  }

  memset(led_drivers[driver_num].framebuffer + old_total,  0, new_total - old_total);
}