sinusoidal 0.10.0

//! The Sinusoidal Framework Rust SDK
//!
//! Apps are defined by implementing the [App] trait and run through a [Sinusoidal] object.
//!
//! Example
//! ```no_run
//! use serde::Deserialize;
//! use sinusoidal::proto::streaming_api::Datagram;
//! use sinusoidal::stream_info::InputStream;
//! use sinusoidal::*;
//! use std::io::{Error, ErrorKind, Write};
//! use std::collections::HashMap;
//! use std::fs;
//!
//! struct MyApp {
//!   settings: MySettings,
//!   out_stream: OutStream,
//!   big_count: i64
//! }
//!
//! impl MyApp {
//!   fn drop(self: Self) -> () {
//!     let mut f = open_persistent_file("big_count").unwrap();
//!     f.write(format!("{}", self.big_count).as_bytes()).unwrap();
//!   }
//! }
//!
//! #[derive(Debug, Clone, Deserialize)]
//! struct MySettings {
//!   #[serde(default)]
//!   input_streams: Vec<InputStream>,
//! }
//!
//! impl App for MyApp {
//!   type Settings = MySettings;
//!   type AppEvent = (); //In this example we don't introduce any custom events
//!
//!   fn new(settings: MySettings, connection: &mut Connection<Self>) -> Result<Self, Error> {
//!     let mut big = 0;
//!     match open_persistent_file("big_count") {
//!       Ok(f)  => {
//!         let contents = String::new();
//!         fs::read_to_string(&contents).unwrap();
//!         big = contents.trim().parse::<i64>().map_err(|e| {Error::new(ErrorKind::InvalidData, e)}).unwrap()
//!         },
//!       Err(_) => (),
//!     }
//!     for stream in settings.input_streams.iter() {
//!       connection.connect_to_stream(stream)?;
//!     }
//!     let out_stream = connection.register_out_stream("Output")?;
//!     Ok(MyApp{ settings, out_stream, big_count: 0 })
//!   }
//!
//!   fn handle_packet(&mut self, _connection: &mut Connection<Self>, id: &StreamId, d: Datagram) -> Result<(), Error> {
//!     println!("StreamId: {id}; Data: {d:?}");
//!     let mut v = 0_i64;
//!     if d.values[0] > 1000 {
//!       v = 1;
//!       self.big_count += 1;
//!     }
//!     self.out_stream
//!       .send(Datagram{ timestamp : d.timestamp
//!                     , sample_count: d.sample_count
//!                     , values: vec![v]
//!                     , clock_synch: d.clock_synch
//!                     , flags: d.flags
//!                     , gm_identity: d.gm_identity })
//!   }
//! }
//!
//! pub fn main() -> Result<(), std::io::Error> {
//!   Sinusoidal::<MyApp>::main()
//! }
//! ```

/// Protobuf definitions
#[macro_use]
pub mod proto;
use crate::proto::control_api::*;
use crate::proto::streaming_api::*;
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};

/// Custom binary encoding/decoding for datagrams
pub mod sv;

pub mod goose;
use goose::GooseEvent;

pub mod stream_info;
pub use stream_info::{InputStream, StreamId, StreamInfo, SysStream};

pub mod utils;

use prost::Message;
use serde::Deserialize;
use serde::de::DeserializeOwned;
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::io::prelude::*;
use std::io::{Error, ErrorKind};
use std::os::unix::net::UnixStream;
use std::str::FromStr;
use std::sync::mpsc;
use std::*;

/// Callbacks implemented by an app.
pub trait App: Sized {
  type Settings: DeserializeOwned;
  type AppEvent: Send + 'static;

  /// Create the app and perform initial connection setup (e.g. subscribing to
  /// streams, registering output streams).
  fn new(settings: Self::Settings, connection: &mut Connection<Self>) -> Result<Self, Error>;

  /// Called when the app receives a packet from a stream-id it has `connect_to_stream`'d to.
  fn handle_packet(
    &mut self,
    _connection: &mut Connection<Self>,
    _stream_id: &StreamId,
    _packet: Datagram,
  ) -> Result<(), Error> {
    Ok(())
  }

  /// Called when app settings are updated. Return Ok if the new settings can be accomodated
  /// without a restart and Err to restart the app with the new settings.
  fn handle_new_settings(
    &mut self,
    _connection: &mut Connection<Self>,
    _settings: Self::Settings,
  ) -> Result<(), Error> {
    Err(Error::other("Unhandled new settings"))
  }

