#include "ina231.h"

#include <app/zedmon/usb.h>
#include <assert.h>
#include <lk/err.h>
#include <dev/i2c.h>
#include <lib/cbuf.h>
#include <lib/console.h>
#include <platform.h>
#include <platform/timer_capture.h>
#include <stdint.h>
#include <stdio.h>
#include <target/gpioconfig.h>

enum {
    INA231_REG_CONFIG           = 0,
    INA231_REG_SHUNT_VOLTAGE    = 1,
    INA231_REG_BUS_VOLTAGE      = 2,
    INA231_REG_POWER            = 3,
    INA231_REG_CURRENT          = 4,
    INA231_REG_CALIBRATION      = 5,
    INA231_REG_MASK_ENABLE      = 6,
    INA231_REG_MASK_ALERT_LIMIT = 7,
};

enum {
    INA231_CONFIG_MODE_POWER_DOWN     = 0 << 0,
    INA231_CONFIG_MODE_TRIG_SHUNT     = 1 << 0,
    INA231_CONFIG_MODE_TRIG_BUS       = 2 << 0,
    INA231_CONFIG_MODE_TRIG_SHUNT_BUS = 3 << 0,
    INA231_CONFIG_MODE_POWER_DOWN2    = 4 << 0,
    INA231_CONFIG_MODE_CONT_SHUNT     = 5 << 0,
    INA231_CONFIG_MODE_CONT_BUS       = 6 << 0,
    INA231_CONFIG_MODE_CONT_SHUNT_BUS = 7 << 0,

    INA231_CONFIG_V_SHUNT_CT_140_US  = 0 << 3,
    INA231_CONFIG_V_SHUNT_CT_204_US  = 1 << 3,
    INA231_CONFIG_V_SHUNT_CT_332_US  = 2 << 3,
    INA231_CONFIG_V_SHUNT_CT_588_US  = 3 << 3,
    INA231_CONFIG_V_SHUNT_CT_1100_US = 4 << 3,
    INA231_CONFIG_V_SHUNT_CT_2116_US = 5 << 3,
    INA231_CONFIG_V_SHUNT_CT_4156_US = 6 << 3,
    INA231_CONFIG_V_SHUNT_CT_8244_US = 7 << 3,

    INA231_CONFIG_V_BUS_CT_140_US  = 0 << 6,
    INA231_CONFIG_V_BUS_CT_204_US  = 1 << 6,
    INA231_CONFIG_V_BUS_CT_332_US  = 2 << 6,
    INA231_CONFIG_V_BUS_CT_588_US  = 3 << 6,
    INA231_CONFIG_V_BUS_CT_1100_US = 4 << 6,
    INA231_CONFIG_V_BUS_CT_2116_US = 5 << 6,
    INA231_CONFIG_V_BUS_CT_4156_US = 6 << 6,
    INA231_CONFIG_V_BUS_CT_8244_US = 7 << 6,

    INA231_CONFIG_AVG_1    = 0 << 9,
    INA231_CONFIG_AVG_4    = 1 << 9,
    INA231_CONFIG_AVG_16   = 2 << 9,
    INA231_CONFIG_AVG_64   = 3 << 9,
    INA231_CONFIG_AVG_128  = 4 << 9,
    INA231_CONFIG_AVG_256  = 5 << 9,
    INA231_CONFIG_AVG_512  = 6 << 9,
    INA231_CONFIG_AVG_1024 = 7 << 9,

    INA231_CONFIG_RST = 1 << 16,
};

enum {
    INA231_ME_LEN  = 1 << 0,  // Alert Latch Enable.
    INA231_ME_APOL = 1 << 1,  // Alert Polarity.
    INA231_ME_OVF  = 1 << 2,  // Math Overflow.
    INA231_ME_CRVF = 1 << 3,  // Conversion Ready Flag.
    INA231_ME_AFF  = 1 << 4, // Alert Function Flag.
    INA231_ME_CNVR = 1 << 10, // Assert Alert on Conversion Ready.
    INA231_ME_POL  = 1 << 11, // Assert Alert on Power Over Limit.
    INA231_ME_BUL  = 1 << 12, // Assert Alert on Bus Under Voltage Limit.
    INA231_ME_BOL  = 1 << 13, // Assert Alert on Bus Over Voltage Limit.
    INA231_ME_SUL  = 1 << 14, // Assert Alert on Shunt Under Voltage Limit.
    INA231_ME_SOL  = 1 << 15, // Assert Alert on Shunt Over Voltage Limit.
};


static const uint8_t INA231_ADDR = 0x40;
static uint32_t ina231_errors;

static status_t ina231_write(uint8_t reg, uint16_t val) {
    uint8_t data[2];
    data[0] = val >> 8;
    data[1] = val & 0xff;

    status_t ret = i2c_write_reg_bytes(INA_BUS, INA231_ADDR, reg, data, 2);
    if (ret != NO_ERROR) {
        return ret;
    }

    return NO_ERROR;
}

static status_t ina231_read(uint8_t reg, uint16_t *val) {
    uint8_t data[2];

    status_t ret = i2c_read_reg_bytes(INA_BUS, INA231_ADDR, reg, data, 2);
    if (ret != NO_ERROR) {
        return ret;
    }

    *val = data[0] << 8 | data[1];
    return NO_ERROR;
}

bool ina231_init() {

    status_t ret = ina231_write(INA231_REG_MASK_ENABLE, INA231_ME_CNVR);
    if (ret != NO_ERROR) {
        return false;
    }

    ret = ina231_write(INA231_REG_CONFIG,
                       INA231_CONFIG_MODE_CONT_SHUNT_BUS
                       | INA231_CONFIG_V_SHUNT_CT_332_US
                       | INA231_CONFIG_V_BUS_CT_332_US
                       | INA231_CONFIG_AVG_1);
    if (ret != NO_ERROR) {
        return false;
    }

    return true;
}


bool ina231_sample(uint16_t *v_shunt, uint16_t *v_bus) {
    status_t ret;
    uint16_t status;
    ret = ina231_read(INA231_REG_MASK_ENABLE, &status);
    if (ret != NO_ERROR) {
        ina231_errors++;
        return false;
    }

    if (!(status & INA231_ME_CRVF)) {
        printf("alert with no crvf!\n");
        return false;
    }

    ret = ina231_read(INA231_REG_SHUNT_VOLTAGE, v_shunt);
    if (ret != NO_ERROR) {
        ina231_errors++;
        return false;
    }

    ret = ina231_read(INA231_REG_BUS_VOLTAGE, v_bus);
    if (ret != NO_ERROR) {
        ina231_errors++;
        return false;
    }
    return true;
}