psdk-fruit-emapi 0.0.1

use std::collections::BTreeMap;

use crate::{
  endian::{read_uint16, read_uint32, uint16, uint32},
  error::{EmapiError, Result},
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TlvEntry {
  pub tag: u8,
  pub value: Vec<u8>,
}

impl TlvEntry {
  pub fn new(tag: u8, value: impl Into<Vec<u8>>) -> Self {
    Self {
      tag,
      value: value.into(),
    }
  }

  pub fn try_new(tag: u16, value: impl Into<Vec<u8>>) -> Result<Self> {
    if tag > 0xFF {
      return Err(EmapiError::Range {
        name: "tag",
        min: 0,
        max: 0xFF,
        value: u64::from(tag),
      });
    }
    Ok(Self::new(tag as u8, value))
  }

  pub fn uint8(tag: u8, value: u8) -> Result<Self> {
    Ok(Self::new(tag, [value]))
  }

  pub fn uint16(tag: u8, value: u16) -> Result<Self> {
    Ok(Self::new(tag, uint16(value)))
  }

  pub fn uint32(tag: u8, value: u32) -> Result<Self> {
    Ok(Self::new(tag, uint32(value)))
  }

  pub fn string(tag: u8, value: &str) -> Result<Self> {
    let mut bytes = value.as_bytes().to_vec();
    bytes.push(0x00);
    Ok(Self::new(tag, bytes))
  }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TlvData {
  entries: BTreeMap<u8, TlvEntry>,
}

impl TlvData {
  pub fn get(&self, tag: u8) -> Option<&TlvEntry> {
    self.entries.get(&tag)
  }

  pub fn contains_tag(&self, tag: u8) -> bool {
    self.entries.contains_key(&tag)
  }

  pub fn uint8(&self, tag: u8) -> Result<Option<u8>> {
    let Some(value) = self.value(tag) else {
      return Ok(None);
    };
    if value.len() != 1 {
      return Err(EmapiError::InvalidTlvType {
        tag,
        expected: "uint8",
        actual_len: value.len(),
      });
    }
    Ok(Some(value[0]))
  }

  pub fn uint16(&self, tag: u8) -> Result<Option<u16>> {
    let Some(value) = self.value(tag) else {
      return Ok(None);
    };
    if value.len() != 2 {
      return Err(EmapiError::InvalidTlvType {
        tag,
        expected: "uint16",
        actual_len: value.len(),
      });
    }
    Ok(Some(read_uint16(value, 0)?))
  }

  pub fn uint32(&self, tag: u8) -> Result<Option<u32>> {
    let Some(value) = self.value(tag) else {
      return Ok(None);
    };
    if value.len() != 4 {
      return Err(EmapiError::InvalidTlvType {
        tag,
        expected: "uint32",
        actual_len: value.len(),
      });
    }
    Ok(Some(read_uint32(value, 0)?))
  }

  pub fn string(&self, tag: u8) -> Result<Option<String>> {
    let Some(value) = self.value(tag) else {
      return Ok(None);
    };
    let end = if value.last() == Some(&0x00) {
      value.len() - 1
    } else {
      value.len()
    };
    String::from_utf8(value[..end].to_vec())
      .map(Some)
      .map_err(|_| EmapiError::InvalidUtf8)
  }

  fn value(&self, tag: u8) -> Option<&[u8]> {
    self.entries.get(&tag).map(|entry| entry.value.as_slice())
  }
}

pub struct Tlv;

impl Tlv {
  pub fn encode(entries: &[TlvEntry]) -> Result<Vec<u8>> {
    let mut result = Vec::new();
    for entry in entries {
      if entry.value.len() > usize::from(u16::MAX) {
        return Err(EmapiError::Range {
          name: "length",
          min: 0,
          max: u64::from(u16::MAX),
          value: entry.value.len() as u64,
        });
      }
      result.push(entry.tag);
      result.extend_from_slice(&uint16(entry.value.len() as u16));
      result.extend_from_slice(&entry.value);
    }
    Ok(result)
  }

  pub fn decode(data: &[u8]) -> Result<TlvData> {
    let mut entries = BTreeMap::new();
    let mut offset = 0;
    while offset < data.len() {
      if offset + 3 > data.len() {
        return Err(EmapiError::TruncatedTlvHeader);
      }
      let tag = data[offset];
      let length = usize::from(read_uint16(data, offset + 1)?);
      let value_start = offset + 3;
      let value_end = value_start + length;
      if value_end > data.len() {
        return Err(EmapiError::TruncatedTlvValue);
      }
      entries.insert(
        tag,
        TlvEntry::new(tag, data[value_start..value_end].to_vec()),
      );
      offset = value_end;
    }
    Ok(TlvData { entries })
  }
}