  /// Called when PTP status is updated. Note that in order to get PTP events you need to be
  /// connected to the `StreamId::SYS_TIME` stream.
  fn handle_ptp_event(
    &mut self,
    _connection: &mut Connection<Self>,
    _ptp_event: PtpEvent,
  ) -> Result<(), Error> {
    Ok(())
  }

  /// Called when Goose status is updated for a GOOSE stream that has been connected to.
  fn handle_goose_packet(
    &mut self,
    _connection: &mut Connection<Self>,
    _goose_event: GooseEvent,
  ) -> Result<(), Error> {
    Ok(())
  }

  fn handle_event(
    &mut self,
    _connection: &mut Connection<Self>,
    _event: Self::AppEvent,
  ) -> Result<(), Error> {
    Ok(())
  }
}

enum SDKMessage<AppEvent> {
  Datagrams(StreamId, Vec<Datagram>),
  SystemEvent(SystemEvent),
  PtpEvent(PtpEvent),
  AppEvent(AppEvent),
  Goose(GooseEvent),
  Trigger(UnqualifiedTriggerName, TriggerPayload),
}

/// Qualified trigger id (i.e. `Trigger:{app_name}:{trigger_name}`)
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct TriggerId(pub String);

/// Unqualified trigger name
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct UnqualifiedTriggerName(pub String);

type TriggerPayload = Value;
type TriggerHandler = Box<dyn FnMut(&TriggerPayload) -> Result<(), Error> + Send>;

impl UnqualifiedTriggerName {
  pub fn qualify(self, app_name: String) -> TriggerId {
    let trigger_name = &self.0;
    TriggerId(format!("Trigger:{app_name}:{trigger_name}"))
  }
}

pub struct Connection<A: App> {
  settings: Settings,
  multiplex_sender: mpsc::Sender<SDKMessage<A::AppEvent>>,
  multiplex_receiver: mpsc::Receiver<SDKMessage<A::AppEvent>>,
  #[allow(dead_code)]
  control_stream: UnixStream,
  subscribed_streams: HashSet<StreamId>,
  trigger_callbacks: HashMap<UnqualifiedTriggerName, Vec<TriggerHandler>>,
}

/// Handles communication with the framework and calls the [App] callbacks.
pub struct Sinusoidal<A: App> {
  app: A,
  connection: Connection<A>,
}

pub const SS_PRIV_DIR: &str = "/sinusoidal/priv/";
pub const SS_DATA_DIR: &str = "/sinusoidal/data/";
pub const SS_LOG_DIR: &str = "/sinusoidal/log/";

/// Open a persistent file that's guaranteed to exist in normal `write` mode.
pub fn open_persistent_file(filename: &str) -> Result<fs::File, std::io::Error> {
  fs::File::options()
    .read(true)
    .write(true)
    .create(true)
    .truncate(false)
    .open(SS_PRIV_DIR.to_owned() + filename)
}

/// Create a new public file in write-only mode, truncating it if it exists.
pub fn create_public_file(filename: &str) -> Result<fs::File, std::io::Error> {
  fs::File::options()
    .write(true)
    .truncate(true)
    .create(true)
    .open(SS_DATA_DIR.to_owned() + filename)
}

/// Create a new log file in write-only append mode, truncating it if it exists.
pub fn create_log_file(filename: &str) -> Result<fs::File, std::io::Error> {
  fs::File::options()
    .append(true)
    .create(true)
    .open(SS_LOG_DIR.to_owned() + filename)
}

fn send16<S: Write>(stream: &mut S, buf: Vec<u8>) -> Result<(), Error> {
  send_framed(stream, buf, 2)
}

fn send32<S: Write>(stream: &mut S, buf: Vec<u8>) -> Result<(), Error> {
  send_framed(stream, buf, 4)
}

fn send_framed<S: Write>(stream: &mut S, buf: Vec<u8>, prefix_bytes: usize) -> Result<(), Error> {
  let len = buf.len();
  let header = match prefix_bytes {
    2 => {
      let len = u16::try_from(len).map_err(|_| {
        Error::new(
          ErrorKind::InvalidInput,
          format!("payload too large for 16-bit length header: {len} bytes"),
        )
      })?;
      len.to_be_bytes().to_vec()
    }
    4 => {
      let len = u32::try_from(len).map_err(|_| {
        Error::new(
          ErrorKind::InvalidInput,
          format!("payload too large for 32-bit length header: {len} bytes"),
        )
      })?;
      len.to_be_bytes().to_vec()
    }
    _ => {
      return Err(Error::new(
        ErrorKind::InvalidInput,
        format!("unsupported length prefix width: {prefix_bytes}"),
      ));
    }
  };
  stream.write_all(&header)?;
  stream.write_all(&buf)?;
  Ok(())
}

