etp 0.0.1-alpha

use tracing::debug;

use crate::{
    Etp, EtpOperations,
    error::{EtpError, EtpResult},
    protocol::{ETP_MESSAGE_HEADER_SIZE, ETP_MESSAGE_HEADER_STATUS_INDEX},
    util::{self, u32_vec_to_le_bytes},
};
use std::collections::HashMap;

#[derive(Debug, Clone, PartialEq)]
pub enum GpioMode {
    Input,
    Output,
}

#[derive(Debug, Clone, PartialEq)]
pub enum GpioType {
    PushPull,
    PullUp,
    PullDown,
    OpenDrain,
}

#[derive(Debug, Clone, PartialEq)]
pub enum InterruptMode {
    None,
    RisingEdge,
    FallingEdge,
    BothEdges,
}

pub type Pin = u8;
pub type Port = Option<char>;
pub type GpioInfo = HashMap<Port, Vec<Pin>>;

#[derive(Debug, Clone)]
pub struct GpioPinConfig {
    pub pin_number: Pin,
    pub port: Port,
    pub mode: GpioMode,
    pub pin_type: Option<GpioType>,
    pub interrupt_mode: Option<InterruptMode>,
}

#[derive(Debug, Clone, PartialEq)]
pub enum GpioPinStatus {
    High,
    Low,
}

pub type GpioPinsStatus = HashMap<(Port, Pin), GpioPinStatus>;

#[repr(u8)]
pub(crate) enum GpioOperations {
    GpioInfo = 0x00,
    GpioInit = 0x01,
    GpioRead = 0x02,
    GpioWrite = 0x03,
}

impl TryFrom<u8> for GpioOperations {
    type Error = EtpError;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            0x00 => Ok(GpioOperations::GpioInfo),
            0x01 => Ok(GpioOperations::GpioInit),
            0x02 => Ok(GpioOperations::GpioRead),
            0x03 => Ok(GpioOperations::GpioWrite),
            _ => Err(EtpError::InvalidByteLength),
        }
    }
}

#[repr(u8)]
#[derive(Debug, Clone)]
pub(crate) enum GpioInfoCommands {
    GpioPortCount = 0x01,
    GpioPinCount = 0x02,
    GpioPorts = 0x03,
    GpioPins = 0x04,
}

impl TryFrom<u8> for GpioInfoCommands {
    type Error = EtpError;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            0x01 => Ok(GpioInfoCommands::GpioPortCount),
            0x02 => Ok(GpioInfoCommands::GpioPinCount),
            0x03 => Ok(GpioInfoCommands::GpioPorts),
            0x04 => Ok(GpioInfoCommands::GpioPins),
            _ => Err(EtpError::InvalidByteLength),
        }
    }
}

impl GpioPinConfig {
    pub fn new(
        pin_number: Pin,
        port: Port,
        mode: GpioMode,
        pin_type: Option<GpioType>,
        interrupt_mode: Option<InterruptMode>,
    ) -> EtpResult<Self> {
        Ok(GpioPinConfig {
            pin_number,
            port,
            mode,
            pin_type,
            interrupt_mode,
        })
    }
}

pub struct Gpio {
    pub pins: Vec<GpioPinConfig>,
}

impl Gpio {
    pub(crate) fn new() -> Self {
        Gpio { pins: Vec::new() }
    }
}

impl Etp {
    pub fn gpio_init(&mut self, pins: Vec<GpioPinConfig>) -> EtpResult<()> {
        let mut dir_mask: u32 = 0;
        let mut dir_value: u32 = 0;
        let mut pull_mask: u32 = 0;
        let mut pull_value: u32 = 0;
        let mut interrupt_mask: u32 = 0;
        let mut interrupt_value: u32 = 0;

        for pin in pins.iter() {
            // Set pin direction
            dir_mask |= 1 << pin.pin_number;
            if pin.mode == GpioMode::Input {
                dir_value |= 1 << pin.pin_number;
            }

            // Set pin type
            pull_mask |= 1 << pin.pin_number;
            if let Some(pin_type) = &pin.pin_type {
                match pin_type {
                    GpioType::PushPull => {}
                    GpioType::PullUp => pull_value |= 1 << pin.pin_number,
                    // TODO: Verify the ones below are correct
                    GpioType::PullDown => pull_value |= 2 << pin.pin_number,
                    GpioType::OpenDrain => pull_value |= 3 << pin.pin_number,
                }
            } else {
                // Do nothing
            }

            // Set interrupt mode
            if let Some(interrupt_mode) = &pin.interrupt_mode {
                match interrupt_mode {
                    InterruptMode::None => {}
                    InterruptMode::RisingEdge => interrupt_value |= 1 << pin.pin_number * 2,
                    InterruptMode::FallingEdge => interrupt_mask |= 2 << pin.pin_number * 2,
                    InterruptMode::BothEdges => interrupt_value |= 3 << pin.pin_number * 2,
                }
            }
        }

        let payload = u32_vec_to_le_bytes(&[
            dir_mask,
            dir_value,
            pull_mask,
            pull_value,
            interrupt_mask,
            interrupt_value,
        ]);

        let cmd = self.frame_cmd_packet(EtpOperations::GpioInit, payload);
        self.send(cmd)?;
        let response = self.receive(ETP_MESSAGE_HEADER_SIZE)?;
        let status = response[ETP_MESSAGE_HEADER_STATUS_INDEX];

        if status != 0 {
            return Err(EtpError::InvalidEtpStatusCode(format!(
                "GPIO Init failed with status: {}",
                status
            )));
        }

        self.gpio.pins = pins;

        Ok(())
    }

