1use std::collections::{BTreeMap, HashMap};
2use std::fs::File;
3use std::io::BufWriter;
4use std::path::Path;
5
6use mcap::records::MessageHeader;
7use mcap::write::Writer;
8
9use clankers_core::{RobotError, RobotResult, Timestamp};
10
11pub struct McapWriter {
13 writer: Writer<BufWriter<File>>,
14 channels: HashMap<String, u16>,
15}
16
17impl McapWriter {
18 pub fn create(path: impl AsRef<Path>) -> RobotResult<Self> {
19 let file = File::create(path.as_ref())
20 .map_err(|e| RobotError::Data(format!("create mcap: {e}")))?;
21 let writer =
22 Writer::new(BufWriter::new(file)).map_err(|e| RobotError::Data(e.to_string()))?;
23 Ok(Self {
24 writer,
25 channels: HashMap::new(),
26 })
27 }
28
29 pub fn write_message(
30 &mut self,
31 topic: &str,
32 schema_name: &str,
33 encoding: &str,
34 data: &[u8],
35 log_time: Timestamp,
36 ) -> RobotResult<()> {
37 let channel_id = self.get_or_create_channel(topic, schema_name, encoding, data)?;
38
39 self.writer
40 .write_to_known_channel(
41 &MessageHeader {
42 channel_id,
43 sequence: 0,
44 log_time: log_time.as_nanos(),
45 publish_time: log_time.as_nanos(),
46 },
47 data,
48 )
49 .map_err(|e| RobotError::Data(e.to_string()))?;
50 Ok(())
51 }
52
53 pub fn finish(mut self) -> RobotResult<()> {
54 self.writer
55 .finish()
56 .map_err(|e| RobotError::Data(e.to_string()))?;
57 Ok(())
58 }
59
60 fn get_or_create_channel(
61 &mut self,
62 topic: &str,
63 schema_name: &str,
64 encoding: &str,
65 schema_data: &[u8],
66 ) -> RobotResult<u16> {
67 if let Some(&id) = self.channels.get(topic) {
68 return Ok(id);
69 }
70 let schema_id = self
71 .writer
72 .add_schema(schema_name, encoding, schema_data)
73 .map_err(|e| RobotError::Data(e.to_string()))?;
74 let channel_id = self
75 .writer
76 .add_channel(schema_id, topic, encoding, &BTreeMap::new())
77 .map_err(|e| RobotError::Data(e.to_string()))?;
78 self.channels.insert(topic.to_string(), channel_id);
79 Ok(channel_id)
80 }
81}