fn recv16<S: Read>(stream: &mut S) -> Result<Vec<u8>, Error> {
  recv_framed(stream, 2)
}

fn recv32<S: Read>(stream: &mut S) -> Result<Vec<u8>, Error> {
  recv_framed(stream, 4)
}

fn recv_framed<S: Read>(stream: &mut S, prefix_bytes: usize) -> Result<Vec<u8>, Error> {
  let n = match prefix_bytes {
    2 => {
      let mut len = [0; 2];
      stream
        .read_exact(&mut len)
        .map_err(|e| Error::other(format!("error while reading size: {e}")))?;
      usize::from(u16::from_be_bytes(len))
    }
    4 => {
      let mut len = [0; 4];
      stream
        .read_exact(&mut len)
        .map_err(|e| Error::other(format!("error while reading size: {e}")))?;
      usize::try_from(u32::from_be_bytes(len)).map_err(|_| {
        Error::new(
          ErrorKind::InvalidData,
          "32-bit length header does not fit in usize",
        )
      })?
    }
    _ => {
      return Err(Error::new(
        ErrorKind::InvalidInput,
        format!("unsupported length prefix width: {prefix_bytes}"),
      ));
    }
  };
  let mut buf = vec![0; n];
  stream
    .read_exact(&mut buf)
    .map_err(|e| Error::other(format!("error while reading {n} bytes of data: {e}")))?;
  Ok(buf)
}

pub struct OutStream {
  pub name: String,
  socket: UnixStream,
}

impl OutStream {
  pub fn send(&mut self, packet: Datagram) -> Result<(), Error> {
    self.send_batch(std::slice::from_ref(&packet))
  }

  /// Send multiple datagrams in a single packet.
  pub fn send_batch(&mut self, packets: &[Datagram]) -> Result<(), Error> {
    if packets.is_empty() {
      return Ok(());
    }
    let bytes = crate::sv::encode_datagrams(packets);
    let mut buf = Vec::new();
    StreamPacketDatagram!(bytes).encode(&mut buf)?;
    send16(&mut self.socket, buf)?;
    Ok(())
  }
}

#[derive(Deserialize, Debug)]
struct Settings {
  name: String,
  input_streams: HashMap<StreamId, StreamInfo>,
}

impl Settings {
  fn new(name: String, streams: Vec<StreamInfo>) -> Result<Self, Error> {
    let mut map = HashMap::new();
    for stream in streams {
      map.insert(
        StreamId::from_str(&stream.name()).map_err(Error::other)?,
        stream,
      );
    }
    Ok(Settings {
      name,
      input_streams: map,
    })
  }
}

impl<A: App> Sinusoidal<A> {
  /// Connect to the framework, create the app, and run the event loop.
  pub fn main() -> Result<(), Error> {
    Self::new()?.run()
  }

  fn new() -> Result<Sinusoidal<A>, Error> {
    let (multiplex_sender, multiplex_receiver) = mpsc::channel();
    let (control_stream, name, input_streams_s, settings_s) =
      Connection::<A>::connect(multiplex_sender.clone())?;
    let settings = Settings::new(name, serde_json::from_str(&input_streams_s)?)?;
    let app_settings = serde_json::from_str(&settings_s)?;
    let mut connection = Connection {
      settings,
      multiplex_sender,
      multiplex_receiver,
      control_stream,
      subscribed_streams: HashSet::new(),
      trigger_callbacks: HashMap::new(),
    };
    let app = A::new(app_settings, &mut connection)?;
    Ok(Sinusoidal { app, connection })
  }

