kcan 0.1.8

use std::collections::BTreeMap;
use std::time::{Duration, Instant};

use crate::bus::{CanBus, CanFrame};
use crate::error::{Error, Result};
use crate::motor::MotorConfig;
use crate::protocol::{
    DirectCommand, MitCommand, MitHelperCommand, MotorFeedback, direct_command_frame,
    mit_command_frame, mit_helper_frame,
};

#[derive(Debug)]
pub struct MotorManager<B> {
    bus: B,
    motors: BTreeMap<u8, MotorConfig>,
    labels: BTreeMap<String, u8>,
    last_feedback: BTreeMap<u8, MotorFeedback>,
}

impl<B: CanBus> MotorManager<B> {
    pub fn new(bus: B, configs: impl IntoIterator<Item = MotorConfig>) -> Result<Self> {
        let mut motors = BTreeMap::new();
        for config in configs {
            config.refresh_period()?;
            let motor_id = config.motor_id;
            if motors.insert(motor_id, config).is_some() {
                return Err(Error::DuplicateMotorId(motor_id));
            }
        }

        Ok(Self {
            bus,
            motors,
            labels: BTreeMap::new(),
            last_feedback: BTreeMap::new(),
        })
    }

    pub fn with_labels(mut self, labels: impl IntoIterator<Item = (String, u8)>) -> Result<Self> {
        for (label, motor_id) in labels {
            if !self.motors.contains_key(&motor_id) {
                return Err(Error::UnknownMotorId(motor_id));
            }
            self.labels.insert(label, motor_id);
        }
        Ok(self)
    }

    pub fn bus(&self) -> &B {
        &self.bus
    }

    pub fn bus_mut(&mut self) -> &mut B {
        &mut self.bus
    }

    pub fn into_inner(self) -> B {
        self.bus
    }

    pub fn motor_ids(&self) -> impl Iterator<Item = u8> + '_ {
        self.motors.keys().copied()
    }

    pub fn config(&self, motor_id: u8) -> Result<&MotorConfig> {
        self.motors
            .get(&motor_id)
            .ok_or(Error::UnknownMotorId(motor_id))
    }

    pub fn id_for_label(&self, label: &str) -> Result<u8> {
        self.labels
            .get(label)
            .copied()
            .ok_or_else(|| Error::UnknownMotorLabel(label.to_owned()))
    }

    pub fn last_feedback(&self, motor_id: u8) -> Option<&MotorFeedback> {
        self.last_feedback.get(&motor_id)
    }

    pub fn send_mit_command(&mut self, motor_id: u8, command: MitCommand) -> Result<()> {
        let config = self.config(motor_id)?;
        let frame = mit_command_frame(config.model, config.motor_id, command)?;
        self.bus.send(&frame)
    }

    pub fn send_direct_command(&mut self, motor_id: u8, command: DirectCommand) -> Result<()> {
        let config = self.config(motor_id)?;
        let frame = direct_command_frame(config.motor_id, command)?;
        self.bus.send(&frame)
    }

    pub fn send_helper(&mut self, motor_id: u8, helper: MitHelperCommand) -> Result<()> {
        let config = self.config(motor_id)?;
        let frame = mit_helper_frame(config.model, config.motor_id, helper)?;
        self.bus.send(&frame)
    }

    pub fn send_keepalive_all(&mut self) -> Result<()> {
        let frames = self
            .motors
            .values()
            .map(|config| mit_command_frame(config.model, config.motor_id, MitCommand::neutral()))
            .collect::<Result<Vec<_>>>()?;

        for frame in frames {
            self.bus.send(&frame)?;
        }

        Ok(())
    }

    pub fn stop_all(&mut self) -> Result<()> {
        let frames = self
            .motors
            .values()
            .flat_map(|config| {
                [
                    mit_command_frame(config.model, config.motor_id, MitCommand::neutral()),
                    mit_helper_frame(config.model, config.motor_id, MitHelperCommand::Disable),
                ]
            })
            .collect::<Result<Vec<_>>>()?;

        for frame in frames {
            self.bus.send(&frame)?;
        }

        Ok(())
    }

    pub fn receive_feedback(&mut self, timeout: Duration) -> Result<Option<(u8, MotorFeedback)>> {
        let deadline = Instant::now() + timeout;

        loop {
            let now = Instant::now();
            if now >= deadline {
                return Ok(None);
            }

            let Some(frame) = self.bus.receive(deadline.saturating_duration_since(now))? else {
                return Ok(None);
            };

            if let Some(parsed) = self.parse_feedback_frame(&frame)? {
                self.last_feedback.insert(parsed.0, parsed.1.clone());
                return Ok(Some(parsed));
            }
        }
    }

    fn parse_feedback_frame(&self, frame: &CanFrame) -> Result<Option<(u8, MotorFeedback)>> {
        for (motor_id, config) in &self.motors {
            if config.feedback_filter().matches(frame) {
                return Ok(Some((
                    *motor_id,
                    config.model.parse_feedback(frame.data())?,
                )));
            }
        }
        Ok(None)
    }
}