#include "aht20.h"

#include <stdio.h>

#include "ch32fun.h"
#include "i2c_bitbang.h"

// Polynomial: x^8 + x^5 + x^4 + 1 (0x31)
static uint8_t aht20_calc_crc(const uint8_t* data, size_t len) {
  uint8_t crc = 0xFF;

  for (size_t i = 0; i < len; i++) {
    crc ^= data[i];
    for (uint8_t j = 0; j < 8; j++) {
      if (crc & 0x80) {
        crc = (crc << 1) ^ 0x31;
      } else {
        crc = (crc << 1);
      }
    }
  }
  return crc;
}

static int aht20_get_status(uint8_t* status) {
  i2c_start();
  if (i2c_write_byte((AHT20_ADDR << 1) | 0) != 0) {
    i2c_stop();
    return AHT20_ERR_I2C;
  }
  i2c_write_byte(AHT20_CMD_STATUS);
  i2c_stop();

  i2c_start();
  if (i2c_write_byte((AHT20_ADDR << 1) | 1) != 0) {
    i2c_stop();
    return AHT20_ERR_I2C;
  }
  *status = i2c_read_byte(1);  // NACK to end read
  i2c_stop();

  return AHT20_OK;
}

uint8_t aht20_init(void) {
  uint8_t status;

  // "Before reading the temperature and humidity value, get a byte of status
  // word by sending 0x71"
  if (aht20_get_status(&status) != AHT20_OK) {
    return AHT20_ERR_INIT;
  }

  // "If the status word and 0x18 are not equal to 0x18..." (Bit 3 CAL, Bit 4
  // Retain)
  if ((status & STATUS_CAL_BIT) == 0) {
    // cal
    i2c_start();
    if (i2c_write_byte((AHT20_ADDR << 1) | 0) != 0) {
      i2c_stop();
      return AHT20_ERR_INIT;
    }
    i2c_write_byte(AHT20_CMD_CALIBRATE);
    i2c_write_byte(0x08);
    i2c_write_byte(0x00);
    i2c_stop();
    Delay_Ms(10);  // "Wait 10ms to send the 0xAC command"
  }

  return AHT20_OK;
}

static inline void aht20_parse_data(uint8_t* buf, aht20_data* result) {
  // extract 20-bit raw values
  // Humidity: Byte 1, Byte 2, Byte 3 (high 4 bits)
  uint32_t raw_hum = ((uint32_t)buf[1] << 12) | ((uint32_t)buf[2] << 4) |
                     ((uint32_t)buf[3] >> 4);

  // Temperature: Byte 3 (low 4 bits), Byte 4, Byte 5
  uint32_t raw_temp = ((uint32_t)(buf[3] & 0x0F) << 16) |
                      ((uint32_t)buf[4] << 8) | (uint32_t)buf[5];

  // Humidity
  // Target: %RH * 100
  // Factor: 10000 / 2^20  ==  625 / 2^16
  // Shift: >> 16
  // Rounding: + (Divisor / 2) = + 32768
  uint32_t hum = (raw_hum * 625) + 32768;
  result->hum_p_x100 = hum >> 16;

  // Temperature
  // Target: DegC * 100
  // Range Factor: 20000 / 2^20  == 625 / 2^15
  // Shift: >> 15
  // Rounding: + (Divisor / 2) = + 16384
  uint32_t temp = (raw_temp * 625) + 16384;
  result->temp_c_x100 = (int32_t)(temp >> 15) - 5000;
}

uint8_t aht20_read(aht20_data* out_data) {
  i2c_start();
  if (i2c_write_byte((AHT20_ADDR << 1) | 0) != 0) {
    i2c_stop();
    return AHT20_ERR_I2C;
  }
  i2c_write_byte(AHT20_CMD_TRIGGER);
  i2c_write_byte(0x33);
  i2c_write_byte(0x00);
  i2c_stop();

  // wait for measurement (datasheet says 80ms)
  Delay_Ms(80);

  uint8_t buf[7];  // status + 5 data + CRC
  uint8_t retry = 3;

  while (retry--) {
    i2c_start();
    if (i2c_write_byte((AHT20_ADDR << 1) | 1) != 0) {
      i2c_stop();
      return AHT20_ERR_I2C;
    }

    buf[0] = i2c_read_byte(0);  // ACK

    // check busy bit
    if ((buf[0] & STATUS_BUSY_BIT) == 0) {
      for (int i = 1; i < 6; i++) {
        buf[i] = i2c_read_byte(0);  // ACK
      }
      buf[6] = i2c_read_byte(1);  // NACK (last byte)
      i2c_stop();
      break;
    }

    // still busy
    i2c_read_byte(1);  // NACK to end this
    i2c_stop();

    Delay_Ms(10);
    if (retry == 0) return AHT20_ERR_BUSY;
  }

  if (aht20_calc_crc(buf, 6) != buf[6]) {
    return AHT20_ERR_CRC;
  }

  aht20_parse_data(buf, out_data);

  return AHT20_OK;
}