  fn run(mut self) -> Result<(), Error> {
    // Set hook for panicking threads
    panic::set_hook(Box::new(|panic_info| {
      println!("{panic_info}");
      process::exit(1);
    }));

    while let Ok(msg) = self.connection.multiplex_receiver.recv() {
      match msg {
        SDKMessage::Datagrams(sid, datagrams) => {
          for datagram in datagrams {
            self
              .app
              .handle_packet(&mut self.connection, &sid, datagram)?;
          }
        }
        SDKMessage::SystemEvent(SystemEvent::Shutdown(_)) => return Ok(()),
        SDKMessage::SystemEvent(SystemEvent::NewSettings(EventNewSettings { new_settings })) => {
          let json = serde_json::from_str(&new_settings)?;
          let s = serde_json::from_value(json)?;
          match self.app.handle_new_settings(&mut self.connection, s) {
            Ok(()) => (),
            Err(err) => return Err(err),
          }
        }
        SDKMessage::PtpEvent(ptp_event) => {
          self.app.handle_ptp_event(&mut self.connection, ptp_event)?
        }
        SDKMessage::AppEvent(event) => self.app.handle_event(&mut self.connection, event)?,
        SDKMessage::Goose(packet) => self.app.handle_goose_packet(&mut self.connection, packet)?,
        SDKMessage::Trigger(trigger_id, payload) => self
          .connection
          .dispatch_trigger_callbacks(&trigger_id, &payload),
      }
    }

    panic!("Multiplex receiver failed")
  }
}

impl<A: App> Connection<A> {
  fn connect(
    multiplex_sender: mpsc::Sender<SDKMessage<A::AppEvent>>,
  ) -> Result<(UnixStream, String, String, String), Error> {
    let mut stream = UnixStream::connect("/sinusoidal/sock/control.sock")?;
    let msg = Request {
      request: Some(Req::Register(RequestRegister {})),
    };
    match Self::request(&mut stream, msg)? {
      Resp::Register(ResponseRegister {
        name,
        input_streams,
        settings,
      }) => {
        thread::spawn(move || Self::event_receiver(multiplex_sender));
        Ok((stream, name, input_streams, settings))
      }
      r => {
        panic!("Protocol violation: Got {r:?} when attempting to register")
      }
    }
  }

  fn event_receiver(sender: mpsc::Sender<SDKMessage<A::AppEvent>>) -> Result<(), Error> {
    let mut event_stream = UnixStream::connect("/sinusoidal/sock/event.sock")?;
    while let Ok(r) = recv32(&mut event_stream) {
      if let Ok(FwMsgEvent!(msg)) = WrappedSystemEvent::decode(&r[..]) {
        sender.send(SDKMessage::SystemEvent(msg)).unwrap();
      }
    }
    panic!("Event socket closed, terminating");
  }

  fn request(control_api: &mut UnixStream, req: Request) -> Result<Resp, Error> {
    println!("Request: {req:?}");
    let mut buf = Vec::new();
    req.encode(&mut buf)?;
    send32(control_api, buf)?;
    let buf = recv32(control_api)?;
    if let FwMsgResponse!(rsp) = Response::decode(&buf[..])? {
      println!("Response: {rsp:?}");
      Ok(rsp)
    } else {
      panic!("Bad response!");
    }
  }

  /// Subscribe to a stream in the framework. Returns `StreamInfo` on
  /// success. Any packets received on the stream will route through `handle_packet`.
  /// Subscribing to the same stream multiple times is idempotent - the second and
  /// subsequent calls for the same stream will return early without creating duplicate
  /// subscription threads.
  pub fn connect_to_stream(&mut self, input_stream: &InputStream) -> Result<StreamInfo, Error> {
    let stream_id = &input_stream.name;

    // Check if already subscribed to this stream to prevent duplicate subscriptions
    if self.subscribed_streams.contains(stream_id) {
      return match self.settings.input_streams.get(stream_id) {
        Some(info) => Ok(info.clone()),
        None => {
          let name = stream_id.to_string();
          Err(Error::other(format!("Unknown stream: {name}")))
        }
      };
    }

    match self.settings.input_streams.get(stream_id) {
      Some(info) => {
        let sender = self.multiplex_sender.clone();
        let stream_id = stream_id.clone();
        thread::spawn(move || Self::stream_receiver(stream_id, sender));
        self.subscribed_streams.insert(input_stream.name.clone());
        Ok(info.clone())
      }
      None => {
        let name = stream_id.to_string();
        Err(Error::other(format!("Unknown stream: {name}")))
      }
    }
  }

