extern crate i2cdev;

use i2cdev::core::*;
use i2cdev::linux::{LinuxI2CDevice, LinuxI2CError, LinuxI2CMessage};

pub enum Channel {
    One,
    Two,
    Three,
}

pub struct Ina3221 {
    slave_addr: u16,
    i2c_device: String,
    shunt_resistor: f64,
}

impl Ina3221 {

    pub fn new(slave_address: u16) -> Ina3221 {
        Ina3221 {
            slave_addr: slave_address,
            i2c_device: format!("/dev/i2c-1"),
            shunt_resistor: 1.0,
        }
    }

    fn write_read_transaction(&mut self, register: u8) -> Result<u16, LinuxI2CError> {
        let mut dev = LinuxI2CDevice::new(&self.i2c_device, self.slave_addr)?;

        let binding = [register];
        let mut read_data = [0; 2];
        let mut msgs = [
            LinuxI2CMessage::write(&binding),
            LinuxI2CMessage::read(&mut read_data)
        ];
        dev.transfer(&mut msgs)?;

        let tmp: u16 = (read_data[0] as u16) << 8 | (read_data[1] as u16);
        Ok(tmp)
    }

    pub fn bus_voltage(&mut self, channel: Channel) -> Result<f64, LinuxI2CError> {
        let register: u8 = match channel {
            Channel::One => 2,
            Channel::Two => 4,
            Channel::Three => 6,
        };
        let mut tmp = self.write_read_transaction(register)?;
        tmp = tmp >> 3;
        let res: f64 = (tmp as f64) * 8.0 / 1000.0;
        Ok(res)
    }

    pub fn shunt_voltage(&mut self, channel: Channel) -> Result<f64, LinuxI2CError> {
        let register: u8 = match channel {
            Channel::One => 1,
            Channel::Two => 3,
            Channel::Three => 5,
        };
        let mut tmp = self.write_read_transaction(register)?;
        tmp = tmp >> 3;
        let res: f64 = (tmp as f64) * 40.0 / 1000000.0;
        Ok(res)
    }

    pub fn current(&mut self, channel: Channel) -> Result<f64, LinuxI2CError> {
        let tmp: f64 = self.shunt_voltage(channel)?;
        let res: f64 = tmp * self.shunt_resistor;
        Ok(res)
    }
}