i2c.rs

use embassy_sync::channel::{Channel, Sender, Receiver};
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embassy_time::{Duration, Ticker};
use core::sync::atomic::{AtomicUsize, Ordering};
use embedded_hal_async::i2c::{AddressMode, Operation};

// Request ID generator
static REQUEST_ID_COUNTER: AtomicUsize = AtomicUsize::new(0);

fn generate_request_id() -> usize {
    REQUEST_ID_COUNTER.fetch_add(1, Ordering::SeqCst)
}

// Error types for STM32
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum Error {
    Bus,
    Arbitration,
    Nack,
    Timeout,
    Crc,
    Overrun,
    ZeroLengthTransfer,
}

// Message types
enum I2CRequest {
    Write {
        req_id: usize,
        address: u16,
        data: Vec<u8>,
    },
    Read {
        req_id: usize,
        address: u16,
        read_length: usize,
    },
    WriteRead {
        req_id: usize,
        address: u16,
        write: Vec<u8>,
        read_length: usize,
    },
    Transaction {
        req_id: usize,
        address: u16,
        operations: Vec<Operation<'static>>,
    },
}

enum I2CResponse {
    WriteResult {
        req_id: usize,
        result: Result<(), Error>,
    },
    ReadResult {
        req_id: usize,
        result: Result<Vec<u8>, Error>,
    },
    WriteReadResult {
        req_id: usize,
        result: Result<Vec<u8>, Error>,
    },
    TransactionResult {
        req_id: usize,
        result: Result<(), Error>,
    },
}

// Channel capacity
const CHANNEL_SIZE: usize = 10;

// Shared channels
static I2C_REQUEST_CHANNEL: Channel<CriticalSectionRawMutex, (usize, I2CRequest), CHANNEL_SIZE> = Channel::new();
static I2C_RESPONSE_CHANNEL: Channel<CriticalSectionRawMutex, (usize, I2CResponse), CHANNEL_SIZE> = Channel::new();

// For STM32, define the specific task for I2C1 peripheral
#[embassy_executor::task]
pub async fn i2c1_service_task(mut i2c: embassy_stm32::peripherals::I2C1) {
    let request_receiver = I2C_REQUEST_CHANNEL.receiver();
    let response_sender = I2C_RESPONSE_CHANNEL.sender();
    
    loop {
        if let Some((client_id, request)) = request_receiver.receive().await {
            match request {
                I2CRequest::Write { req_id, address, data } => {
                    let result = i2c.write(address, &data).await;
                    let response = I2CResponse::WriteResult { req_id, result };
                    response_sender.send((client_id, response)).await;
                },
                
                I2CRequest::Read { req_id, address, read_length } => {
                    let mut buffer = vec![0; read_length];
                    let result = i2c.read(address, &mut buffer).await
                        .map(|_| buffer);
                    
                    let response = I2CResponse::ReadResult { req_id, result };
                    response_sender.send((client_id, response)).await;
                },
                
                I2CRequest::WriteRead { req_id, address, write, read_length } => {
                    let mut buffer = vec![0; read_length];
                    let result = i2c.write_read(address, &write, &mut buffer).await
                        .map(|_| buffer);
                    
                    let response = I2CResponse::WriteReadResult { req_id, result };
                    response_sender.send((client_id, response)).await;
                },
                
                I2CRequest::Transaction { req_id, address, operations } => {
                    let mut result = Ok(());
                    
                    for op in &operations {
                        match op {
                            Operation::Read(buf) => {
                                let mut read_buf = vec![0u8; buf.len()];
                                match i2c.read(address, &mut read_buf).await {
                                    Ok(_) => {},
                                    Err(e) => {
                                        result = Err(e);
                                        break;
                                    }
                                }
                            },
                            Operation::Write(buf) => {
                                match i2c.write(address, buf).await {
                                    Ok(_) => {},
                                    Err(e) => {
                                        result = Err(e);
                                        break;
                                    }
                                }
                            },
                        }
                    }
                    
                    let response = I2CResponse::TransactionResult { req_id, result };
                    response_sender.send((client_id, response)).await;
                },
            }
        }
    }
}

// The I2C client that mimics embassy-stm32's I2C
pub struct I2c<'d> {
    client_id: usize,
    request_sender: Sender<'static, CriticalSectionRawMutex, (usize, I2CRequest), CHANNEL_SIZE>,
    response_receiver: Receiver<'static, CriticalSectionRawMutex, (usize, I2CResponse), CHANNEL_SIZE>,
    _phantom: core::marker::PhantomData<&'d ()>,
}

impl<'d> I2c<'d> {
    pub fn new() -> Self {
        static CLIENT_ID_COUNTER: AtomicUsize = AtomicUsize::new(0);
        let client_id = CLIENT_ID_COUNTER.fetch_add(1, Ordering::SeqCst);

        Self {
            client_id,
            request_sender: I2C_REQUEST_CHANNEL.sender(),
            response_receiver: I2C_RESPONSE_CHANNEL.receiver(),
            _phantom: core::marker::PhantomData,
        }
    }

    // Wait for our specific client response
    async fn wait_for_response(&self) -> I2CResponse {
        loop {
            let (resp_client_id, response) = self.response_receiver.receive().await;
            if resp_client_id == self.client_id {
                return response;
            }
        }
    }
    