    pub fn gpio_get_info(&self) -> EtpResult<GpioInfo> {
        // Get GPIO port count
        let cmd = self.frame_cmd_packet(
            EtpOperations::GetGpioInfo,
            vec![GpioInfoCommands::GpioPortCount as u8],
        );
        self.send(cmd)?;
        let response = self.receive(ETP_MESSAGE_HEADER_SIZE + 2)?;
        let port_count = response[ETP_MESSAGE_HEADER_SIZE + 1];
        debug!("Port count: {}", port_count);

        // Get GPIO pin count
        let cmd = self.frame_cmd_packet(
            EtpOperations::GetGpioInfo,
            vec![GpioInfoCommands::GpioPins as u8],
        );
        self.send(cmd)?;
        let response = self.receive(ETP_MESSAGE_HEADER_SIZE + 2)?;
        let pin_info = response[(ETP_MESSAGE_HEADER_SIZE + 1)..].to_vec();
        // Pad with 0s to make it a multiple of 5 [1 byte for port + 4 bytes for pin]
        let pin_info = pin_info
            .chunks(5)
            .map(|chunk| {
                let mut padded_chunk = chunk.to_vec();
                while padded_chunk.len() < 5 {
                    padded_chunk.push(0);
                }
                padded_chunk
            })
            .collect::<Vec<_>>();

        let mut gpio_info: GpioInfo = HashMap::new();

        for chunk in pin_info.iter() {
            let port = chunk[0];
            let port: Port = if port == '_' as u8 {
                None
            } else {
                Some(port as char)
            };
            let pin_mask: u32 = u32::from_le_bytes([chunk[1], chunk[2], chunk[3], chunk[4]]);
            let pins: Vec<Pin> = util::mask_to_bits(pin_mask);

            gpio_info.insert(port, pins);
        }

        Ok(gpio_info)
    }

    pub fn gpio_read(&mut self, port: Port, pins: Vec<Pin>) -> EtpResult<GpioPinsStatus> {
        let mut pin_mask: u32 = 0;
        for pin in pins.iter() {
            pin_mask |= 1 << pin;
        }

        let port_char = match port {
            Some(p) => p as u8,
            None => '_' as u8,
        };
        let pin_mask_bytes = u32_vec_to_le_bytes(&[pin_mask]);
        let payload = [
            port_char,
            pin_mask_bytes[0],
            pin_mask_bytes[1],
            pin_mask_bytes[2],
            pin_mask_bytes[3],
        ];

        let cmd = self.frame_cmd_packet(EtpOperations::GpioRead, payload.to_vec());
        self.send(cmd)?;
        let response = self.receive(ETP_MESSAGE_HEADER_SIZE + 9)?;
        let status = response[ETP_MESSAGE_HEADER_STATUS_INDEX];
        if status != 0 {
            return Err(EtpError::InvalidEtpStatusCode(format!(
                "GPIO Read failed with status: {}",
                status
            )));
        }

        let data = response[ETP_MESSAGE_HEADER_SIZE..].to_vec();
        debug!("GPIO Read data: {:?}", data);

        // First byte is the port, next 4 bytes represent the requested pins, and the last 4 bytes represent
        // status of the requested pins
        let requested_pins_status: [u8; 4] = [data[5], data[6], data[7], data[8]];
        let pins_status = u32::from_le_bytes(requested_pins_status);

        let mut gpio_pins_status: GpioPinsStatus = HashMap::new();
        for pin in pins.iter() {
            let pin_mask = 1 << pin;
            let pin_status = if (pins_status & pin_mask) != 0 {
                GpioPinStatus::High
            } else {
                GpioPinStatus::Low
            };
            gpio_pins_status.insert((port, *pin), pin_status);
        }

        Ok(gpio_pins_status)
    }

    pub fn gpio_write(&mut self, pins_status: GpioPinsStatus) -> EtpResult<()> {
        // Separate each port
        let mut port_pins: HashMap<Port, Vec<(Pin, GpioPinStatus)>> = HashMap::new();
        for ((port, pin), status) in pins_status.iter() {
            port_pins
                .entry(*port)
                .or_insert_with(Vec::new)
                .push((*pin, status.clone()));
        }

        // Create payload for each port
        for (port, pins) in port_pins.iter() {
            let mut pin_mask: u32 = 0;
            let mut pin_value: u32 = 0;
            for (pin, status) in pins.iter() {
                pin_mask |= 1 << pin;
                if *status == GpioPinStatus::High {
                    pin_value |= 1 << pin;
                }
            }

            let port_char = match port {
                Some(p) => *p as u8,
                None => '_' as u8,
            };

            let pin_mask_bytes = u32_vec_to_le_bytes(&[pin_mask]);
            let pin_value_bytes = u32_vec_to_le_bytes(&[pin_value]);
            let payload = [
                port_char,
                pin_mask_bytes[0],
                pin_mask_bytes[1],
                pin_mask_bytes[2],
                pin_mask_bytes[3],
                pin_value_bytes[0],
                pin_value_bytes[1],
                pin_value_bytes[2],
                pin_value_bytes[3],
            ];

            let cmd = self.frame_cmd_packet(EtpOperations::GpioWrite, payload.to_vec());
            self.send(cmd)?;

            let response = self.receive(ETP_MESSAGE_HEADER_SIZE)?;
            let status = response[ETP_MESSAGE_HEADER_STATUS_INDEX];
            if status != 0 {
                return Err(EtpError::InvalidEtpStatusCode(format!(
                    "GPIO Write failed with status: {} for port {:?}",
                    status, port
                )));
            }
        }

        Ok(())
    }
}