#include "rpi.h"
#include "main.h"
#include "stm32g4xx_hal_def.h"
#include "stm32g4xx_hal_gpio.h"
#include "stm32g4xx_hal_i2c.h"
#include <stdint.h>

uint8_t i2c_rx_buf[I2C_RX_BUF_SIZE];
uint8_t i2c_tx_buf[I2C_TX_BUF_MAX_SIZE];
uint8_t i2c_tx_len;
RPiI2CState rpi_i2c_state;

static I2C_HandleTypeDef *bus;
static uint8_t current_command;

void rpi_init(I2C_HandleTypeDef *handle) {
  bus = handle;
  rpi_i2c_state = RPI_I2C_IDLE;
  if (HAL_I2C_EnableListen_IT(bus) != HAL_OK) {
    Error_Handler();
  }

  i2c_tx_len = I2C_TX_BUF_MAX_SIZE;
  for (int i = 0; i < I2C_TX_BUF_MAX_SIZE; i++) {
    i2c_tx_buf[i] = i;
  }
}

void HAL_I2C_SlaveTxCpltCallback(I2C_HandleTypeDef *handle) {
  switch (rpi_i2c_state) {
  case RPI_I2C_SENDING_DATA_SIZE:
    if (HAL_I2C_Slave_Seq_Transmit_IT(bus, i2c_tx_buf, i2c_tx_len,
                                      I2C_LAST_FRAME) != HAL_OK) {
      Error_Handler();
    }
    rpi_i2c_state = RPI_I2C_SENDING_DATA;
    break;
  case RPI_I2C_SENDING_DATA:
    rpi_i2c_state = RPI_I2C_IDLE;
    // if (HAL_I2C_EnableListen_IT(bus) != HAL_OK) {
    //   Error_Handler();
    // }
    break;
  default:
    // TODO: How do we handle this?
    return;
  }
}

void HAL_I2C_SlaveRxCpltCallback(I2C_HandleTypeDef *handle) {
  if (rpi_i2c_state != RPI_I2C_WAITING_FOR_COMMAND) {
    // TODO: How do we handle this?
    return;
  }

  // TODO: Check comm byte
  current_command = i2c_rx_buf[0];
  rpi_i2c_state = RPI_I2C_WAITING_FOR_MASTER_RECEIVE;
}

void HAL_I2C_AddrCallback(I2C_HandleTypeDef *handle, uint8_t TransferDirection,
                          uint16_t AddrMatchCode) {
  if (handle != bus) {
    HAL_GPIO_WritePin(LED_SER_GPIO_Port, LED_SER_Pin, GPIO_PIN_RESET);
    return;
  }

  // TransferDirection is from master POV, so we need to receive when the
  // direction is transmit
  if (TransferDirection == I2C_DIRECTION_TRANSMIT) {
    if (rpi_i2c_state != RPI_I2C_IDLE) {
      // TODO: How do we handle this?
      return;
    }

    if (HAL_I2C_Slave_Seq_Receive_IT(bus, i2c_rx_buf, I2C_RX_BUF_SIZE,
                                     I2C_FIRST_FRAME) != HAL_OK) {
      Error_Handler();
    }
    rpi_i2c_state = RPI_I2C_WAITING_FOR_COMMAND;
  } else {
    if (rpi_i2c_state != RPI_I2C_WAITING_FOR_MASTER_RECEIVE) {
      // TODO: How do we handle this?
      return;
    }

    if (HAL_I2C_Slave_Seq_Transmit_IT(bus, &i2c_tx_len, 1, I2C_NEXT_FRAME) !=
        HAL_OK) {
      Error_Handler();
    }
    rpi_i2c_state = RPI_I2C_SENDING_DATA_SIZE;
  }
}

void HAL_I2C_ListenCpltCallback(I2C_HandleTypeDef *handle) {
  if (handle != bus) {
    return;
  }

  if (rpi_i2c_state == RPI_I2C_IDLE) {
    HAL_I2C_EnableListen_IT(bus);
  }
}