  fn stream_receiver(stream_id: StreamId, sender: mpsc::Sender<SDKMessage<A::AppEvent>>) -> () {
    let sock_name = socket_token(&stream_id.to_string());
    let sock_file = format!("/sinusoidal/sock/{sock_name}.sock");
    let mut stream = UnixStream::connect(sock_file)
      .unwrap_or_else(|e| panic!("Failed to connect to {stream_id}: {e}"));
    loop {
      let r = recv16(&mut stream)
        .unwrap_or_else(|e| panic!("Failed to receive packet on {stream_id}: {e}"));
      let msg = StreamPacket::decode(&r[..]).expect("StreamPacket decode failure");
      match msg {
        StreamPacket {
          payload: Some(stream_packet::Payload::Datagrams(bytes)),
        } => {
          let datagrams = crate::sv::decode_datagrams(&bytes)
            .unwrap_or_else(|e| panic!("Failed to decode datagrams on {stream_id}: {e}"));
          sender
            .send(SDKMessage::Datagrams(stream_id.clone(), datagrams))
            .unwrap();
        }
        StreamPacketPtp!(ptp) => {
          if let StreamId::Sys {
            stream: SysStream::Time,
          } = stream_id
          {
            sender.send(SDKMessage::PtpEvent(ptp)).unwrap()
          }
        }
        StreamPacketGoose!(goose) => {
          if let StreamId::Goose { go_id: _ } = stream_id {
            sender
              .send(SDKMessage::Goose(GooseEvent::new(goose)))
              .unwrap()
          }
        }
        _ => {} // TODO: non-datagram messages
      };
    }
  }

  fn trigger_receiver(
    app_name: String,
    trigger: &UnqualifiedTriggerName,
    sender: mpsc::Sender<SDKMessage<A::AppEvent>>,
  ) {
    let trigger_id = trigger.clone().qualify(app_name);
    let sock_name = socket_token(&trigger_id.0);
    let sock_file = format!("/sinusoidal/sock/{sock_name}.in.sock");
    let mut stream = UnixStream::connect(&sock_file)
      .unwrap_or_else(|e| panic!("Failed to connect to {trigger:?}: {e}"));

    loop {
      let msg = match recv16(&mut stream) {
        Ok(raw) => raw,
        Err(err) => {
          eprintln!("Error in trigger receiver for {trigger:?}: {err}");
          break;
        }
      };
      let payload: TriggerPayload = match serde_json::from_slice(&msg) {
        Ok(payload) => payload,
        Err(err) => {
          eprintln!("Error decoding trigger payload for {trigger:?}: {err}");
          continue;
        }
      };
      if sender
        .send(SDKMessage::Trigger(trigger.clone(), payload))
        .is_err()
      {
        break;
      }
    }
  }

  /// Register a callback executed whenever `trigger_id` receives a payload.
  ///
  /// Trigger callbacks are dispatched on the app event loop thread.
  pub fn on_trigger<F, T>(&mut self, trigger: &UnqualifiedTriggerName, mut handler: F)
  where
    F: FnMut(&T) + Send + 'static,
    T: serde::de::DeserializeOwned,
  {
    let handler = Box::new(move |payload: &TriggerPayload| {
      let payload: T = <T as serde::Deserialize>::deserialize(payload)?;
      handler(&payload);
      Ok(())
    });

    match self.trigger_callbacks.entry(trigger.clone()) {
      Entry::Vacant(entry) => {
        entry.insert(vec![Box::new(handler)]);
        let trigger = trigger.clone();
        let sender = self.multiplex_sender.clone();
        let app_name = self.app_name().to_string();
        thread::spawn(move || Self::trigger_receiver(app_name, &trigger, sender));
      }
      Entry::Occupied(mut entry) => {
        entry.get_mut().push(Box::new(handler));
      }
    }
  }

  /// Send a JSON trigger payload to a remote trigger.
  pub fn send_trigger<T>(&mut self, trigger_id: &TriggerId, payload: &T) -> Result<(), Error>
  where
    T: serde::Serialize,
  {
    let sock_name = socket_token(&trigger_id.0);
    let sock_file = format!("/sinusoidal/sock/{sock_name}.out.sock");
    let mut stream = UnixStream::connect(sock_file)?;
    let payload = serde_json::to_vec(payload)?;
    send16(&mut stream, payload)
  }

  fn dispatch_trigger_callbacks(
    &mut self,
    trigger: &UnqualifiedTriggerName,
    payload: &TriggerPayload,
  ) {
    if let Some(callbacks) = self.trigger_callbacks.get_mut(trigger) {
      for callback in callbacks.iter_mut() {
        if let Err(err) = callback(payload) {
          eprintln!("Error in trigger callback for {trigger:?}: {err}");
        }
      }
    }
  }

  /// Register an output stream.
  pub fn register_out_stream(&mut self, stream_name: &str) -> Result<OutStream, Error> {
    let stream_id = StreamId::App {
      app_name: self.settings.name.clone(),
      stream_name: stream_name.to_string(),
    };
    let sock_name = socket_token(&stream_id.to_string());
    let socket = UnixStream::connect(format!("/sinusoidal/sock/{sock_name}.out.sock"))?;
    let stream = OutStream {
      name: stream_name.to_string(),
      socket,
    };
    Ok(stream)
  }

