#ifndef _HW_I2C_H
#define _HW_I2C_H

#include <stdint.h>

#include "ch32fun.h"
#include "ch32v20xhw.h"

// config
#define I2C_CLKRATE 400000
#define I2C_PRERATE 2000000
#define I2C_TIMEOUT_MAX 250000

#define I2C_DIRECTION_TX 0
#define I2C_DIRECTION_RX 1

enum I2C_Error { I2C_OK = 0, I2C_ERR_TIMEOUT, I2C_ERR_BUSY, I2C_ERR_NACK };

static inline uint8_t i2c_check_event(uint32_t event_mask) {
  uint32_t status = I2C1->STAR1 | (I2C1->STAR2 << 16);
  return (status & event_mask) == event_mask;
}

static inline uint8_t i2c_wait_event(uint32_t event_mask) {
  int32_t timeout = I2C_TIMEOUT_MAX;
  while (!i2c_check_event(event_mask) && (timeout-- > 0));
  return timeout > 0 ? I2C_OK : I2C_ERR_TIMEOUT;
}

static inline uint8_t i2c_wait_flag(uint32_t flag) {
  int32_t timeout = I2C_TIMEOUT_MAX;
  while (!(I2C1->STAR1 & flag) && (timeout-- > 0));
  return timeout > 0 ? I2C_OK : I2C_ERR_TIMEOUT;
}

static inline uint8_t i2c_start(uint8_t addr, uint8_t direction) {
  // wait until bus is not busy
  int32_t timeout = I2C_TIMEOUT_MAX;
  while ((I2C1->STAR2 & I2C_STAR2_BUSY) && (timeout-- > 0));
  if (timeout <= 0) return I2C_ERR_TIMEOUT;

  // gen START
  I2C1->CTLR1 |= I2C_CTLR1_START;
  if (i2c_wait_event(I2C_EVENT_MASTER_MODE_SELECT) != I2C_OK)
    return I2C_ERR_TIMEOUT;

  // send address
  I2C1->DATAR = (addr << 1) | direction;

  uint32_t event = (direction == I2C_DIRECTION_TX)
                       ? I2C_EVENT_MASTER_TRANSMITTER_MODE_SELECTED
                       : I2C_EVENT_MASTER_RECEIVER_MODE_SELECTED;

  return i2c_wait_event(event);
}

static inline void i2c_stop(void) { I2C1->CTLR1 |= I2C_CTLR1_STOP; }

static inline uint8_t i2c_write_byte(uint8_t data) {
  I2C1->DATAR = data;
  return i2c_wait_flag(I2C_STAR1_TXE);
}

static inline uint8_t i2c_write(uint8_t addr, const uint8_t* data,
                                uint8_t length) {
  uint8_t err;

  if ((err = i2c_start(addr, I2C_DIRECTION_TX)) != I2C_OK) return err;

  for (uint8_t i = 0; i < length; i++) {
    if ((err = i2c_write_byte(data[i])) != I2C_OK) return err;
  }

  if (i2c_wait_event(I2C_EVENT_MASTER_BYTE_TRANSMITTED) != I2C_OK)
    return I2C_ERR_TIMEOUT;

  i2c_stop();
  return I2C_OK;
}

static inline uint8_t i2c_read(uint8_t addr, uint8_t* buffer, uint8_t length) {
  uint8_t err;

  if ((err = i2c_start(addr, I2C_DIRECTION_RX)) != I2C_OK) return err;

  for (uint8_t i = 0; i < length; i++) {
    if (i == length - 1) {
      I2C1->CTLR1 &= ~I2C_CTLR1_ACK;  // NACK last byte
    }

    if (i2c_wait_flag(I2C_STAR1_RXNE) != I2C_OK) {
      I2C1->CTLR1 |= I2C_CTLR1_ACK;
      return I2C_ERR_TIMEOUT;
    }

    buffer[i] = I2C1->DATAR;
  }

  I2C1->CTLR1 |= I2C_CTLR1_ACK;  // re-enable ACK
  i2c_stop();
  return I2C_OK;
}

static inline uint8_t i2c_write_read(uint8_t addr, const uint8_t* tx_data,
                                     uint8_t tx_len, uint8_t* rx_data,
                                     uint8_t rx_len) {
  uint8_t err;

  // Write phase
  if ((err = i2c_start(addr, I2C_DIRECTION_TX)) != I2C_OK) return err;

  for (uint8_t i = 0; i < tx_len; i++) {
    if ((err = i2c_write_byte(tx_data[i])) != I2C_OK) return err;
  }

  if (i2c_wait_event(I2C_EVENT_MASTER_BYTE_TRANSMITTED) != I2C_OK)
    return I2C_ERR_TIMEOUT;

  // repeated START for read phase
  I2C1->CTLR1 |= I2C_CTLR1_START;
  if (i2c_wait_event(I2C_EVENT_MASTER_MODE_SELECT) != I2C_OK)
    return I2C_ERR_TIMEOUT;

  I2C1->DATAR = (addr << 1) | I2C_DIRECTION_RX;
  if (i2c_wait_event(I2C_EVENT_MASTER_RECEIVER_MODE_SELECTED) != I2C_OK)
    return I2C_ERR_TIMEOUT;

  // read phase
  for (uint8_t i = 0; i < rx_len; i++) {
    if (i == rx_len - 1) {
      I2C1->CTLR1 &= ~I2C_CTLR1_ACK;  // NACK last byte
    }

    if (i2c_wait_flag(I2C_STAR1_RXNE) != I2C_OK) {
      I2C1->CTLR1 |= I2C_CTLR1_ACK;
      return I2C_ERR_TIMEOUT;
    }

    rx_data[i] = I2C1->DATAR;
  }

  I2C1->CTLR1 |= I2C_CTLR1_ACK;  // re-enable ACK
  i2c_stop();
  return I2C_OK;
}

static inline void i2c_init(void) {
  uint16_t tempreg;

  RCC->APB2PCENR |= RCC_APB2Periph_GPIOB | RCC_APB2Periph_AFIO;
  RCC->APB1PCENR |= RCC_APB1Periph_I2C1;

  // PB6 (SCL) and PB7 (SDA)
  GPIOB->CFGLR &= ~(0xFF << (4 * 6));
  GPIOB->CFGLR |= ((GPIO_Speed_10MHz | GPIO_CNF_OUT_OD_AF) << (4 * 6)) |
                  ((GPIO_Speed_10MHz | GPIO_CNF_OUT_OD_AF) << (4 * 7));

  // Reset I2C1 to init all regs
  RCC->APB1PRSTR |= RCC_APB1Periph_I2C1;
  RCC->APB1PRSTR &= ~RCC_APB1Periph_I2C1;

  I2C1->CTLR1 |= I2C_CTLR1_SWRST;
  I2C1->CTLR1 &= ~I2C_CTLR1_SWRST;

  // set I2C freq
  tempreg = I2C1->CTLR2;
  tempreg &= ~I2C_CTLR2_FREQ;
  tempreg |= (FUNCONF_SYSTEM_CORE_CLOCK / I2C_PRERATE) & I2C_CTLR2_FREQ;
  I2C1->CTLR2 = tempreg;

  // Fast Mode 400kHz
  tempreg = (FUNCONF_SYSTEM_CORE_CLOCK / (3 * I2C_CLKRATE)) & I2C_CKCFGR_CCR;
  tempreg |= I2C_CKCFGR_FS;
  I2C1->CKCFGR = tempreg;

  // en I2C and ACK
  I2C1->CTLR1 |= I2C_CTLR1_PE | I2C_CTLR1_ACK;
}

#endif  // _HW_I2C_H