1use std::path::PathBuf;
52use std::process::{Command, Stdio};
53use std::time::{Duration, Instant};
54
55use serde::{Deserialize, Serialize};
56
57use crate::adapter::{deserialize, serialize, Codec};
58use crate::subprocess::{reshape_channel_major, write_edf_bytes, SampleDtype, ScratchDir};
59
60const ENVELOPE_MAGIC: &[u8; 4] = b"ECSX";
62const ENVELOPE_HEADER_LEN: usize = 21;
64pub const DEFAULT_TIMEOUT_SECS: u64 = 600;
66
67#[derive(Clone, Copy, Debug, PartialEq, Eq, Default, Serialize, Deserialize)]
69#[serde(rename_all = "lowercase")]
70pub enum InputFormat {
71 #[default]
73 Edf,
74 Ecs0,
76}
77
78#[derive(Clone, Copy, Debug, PartialEq, Eq, Default, Serialize, Deserialize)]
80#[serde(rename_all = "lowercase")]
81pub enum OutputFormat {
82 #[default]
85 Raw,
86 Ecs0,
88}
89
90#[derive(Clone, Debug)]
96pub struct ExternalCodec {
97 name: String,
99 cmd: PathBuf,
102 prefix_args: Vec<String>,
104 encode_args: Vec<String>,
106 decode_args: Vec<String>,
108 env: Vec<(String, String)>,
110 declared_lossless: bool,
112 sample_dtype: SampleDtype,
114 input_format: InputFormat,
116 output_format: OutputFormat,
118 timeout: Duration,
120}
121
122pub fn default_encode_args() -> Vec<String> {
124 ["encode", "{input}", "{output}"]
125 .iter()
126 .map(|s| s.to_string())
127 .collect()
128}
129
130pub fn default_decode_args() -> Vec<String> {
134 [
135 "decode",
136 "{input}",
137 "{output}",
138 "--channels",
139 "{channels}",
140 "--samples",
141 "{samples}",
142 "--rate",
143 "{rate}",
144 "--dtype",
145 "{dtype}",
146 ]
147 .iter()
148 .map(|s| s.to_string())
149 .collect()
150}
151
152impl ExternalCodec {
153 pub fn new(name: impl Into<String>, cmd: impl Into<PathBuf>) -> Self {
155 Self {
156 name: name.into(),
157 cmd: cmd.into(),
158 prefix_args: Vec::new(),
159 encode_args: default_encode_args(),
160 decode_args: default_decode_args(),
161 env: Vec::new(),
162 declared_lossless: false,
163 sample_dtype: SampleDtype::default(),
164 input_format: InputFormat::default(),
165 output_format: OutputFormat::default(),
166 timeout: Duration::from_secs(DEFAULT_TIMEOUT_SECS),
167 }
168 }
169
170 pub fn with_prefix_args(mut self, args: Vec<String>) -> Self {
172 self.prefix_args = args;
173 self
174 }
175
176 pub fn with_templates(mut self, encode: Vec<String>, decode: Vec<String>) -> Self {
178 self.encode_args = encode;
179 self.decode_args = decode;
180 self
181 }
182
183 pub fn with_env(mut self, env: Vec<(String, String)>) -> Self {
185 self.env = env;
186 self
187 }
188
189 pub fn with_declared_lossless(mut self, v: bool) -> Self {
191 self.declared_lossless = v;
192 self
193 }
194
195 pub fn with_sample_dtype(mut self, d: SampleDtype) -> Self {
197 self.sample_dtype = d;
198 self
199 }
200
201 pub fn with_formats(mut self, input: InputFormat, output: OutputFormat) -> Self {
203 self.input_format = input;
204 self.output_format = output;
205 self
206 }
207
208 pub fn with_timeout(mut self, timeout: Duration) -> Self {
210 self.timeout = timeout;
211 self
212 }
213
214 pub fn cmd(&self) -> &PathBuf {
216 &self.cmd
217 }
218}
219
220fn dtype_code(d: SampleDtype) -> u8 {
222 match d {
223 SampleDtype::I16 => 0,
224 SampleDtype::I32 => 1,
225 SampleDtype::I64 => 2,
226 }
227}
228
229fn code_dtype(c: u8) -> Option<SampleDtype> {
231 match c {
232 0 => Some(SampleDtype::I16),
233 1 => Some(SampleDtype::I32),
234 2 => Some(SampleDtype::I64),
235 _ => None,
236 }
237}
238
239fn wrap_envelope(
241 payload: &[u8],
242 n_chan: u32,
243 n_samples: u32,
244 dtype: SampleDtype,
245 fs: f64,
246) -> Vec<u8> {
247 let mut out = Vec::with_capacity(ENVELOPE_HEADER_LEN + payload.len());
248 out.extend_from_slice(ENVELOPE_MAGIC);
249 out.extend_from_slice(&n_chan.to_le_bytes());
250 out.extend_from_slice(&n_samples.to_le_bytes());
251 out.push(dtype_code(dtype));
252 out.extend_from_slice(&fs.to_le_bytes());
253 out.extend_from_slice(payload);
254 out
255}
256
257fn parse_envelope(blob: &[u8]) -> Option<(usize, usize, SampleDtype, f64, &[u8])> {
260 if blob.len() < ENVELOPE_HEADER_LEN || &blob[0..4] != ENVELOPE_MAGIC {
261 return None;
262 }
263 let n_chan = u32::from_le_bytes([blob[4], blob[5], blob[6], blob[7]]) as usize;
264 let n_samples = u32::from_le_bytes([blob[8], blob[9], blob[10], blob[11]]) as usize;
265 let dtype = code_dtype(blob[12])?;
266 let fs = f64::from_le_bytes([
267 blob[13], blob[14], blob[15], blob[16], blob[17], blob[18], blob[19], blob[20],
268 ]);
269 Some((n_chan, n_samples, dtype, fs, &blob[ENVELOPE_HEADER_LEN..]))
270}
271
272fn substitute(args: &[String], map: &[(&str, String)]) -> Vec<String> {
274 args.iter()
275 .map(|a| {
276 let mut s = a.clone();
277 for (k, v) in map {
278 if s.contains(k) {
279 s = s.replace(k, v);
280 }
281 }
282 s
283 })
284 .collect()
285}
286
287impl ExternalCodec {
288 fn subst_map(
290 &self,
291 input: &str,
292 output: &str,
293 n_chan: usize,
294 n_samples: usize,
295 fs: f64,
296 ) -> Vec<(&'static str, String)> {
297 vec![
298 ("{input}", input.to_string()),
299 ("{output}", output.to_string()),
300 ("{channels}", n_chan.to_string()),
301 ("{samples}", n_samples.to_string()),
302 ("{rate}", format!("{fs}")),
303 ("{dtype}", self.sample_dtype.as_token().to_string()),
304 ]
305 }
306
307 fn run(
312 &self,
313 args: &[String],
314 n_chan: usize,
315 n_samples: usize,
316 fs: f64,
317 workdir: &std::path::Path,
318 ) -> bool {
319 let mut cmd = Command::new(&self.cmd);
320 cmd.args(&self.prefix_args)
321 .args(args)
322 .stdin(Stdio::null())
323 .stdout(Stdio::null())
324 .stderr(Stdio::null())
325 .env("ECS_CHANNELS", n_chan.to_string())
328 .env("ECS_SAMPLES", n_samples.to_string())
329 .env("ECS_RATE", format!("{fs}"))
330 .env("ECS_DTYPE", self.sample_dtype.as_token())
331 .env("ECS_WORKDIR", workdir);
332 for (k, v) in &self.env {
333 cmd.env(k, v);
334 }
335
336 let mut child = match cmd.spawn() {
337 Ok(c) => c,
338 Err(_) => return false,
339 };
340 let start = Instant::now();
341 loop {
342 match child.try_wait() {
343 Ok(Some(status)) => return status.success(),
344 Ok(None) => {
345 if start.elapsed() >= self.timeout {
346 let _ = child.kill();
347 let _ = child.wait();
348 return false;
349 }
350 std::thread::sleep(Duration::from_millis(5));
351 }
352 Err(_) => {
353 let _ = child.kill();
354 let _ = child.wait();
355 return false;
356 }
357 }
358 }
359 }
360
361 fn try_encode(&self, signal: &[Vec<i64>], fs: f64) -> Vec<u8> {
364 let n_chan = signal.len();
365 if n_chan == 0 || n_chan > u32::MAX as usize {
366 return Vec::new();
367 }
368 let n_samples = signal[0].len();
373 let uniform = signal.iter().all(|c| c.len() == n_samples);
374 if self.output_format == OutputFormat::Raw && (!uniform || n_samples == 0) {
375 return Vec::new();
376 }
377 if n_samples > u32::MAX as usize {
378 return Vec::new();
379 }
380
381 let input_bytes = match self.input_format {
383 InputFormat::Edf => match write_edf_bytes(signal, fs) {
384 Some(b) => b,
385 None => return Vec::new(),
386 },
387 InputFormat::Ecs0 => serialize(signal),
388 };
389
390 let dir = match ScratchDir::new("ext_enc") {
391 Ok(d) => d,
392 Err(_) => return Vec::new(),
393 };
394 let in_name = match self.input_format {
395 InputFormat::Edf => "in.edf",
396 InputFormat::Ecs0 => "in.ecs0",
397 };
398 let in_path = dir.join(in_name);
399 let out_path = dir.join("blob.bin");
400 if std::fs::write(&in_path, &input_bytes).is_err() {
401 return Vec::new();
402 }
403
404 let map = self.subst_map(
405 &in_path.to_string_lossy(),
406 &out_path.to_string_lossy(),
407 n_chan,
408 n_samples,
409 fs,
410 );
411 let args = substitute(&self.encode_args, &map);
412 if !self.run(&args, n_chan, n_samples, fs, &dir.path) {
413 return Vec::new();
414 }
415
416 let payload = match std::fs::read(&out_path) {
417 Ok(p) => p,
418 Err(_) => return Vec::new(),
419 };
420 wrap_envelope(&payload, n_chan as u32, n_samples as u32, self.sample_dtype, fs)
421 }
422
423 fn try_decode(&self, blob: &[u8]) -> Vec<Vec<i64>> {
426 let (n_chan, n_samples, dtype, fs, payload) = match parse_envelope(blob) {
427 Some(parts) => parts,
428 None => return Vec::new(),
429 };
430
431 let dir = match ScratchDir::new("ext_dec") {
432 Ok(d) => d,
433 Err(_) => return Vec::new(),
434 };
435 let in_path = dir.join("blob.bin");
436 let out_name = match self.output_format {
437 OutputFormat::Raw => "out.raw",
438 OutputFormat::Ecs0 => "out.ecs0",
439 };
440 let out_path = dir.join(out_name);
441 if std::fs::write(&in_path, payload).is_err() {
442 return Vec::new();
443 }
444
445 let map = self.subst_map(
446 &in_path.to_string_lossy(),
447 &out_path.to_string_lossy(),
448 n_chan,
449 n_samples,
450 fs,
451 );
452 let args = substitute(&self.decode_args, &map);
453 if !self.run(&args, n_chan, n_samples, fs, &dir.path) {
454 return Vec::new();
455 }
456
457 let out = match std::fs::read(&out_path) {
458 Ok(o) => o,
459 Err(_) => return Vec::new(),
460 };
461 match self.output_format {
462 OutputFormat::Raw => {
463 let _ = dtype; reshape_channel_major(&out, n_chan, n_samples, self.sample_dtype).unwrap_or_default()
467 }
468 OutputFormat::Ecs0 => deserialize(&out),
469 }
470 }
471}
472
473impl Codec for ExternalCodec {
474 fn name(&self) -> &str {
475 &self.name
476 }
477
478 fn declared_lossless(&self) -> bool {
479 self.declared_lossless
480 }
481
482 fn encode(&self, signal: &[Vec<i64>], fs: f64) -> Vec<u8> {
483 let rate = if fs.is_finite() && fs > 0.0 { fs } else { 1.0 };
484 self.try_encode(signal, rate)
485 }
486
487 fn decode(&self, blob: &[u8]) -> Vec<Vec<i64>> {
488 self.try_decode(blob)
489 }
490}
491
492pub fn resolve_cmd(name: &str, cmd: &str) -> Option<PathBuf> {
499 let envvar = format!("ECS_CODEC_{}_BIN", name.to_uppercase().replace('-', "_"));
500 if let Some(p) = std::env::var_os(&envvar) {
501 let pb = PathBuf::from(p);
502 if pb.is_file() {
503 return Some(pb);
504 }
505 }
506 let pb = PathBuf::from(cmd);
507 if cmd.contains('/') || cmd.contains('\\') {
508 if pb.is_file() {
509 return Some(pb);
510 }
511 return None;
512 }
513 Some(pb)
515}
516
517#[cfg(test)]
518mod tests {
519 use super::*;
520
521 #[test]
522 fn envelope_round_trips() {
523 let payload = b"arbitrary codec bytes";
524 let blob = wrap_envelope(payload, 4, 256, SampleDtype::I32, 250.0);
525 let (n_chan, n_samples, dtype, fs, back) =
526 parse_envelope(&blob).expect("valid envelope parses");
527 assert_eq!(n_chan, 4);
528 assert_eq!(n_samples, 256);
529 assert_eq!(dtype, SampleDtype::I32);
530 assert_eq!(fs, 250.0);
531 assert_eq!(back, payload);
532 }
533
534 #[test]
535 fn parse_envelope_rejects_bad_input() {
536 assert!(parse_envelope(b"").is_none());
537 assert!(parse_envelope(b"XXXX too short").is_none());
538 let mut bad = vec![0u8; ENVELOPE_HEADER_LEN + 4];
540 bad[0] = b'N';
541 assert!(parse_envelope(&bad).is_none());
542 let mut bad2 = wrap_envelope(b"x", 1, 1, SampleDtype::I16, 1.0);
544 bad2[12] = 3;
545 assert!(parse_envelope(&bad2).is_none());
546 }
547
548 #[test]
549 fn decode_of_garbage_is_empty() {
550 let codec = ExternalCodec::new("x", "/nonexistent/codec");
552 assert!(codec.decode(b"").is_empty());
553 assert!(codec.decode(b"not an envelope").is_empty());
554 }
555
556 #[test]
557 fn empty_signal_encodes_empty() {
558 let codec = ExternalCodec::new("x", "/nonexistent/codec");
561 assert!(codec.encode(&[], 256.0).is_empty());
562 }
563
564 #[test]
565 fn name_and_claim_are_pure() {
566 let codec = ExternalCodec::new("mycodec", "python3").with_declared_lossless(true);
567 assert_eq!(codec.name(), "mycodec");
568 assert!(codec.declared_lossless());
569 }
570
571 #[test]
572 fn default_templates_carry_placeholders() {
573 let enc = default_encode_args();
574 assert_eq!(enc[0], "encode");
575 assert!(enc.iter().any(|a| a == "{input}"));
576 assert!(enc.iter().any(|a| a == "{output}"));
577 let dec = default_decode_args();
578 assert_eq!(dec[0], "decode");
579 for tok in ["{input}", "{output}", "{channels}", "{samples}", "{rate}", "{dtype}"] {
580 assert!(dec.iter().any(|a| a == tok), "decode template missing {tok}");
581 }
582 }
583
584 #[test]
585 fn substitute_replaces_tokens() {
586 let tmpl = default_decode_args();
587 let map = vec![
588 ("{input}", "/tmp/a".to_string()),
589 ("{output}", "/tmp/b".to_string()),
590 ("{channels}", "3".to_string()),
591 ("{samples}", "8".to_string()),
592 ("{rate}", "256".to_string()),
593 ("{dtype}", "i32".to_string()),
594 ];
595 let out = substitute(&tmpl, &map);
596 assert!(out.contains(&"/tmp/a".to_string()));
597 assert!(out.contains(&"3".to_string()));
598 assert!(out.contains(&"i32".to_string()));
599 assert!(!out.iter().any(|a| a.contains('{')), "no placeholder left");
600 }
601
602 #[test]
603 fn resolve_cmd_bare_name_is_trusted() {
604 assert_eq!(resolve_cmd("c", "python3"), Some(PathBuf::from("python3")));
606 assert!(resolve_cmd("c", "/definitely/not/here/codec").is_none());
608 }
609}