  pub fn get_event_sender(&mut self) -> Box<dyn Fn(A::AppEvent) + Send> {
    let tx = self.multiplex_sender.clone();
    Box::new(move |e| tx.send(SDKMessage::AppEvent(e)).unwrap())
  }

  pub fn app_name(&self) -> &str {
    &self.settings.name
  }
}

/// Compute a fixed-size, file-system-safe token for a socket name.
///
/// Takes the first 16 bytes of the SHA-256 digest of `name` and encodes them
/// with URL-safe base64 (no padding). The framework and all SDKs use this
/// function to derive Unix socket file names from stream/trigger IDs, avoiding
/// path-length limits imposed by the OS.
pub fn socket_token(name: &str) -> String {
  let digest = Sha256::digest(name.as_bytes());
  URL_SAFE_NO_PAD.encode(&digest[..16])
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ClockSynch {
  NoSync,
  Local,
  Global,
  LocalSet(u8),
}

impl ClockSynch {
  pub fn decode(raw: u32) -> Result<Option<Self>, Error> {
    match raw {
      0 => Ok(Some(Self::NoSync)),
      1 => Ok(Some(Self::Local)),
      2 => Ok(Some(Self::Global)),
      3 => Ok(None),
      5..=254 => Ok(Some(Self::LocalSet(raw as u8))),
      _ => Err(Error::new(
        std::io::ErrorKind::InvalidData,
        format!("invalid clock_synch value {raw}"),
      )),
    }
  }

  pub fn encode(value: Option<Self>) -> u32 {
    match value {
      Some(Self::NoSync) => 0,
      Some(Self::Local) => 1,
      Some(Self::Global) => 2,
      None => 3,
      Some(Self::LocalSet(level @ 5..=254)) => u32::from(level),
      Some(Self::LocalSet(level)) => panic!("invalid local clock_synch payload {level}"),
    }
  }
}

fn clock_synch_rank(clock_synch: ClockSynch) -> u8 {
  match clock_synch {
    ClockSynch::NoSync => 0,
    ClockSynch::Local => 1,
    ClockSynch::Global => 255,
    ClockSynch::LocalSet(level) => level,
  }
}

// The idea is to keep the _worst_ sync level
pub fn merge_clock_synch(
  left: Option<ClockSynch>,
  right: Option<ClockSynch>,
) -> Option<ClockSynch> {
  match (left, right) {
    (None, other) | (other, None) => other,
    (Some(left), Some(right)) => {
      if clock_synch_rank(left) <= clock_synch_rank(right) {
        Some(left)
      } else {
        Some(right)
      }
    }
  }
}

pub fn decode_datagram_flags(raw_flags: &[u8], value_count: usize) -> Result<Vec<u16>, Error> {
  match raw_flags.len() {
    // Empty flags means all flags are 0
    0 => Ok(vec![0; value_count]),
    // Two empty bytes means no flags (saves 2 bytes in the most common case).
    2 if raw_flags[0] == 0 && raw_flags[1] == 0 => Ok(vec![]),
    2 if value_count > 0 => {
      let shared_flags = u16::from_be_bytes([raw_flags[0], raw_flags[1]]);
      Ok(vec![shared_flags; value_count])
    }
    len if len == value_count * 2 => Ok(
      raw_flags
        .chunks_exact(2)
        .map(|chunk| u16::from_be_bytes([chunk[0], chunk[1]]))
        .collect(),
    ),
    len if len % 2 != 0 => Err(Error::new(
      std::io::ErrorKind::InvalidData,
      format!("invalid flags payload length {len}"),
    )),
    len => Err(Error::new(
      std::io::ErrorKind::InvalidData,
      format!("flags payload length {len} does not match value count {value_count}"),
    )),
  }
}

pub fn encode_datagram_flags(flags: &[u16]) -> Vec<u8> {
  match flags.split_first() {
    // No flags is encoded as two 0 bytes
    None => vec![0, 0],
    // And all flags 0 is encoded as the empty vector
    Some((&first_flag, rest)) if first_flag == 0 && rest.iter().all(|&flag| flag == 0) => {
      Vec::new()
    }
    Some((&first_flag, rest)) if rest.iter().all(|&flag| flag == first_flag) => {
      first_flag.to_be_bytes().to_vec()
    }
    _ => flags.iter().flat_map(|flag| flag.to_be_bytes()).collect(),
  }
}

// Merge (bitwise or) flags all flags are expanded.
pub fn merge_datagram_flags(left: &[u16], right: &[u16]) -> Result<Vec<u16>, Error> {
  match (left.is_empty(), right.is_empty()) {
    (true, true) => Ok(Vec::new()),
    (true, false) => Ok(right.to_vec()),
    (false, true) => Ok(left.to_vec()),
    (false, false) if left.len() == right.len() => Ok(
      left
        .iter()
        .zip(right.iter())
        .map(|(left, right)| left | right)
        .collect(),
    ),
    (false, false) => Err(Error::new(
      std::io::ErrorKind::InvalidInput,
      format!(
        "cannot merge flag vectors of different lengths ({} vs {})",
        left.len(),
        right.len()
      ),
    )),
  }
}

pub fn collapse_datagram_flags(flags: &[u16]) -> u16 {
  flags.iter().copied().fold(0, |acc, flag| acc | flag)
}

impl Datagram {
  pub fn decoded_clock_synch(&self) -> Result<Option<ClockSynch>, Error> {
    ClockSynch::decode(self.clock_synch)
  }

  pub fn set_clock_synch(&mut self, clock_synch: Option<ClockSynch>) {
    self.clock_synch = ClockSynch::encode(clock_synch);
  }

  pub fn decoded_flags(&self) -> Result<Vec<u16>, Error> {
    decode_datagram_flags(&self.flags, self.values.len())
  }

  pub fn set_flags(&mut self, flags: &[u16]) -> Result<(), Error> {
    if !flags.is_empty() && flags.len() != self.values.len() {
      return Err(Error::new(
        std::io::ErrorKind::InvalidInput,
        format!(
          "flags length {} does not match value count {}",
          flags.len(),
          self.values.len()
        ),
      ));
    }
    self.flags = encode_datagram_flags(flags);
    Ok(())
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use quickcheck::TestResult;
  use std::io::Cursor;

  #[quickcheck_macros::quickcheck]
  fn send16_recv16_roundtrip(payload: Vec<u8>) -> TestResult {
    if payload.len() > u16::MAX as usize {
      return TestResult::discard();
    }

    let mut stream = Cursor::new(Vec::new());
    send16(&mut stream, payload.clone()).unwrap();
    stream.set_position(0);
    let received = recv16(&mut stream).unwrap();
    TestResult::from_bool(payload == received)
  }

  #[quickcheck_macros::quickcheck]
  fn send32_recv32_roundtrip(payload: Vec<u8>) -> TestResult {
    let mut stream = Cursor::new(Vec::new());
    send32(&mut stream, payload.clone()).unwrap();
    stream.set_position(0);
    let received = recv32(&mut stream).unwrap();
    TestResult::from_bool(payload == received)
  }

  #[test]
  fn clock_synch_roundtrips() {
    assert_eq!(ClockSynch::decode(ClockSynch::encode(None)).unwrap(), None);
    assert_eq!(
      ClockSynch::decode(ClockSynch::encode(Some(ClockSynch::NoSync))).unwrap(),
      Some(ClockSynch::NoSync)
    );
    assert_eq!(
      ClockSynch::decode(ClockSynch::encode(Some(ClockSynch::Local))).unwrap(),
      Some(ClockSynch::Local)
    );
    assert_eq!(
      ClockSynch::decode(ClockSynch::encode(Some(ClockSynch::Global))).unwrap(),
      Some(ClockSynch::Global)
    );
    assert_eq!(
      ClockSynch::decode(ClockSynch::encode(Some(ClockSynch::LocalSet(5)))).unwrap(),
      Some(ClockSynch::LocalSet(5))
    );
    assert_eq!(
      ClockSynch::decode(ClockSynch::encode(Some(ClockSynch::LocalSet(254)))).unwrap(),
      Some(ClockSynch::LocalSet(254))
    );
  }

  #[test]
  fn clock_synch_rejects_invalid_values() {
    assert!(ClockSynch::decode(4).is_err());
    assert!(ClockSynch::decode(255).is_err());
  }

  #[test]
  fn decode_datagram_flags_handles_compact_encodings() {
    assert_eq!(decode_datagram_flags(&[], 3).unwrap(), vec![0, 0, 0]);
    assert_eq!(
      decode_datagram_flags(&0x1234_u16.to_be_bytes(), 3).unwrap(),
      vec![0x1234, 0x1234, 0x1234]
    );
    assert_eq!(
      decode_datagram_flags(&[0x00, 0x01, 0x00, 0x02, 0x00, 0x03], 3).unwrap(),
      vec![1, 2, 3]
    );
  }

  #[test]
  fn decode_datagram_flags_rejects_invalid_lengths() {
    assert!(decode_datagram_flags(&[0x00], 1).is_err());
    assert!(decode_datagram_flags(&[0x00, 0x01, 0x00, 0x02], 3).is_err());
    assert!(decode_datagram_flags(&[0x00, 0x01], 0).is_err());
  }

  #[test]
  fn encode_datagram_flags_compacts_zero_and_shared_values() {
    assert_eq!(encode_datagram_flags(&[]), vec![0, 0]);
    assert_eq!(encode_datagram_flags(&[0, 0, 0]), Vec::<u8>::new());
    assert_eq!(encode_datagram_flags(&[0x1234, 0x1234]), vec![0x12, 0x34]);
    assert_eq!(
      encode_datagram_flags(&[1, 2, 3]),
      vec![0x00, 0x01, 0x00, 0x02, 0x00, 0x03]
    );
  }

  #[test]
  fn datagram_helpers_roundtrip() {
    let mut datagram = Datagram {
      timestamp: 0,
      sample_count: 3,
      values: vec![10, 20, 30],
      clock_synch: 3,
      flags: Vec::new(),
      gm_identity: vec![1, 2, 3],
    };

    datagram.set_clock_synch(Some(ClockSynch::LocalSet(9)));
    datagram.set_flags(&[0x00ff, 0x00ff, 0x00ff]).unwrap();

    assert_eq!(
      datagram.decoded_clock_synch().unwrap(),
      Some(ClockSynch::LocalSet(9))
    );
    assert_eq!(datagram.decoded_flags().unwrap(), vec![0x00ff; 3]);
    assert_eq!(datagram.flags, vec![0x00, 0xff]);
  }

  #[test]
  fn datagram_set_flags_requires_one_entry_per_value() {
    let mut datagram = Datagram {
      timestamp: 0,
      sample_count: 2,
      values: vec![10, 20],
      clock_synch: 3,
      flags: Vec::new(),
      gm_identity: Vec::new(),
    };

    assert!(datagram.set_flags(&[0x0011]).is_err());
  }

  #[test]
  fn datagram_set_flags_allows_no_flags_sentinel() {
    let mut datagram = Datagram {
      timestamp: 0,
      sample_count: 2,
      values: vec![10, 20],
      clock_synch: 3,
      flags: Vec::new(),
      gm_identity: Vec::new(),
    };

    datagram.set_flags(&[]).unwrap();

    assert_eq!(datagram.flags, vec![0, 0]);
    assert_eq!(datagram.decoded_flags().unwrap(), Vec::<u16>::new());
  }

  #[test]
  fn merge_datagram_flags_preserves_empty_as_identity() {
    assert_eq!(merge_datagram_flags(&[], &[]).unwrap(), Vec::<u16>::new());
    assert_eq!(merge_datagram_flags(&[], &[1, 2]).unwrap(), vec![1, 2]);
    assert_eq!(merge_datagram_flags(&[1, 2], &[]).unwrap(), vec![1, 2]);
    assert_eq!(merge_datagram_flags(&[1, 2], &[4, 8]).unwrap(), vec![5, 10]);
    assert!(merge_datagram_flags(&[1], &[1, 2]).is_err());
  }

  #[test]
  fn merge_clock_synch_prefers_more_constrained_value() {
    assert_eq!(
      merge_clock_synch(Some(ClockSynch::Global), Some(ClockSynch::Local)),
      Some(ClockSynch::Local)
    );
    assert_eq!(
      merge_clock_synch(Some(ClockSynch::LocalSet(5)), Some(ClockSynch::LocalSet(9))),
      Some(ClockSynch::LocalSet(5))
    );
    assert_eq!(
      merge_clock_synch(None, Some(ClockSynch::Global)),
      Some(ClockSynch::Global)
    );
  }

  #[test]
  fn socket_token_matches_framework_rules() {
    assert_eq!(socket_token("SV:4001:Mu1"), "Q16cYaFEBkUMxEZoZs6y9A");
    assert_eq!(
      socket_token("App:EM_SIP1MOD2MU103:FundamentalPhaseToPhaseVoltage-Total"),
      "7sZ4vTWZUNLcFDe-CzroSA"
    );
    assert_eq!(
      socket_token("Trigger:PMU_SIP1MOD2MU103:trip"),
      "Xhv9u8cakYL-UFbl2-NdKA"
    );
  }
}