    // Matching embassy-stm32 I2C API
    pub async fn read(&mut self, address: u8, buffer: &mut [u8]) -> Result<(), Error> {
        let req_id = generate_request_id();
        
        if buffer.is_empty() {
            return Err(Error::ZeroLengthTransfer);
        }
        
        let request = I2CRequest::Read {
            req_id,
            address: address as u16,
            read_length: buffer.len(),
        };
        
        self.request_sender.send((self.client_id, request)).await;
        
        loop {
            match self.wait_for_response().await {
                I2CResponse::ReadResult { req_id: resp_id, result } if resp_id == req_id => {
                    match result {
                        Ok(data) => {
                            let len = data.len().min(buffer.len());
                            buffer[..len].copy_from_slice(&data[..len]);
                            return Ok(());
                        },
                        Err(e) => return Err(e),
                    }
                },
                _ => continue, // Wait for the right response
            }
        }
    }
    
    pub async fn write(&mut self, address: u8, data: &[u8]) -> Result<(), Error> {
        let req_id = generate_request_id();
        
        if data.is_empty() {
            return Err(Error::ZeroLengthTransfer);
        }
        
        let request = I2CRequest::Write {
            req_id,
            address: address as u16,
            data: data.to_vec(),
        };
        
        self.request_sender.send((self.client_id, request)).await;
        
        loop {
            match self.wait_for_response().await {
                I2CResponse::WriteResult { req_id: resp_id, result } if resp_id == req_id => {
                    return result;
                },
                _ => continue,
            }
        }
    }
    
    pub async fn write_read(&mut self, address: u8, write: &[u8], read: &mut [u8]) -> Result<(), Error> {
        let req_id = generate_request_id();
        
        if write.is_empty() || read.is_empty() {
            return Err(Error::ZeroLengthTransfer);
        }
        
        let request = I2CRequest::WriteRead {
            req_id,
            address: address as u16,
            write: write.to_vec(),
            read_length: read.len(),
        };
        
        self.request_sender.send((self.client_id, request)).await;
        
        loop {
            match self.wait_for_response().await {
                I2CResponse::WriteReadResult { req_id: resp_id, result } if resp_id == req_id => {
                    match result {
                        Ok(data) => {
                            let len = data.len().min(read.len());
                            read[..len].copy_from_slice(&data[..len]);
                            return Ok(());
                        },
                        Err(e) => return Err(e),
                    }
                },
                _ => continue,
            }
        }
    }
}

// Implement the embedded-hal-async I2c trait
impl<'d> embedded_hal_async::i2c::I2c for I2c<'d> {
    type Error = Error;

    async fn read(&mut self, address: u8, read: &mut [u8]) -> Result<(), Self::Error> {
        self.read(address, read).await
    }

    async fn write(&mut self, address: u8, write: &[u8]) -> Result<(), Self::Error> {
        self.write(address, write).await
    }

    async fn write_read(&mut self, address: u8, write: &[u8], read: &mut [u8]) -> Result<(), Self::Error> {
        self.write_read(address, write, read).await
    }

    async fn transaction(&mut self, address: u8, operations: &mut [Operation<'_>]) -> Result<(), Self::Error> {
        let req_id = generate_request_id();
        
        // Convert operations to owned versions
        let owned_ops: Vec<Operation<'static>> = operations
            .iter()
            .map(|op| match op {
                Operation::Read(_) => {
                    if let Operation::Read(buf) = op {
                        Operation::Read(vec![0; buf.len()])
                    } else {
                        unreachable!()
                    }
                },
                Operation::Write(data) => {
                    Operation::Write(data.to_vec())
                },
            })
            .collect();
        
        let request = I2CRequest::Transaction {
            req_id,
            address: address as u16,
            operations: owned_ops,
        };
        
        self.request_sender.send((self.client_id, request)).await;
        
        loop {
            match self.wait_for_response().await {
                I2CResponse::TransactionResult { req_id: resp_id, result } if resp_id == req_id => {
                    return result;
                },
                _ => continue,
            }
        }
    }
}

// Now we need a separate task function for each I2C peripheral
// Example for I2C2
#[embassy_executor::task]
pub async fn i2c2_service_task(mut i2c: embassy_stm32::peripherals::I2C2) {
    // Same implementation as i2c1_service_task
    let request_receiver = I2C_REQUEST_CHANNEL.receiver();
    let response_sender = I2C_RESPONSE_CHANNEL.sender();
    
    // Identical loop as in i2c1_service_task
    // ...
}

// Example of using the I2C client in a task
#[embassy_executor::task]
pub async fn sensor_task() {
    let mut i2c = I2c::new();
    let mut ticker: Ticker = Ticker::every(Duration::from_millis(10));
    let mut display_cnt: i32 = 0;
    
    const SENSOR_ADDR: u8 = 0x48;
    const SENSOR_REG: u8 = 0x00;
    
    loop {
        let mut buffer = [0u8; 2];
        match i2c.write_read(SENSOR_ADDR, &[SENSOR_REG], &mut buffer).await {
            Ok(_) => {
                display_cnt += 1;
            },
            Err(_) => {
                // Handle error
            }
        }
        
        ticker.next().await;
    }
}