km003c-lib 0.4.0

use crate::adc::{AdcDataRaw, AdcDataSimple};
use crate::adcqueue::{AdcQueueData, AdcQueueRawData, GraphSampleRate};
use crate::auth::{self, AuthCredential, StreamingAuthResult};
use crate::constants::*;
use crate::error::KMError;
use crate::offline::{LogMetadata, LogMetadataResponse};
use crate::packet::{
    Attribute, AttributeSet, CtrlHeader, DataHeader, LogicalPacket, PacketType, RawPacket, StreamingAuthHeader,
};
use crate::pd::{PdEventStream, PdStatus, PdStatusRaw};
use crate::pd_trace::PdTrace;
use crate::settings::Settings;
use bytes::Bytes;
use num_enum::FromPrimitive;
use zerocopy::{FromBytes, IntoBytes};

const PD_MONITOR_ENABLED_PARAMETER: u16 = 1;
const PD_MONITOR_DISABLED_PARAMETER: u16 = 0;

/// Represents parsed payload data from logical packets
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "python", derive(pyo3::IntoPyObject))]
pub enum PayloadData {
    Adc(AdcDataSimple),
    AdcQueue(AdcQueueData),
    AdcQueueRaw(AdcQueueRawData),
    PdStatus(PdStatus),
    PdEvents(PdEventStream),
    PdTrace(PdTrace),
    Settings(Settings),
    LogMetadata(LogMetadataResponse),
    Unknown { attribute: Attribute, data: Vec<u8> },
}

#[derive(Debug, Clone, PartialEq)]
pub enum Packet {
    /// Data response with parsed payload data
    DataResponse { payloads: Vec<PayloadData> },
    /// Request data with attribute set
    GetData { attributes: AttributeSet },
    /// Start AdcQueue graph mode with sample rate.
    ///
    /// `CtrlHeader` places the logical rate index (0=2SPS, 1=10SPS, 2=50SPS,
    /// 3=1000SPS) in the wire header, producing byte 2 values 0, 2, 4, and 6
    /// respectively. A header carrying an index outside that range parses as
    /// [`Packet::Generic`] rather than being coerced into a known rate.
    StartGraph { rate: GraphSampleRate },
    /// Stop AdcQueue graph mode
    StopGraph,
    /// Accept response
    Accept { id: u8 },
    /// Reject response
    Reject { id: u8 },
    /// NotReadable response - memory address not accessible
    NotReadable { id: u8 },
    /// Connect command
    Connect,
    /// Disconnect command
    Disconnect,
    /// Enable PD monitor/sniffer
    EnablePdMonitor,
    /// Disable PD monitor/sniffer
    DisablePdMonitor,
    /// MemoryRead command (0x44) - read device memory with encrypted payload
    MemoryRead {
        /// Memory address to read from
        address: u32,
        /// Number of bytes to read
        size: u32,
    },
    /// StreamingAuth command (0x4C) - authenticate for AdcQueue streaming
    StreamingAuth {
        /// Device or calibration credential to authenticate with
        credential: AuthCredential,
    },
    /// StreamingAuth response (0x4C) - authentication result
    StreamingAuthResponse(StreamingAuthResult),
    /// Generic packet for types we haven't specifically implemented yet
    Generic(RawPacket),
}

/// Define a `Packet` accessor that returns the first payload of one variant.
macro_rules! payload_accessor {
    ($(#[$doc:meta])* $name:ident -> $target:ty, $variant:ident) => {
        $(#[$doc])*
        pub fn $name(&self) -> Option<&$target> {
            self.find_payload(|payload| match payload {
                PayloadData::$variant(inner) => Some(inner),
                _ => None,
            })
        }
    };
}

impl Packet {
    /// Return the first payload for which `select` yields a value.
    ///
    /// Only [`Packet::DataResponse`] carries payloads; every other variant
    /// yields `None`.
    pub fn find_payload<'a, T: ?Sized>(&'a self, select: impl Fn(&'a PayloadData) -> Option<&'a T>) -> Option<&'a T> {
        match self {
            Self::DataResponse { payloads } => payloads.iter().find_map(select),
            _ => None,
        }
    }

    payload_accessor!(
        /// Get ADC data from the packet, if present.
        get_adc -> AdcDataSimple, Adc
    );
    payload_accessor!(
        /// Get AdcQueue data from the packet, if present.
        get_adc_queue -> AdcQueueData, AdcQueue
    );
    payload_accessor!(
        /// Get losslessly parsed AdcQueue data whose graph rate is unknown.
        get_adc_queue_raw -> AdcQueueRawData, AdcQueueRaw
    );
    payload_accessor!(
        /// Get PD status from the packet, if present.
        get_pd_status -> PdStatus, PdStatus
    );
    payload_accessor!(
        /// Get PD events from the packet, if present.
        get_pd_events -> PdEventStream, PdEvents
    );
    payload_accessor!(
        /// Get the internal PD state-machine trace, if present.
        get_pd_trace -> PdTrace, PdTrace
    );
    payload_accessor!(
        /// Get the read-only settings payload, if present.
        get_settings -> Settings, Settings
    );
    payload_accessor!(
        /// Get the response to a LogMetadata request, if present.
        get_log_metadata -> LogMetadataResponse, LogMetadata
    );

    /// Check if packet has a specific payload type
    pub fn has_payload(&self, attr: Attribute) -> bool {
        match self {
            Self::DataResponse { payloads } => payloads.iter().any(|p| match p {
                PayloadData::Adc(_) => attr == Attribute::Adc,
                PayloadData::AdcQueue(_) => attr == Attribute::AdcQueue,
                PayloadData::AdcQueueRaw(_) => attr == Attribute::AdcQueue,
                PayloadData::PdStatus(_) | PayloadData::PdEvents(_) => attr == Attribute::PdPacket,
                PayloadData::PdTrace(_) => attr == Attribute::PdTrace,
                PayloadData::Settings(_) => attr == Attribute::Settings,
                PayloadData::LogMetadata(_) => attr == Attribute::LogMetadata,
                PayloadData::Unknown { attribute, .. } => *attribute == attr,
            }),
            _ => false,
        }
    }
}

impl TryFrom<RawPacket> for Packet {
    type Error = KMError;

    fn try_from(raw_packet: RawPacket) -> Result<Self, Self::Error> {
        Self::from_raw(raw_packet, None)
    }
}

impl Packet {
    /// Parse a raw packet using the configured AdcQueue graph rate.
    ///
    /// Live device users should supply the rate used with `StartGraph` because
    /// auxiliary AdcQueue voltage fields use rate-dependent units.
    pub fn from_raw_with_graph_rate(raw_packet: RawPacket, rate: GraphSampleRate) -> Result<Self, KMError> {
        Self::from_raw(raw_packet, Some(rate))
    }

    fn from_raw(raw_packet: RawPacket, graph_rate: Option<GraphSampleRate>) -> Result<Self, KMError> {
        match raw_packet {
            RawPacket::Ctrl { header, payload } => {
                let packet_type = PacketType::from_primitive(header.packet_type());
                let attribute_set = AttributeSet::from_raw(header.attribute());

                match packet_type {
                    PacketType::GetData => Ok(Packet::GetData {
                        attributes: attribute_set,
                    }),
                    // An unknown rate index is preserved verbatim as a generic
                    // packet rather than coerced into one of the four rates.
                    PacketType::StartGraph => match GraphSampleRate::try_from(attribute_set.raw()) {
                        Ok(rate) => Ok(Packet::StartGraph { rate }),
                        Err(_) => Ok(Packet::Generic(RawPacket::Ctrl { header, payload })),
                    },
                    PacketType::StopGraph => Ok(Packet::StopGraph),
                    PacketType::Accept => Ok(Packet::Accept { id: header.id() }),
                    PacketType::Rejected => Ok(Packet::Reject { id: header.id() }),
                    PacketType::NotReadable => Ok(Packet::NotReadable { id: header.id() }),
                    PacketType::Connect => Ok(Packet::Connect),
                    PacketType::Disconnect => Ok(Packet::Disconnect),
                    PacketType::EnablePdMonitor => Ok(Packet::EnablePdMonitor),
                    PacketType::DisablePdMonitor => Ok(Packet::DisablePdMonitor),
                    _ => Ok(Packet::Generic(RawPacket::Ctrl { header, payload })),
                }
            }
            RawPacket::SimpleData { header, payload } => {
                let packet_type = PacketType::from_primitive(header.packet_type());

                match packet_type {
                    PacketType::MemoryRead => {
                        if !header.reserved_flag()
                            && header.into_bytes()[2..4] == [0x01, 0x01]
                            && let Some((address, size)) = auth::parse_memory_read_payload(&payload)
                        {
                            Ok(Packet::MemoryRead { address, size })
                        } else {
                            Ok(Packet::Generic(RawPacket::SimpleData { header, payload }))
                        }
                    }
                    PacketType::StreamingAuth => {
                        let auth_header = StreamingAuthHeader::from_bytes(header.into_bytes());
                        if auth_header.attribute() == 0x0200
                            && let Some(credential) = auth::parse_streaming_auth_request_payload(&payload)
                        {
                            Ok(Packet::StreamingAuth { credential })
                        } else if let Some(result) = auth::parse_streaming_auth_response_parts(auth_header, &payload) {
                            Ok(Packet::StreamingAuthResponse(result))
                        } else {
                            Ok(Packet::Generic(RawPacket::SimpleData { header, payload }))
                        }
                    }
                    _ => Ok(Packet::Generic(RawPacket::SimpleData { header, payload })),
                }
            }
            RawPacket::Data { logical_packets, .. } => {
                // Parse logical packets into PayloadData
                let mut payloads = Vec::new();

                for lp in logical_packets {
                    let payload_data = match lp.attribute {
                        Attribute::Adc => {
                            // Parse ADC data (44 bytes)
                            if lp.payload.len() < ADC_DATA_SIZE {
                                return Err(KMError::InvalidPacket(format!(
                                    "ADC payload too small: expected {}, got {}",
                                    ADC_DATA_SIZE,
                                    lp.payload.len()
                                )));
                            }

                            let adc_data_raw = AdcDataRaw::ref_from_bytes(&lp.payload[..ADC_DATA_SIZE])
                                .map_err(|_| KMError::InvalidPacket("Failed to parse ADC data".to_string()))?;
                            let adc_data = AdcDataSimple::from(*adc_data_raw);
                            PayloadData::Adc(adc_data)
                        }
                        Attribute::AdcQueue => {
                            // Parse AdcQueue data (multiple 20-byte samples)
                            // Note: Extended header size field (typically 20) indicates size per sample,
                            // not total payload size. Actual payload contains N samples.
                            if let Some(rate) = graph_rate {
                                PayloadData::AdcQueue(AdcQueueData::from_bytes_with_rate(lp.payload.as_ref(), rate)?)
                            } else {
                                PayloadData::AdcQueueRaw(AdcQueueRawData::from_bytes(lp.payload.as_ref())?)
                            }
                        }
                        Attribute::PdPacket => {
                            // Determine if this is PD status or PD events
                            if lp.payload.len() == PD_STATUS_SIZE {
                                // PD Status (12 bytes)
                                let pd_status_raw = PdStatusRaw::ref_from_bytes(lp.payload.as_ref())
                                    .map_err(|_| KMError::InvalidPacket("Failed to parse PD status".to_string()))?;
                                PayloadData::PdStatus(PdStatus::from(*pd_status_raw))
                            } else {
                                // PD Event Stream
                                let pd_events = PdEventStream::from_bytes(Bytes::from(lp.payload))?;
                                PayloadData::PdEvents(pd_events)
                            }
                        }
                        Attribute::PdTrace => PayloadData::PdTrace(PdTrace::from_bytes(&lp.payload)?),
                        Attribute::Settings => PayloadData::Settings(Settings::from_bytes(lp.payload.as_ref())?),
                        Attribute::LogMetadata => {
                            if lp.payload.is_empty() {
                                PayloadData::LogMetadata(LogMetadataResponse::Empty)
                            } else {
                                if !lp.payload.len().is_multiple_of(crate::offline::LOG_METADATA_SIZE) {
                                    return Err(KMError::InvalidPacket(format!(
                                        "LogMetadata payload length must be a multiple of {}, got {}",
                                        crate::offline::LOG_METADATA_SIZE,
                                        lp.payload.len()
                                    )));
                                }
                                let entries = lp
                                    .payload
                                    .as_chunks::<{ crate::offline::LOG_METADATA_SIZE }>()
                                    .0
                                    .iter()
                                    .map(|chunk| LogMetadata::from_bytes(chunk))
                                    .collect::<Result<Vec<_>, _>>()?;
                                PayloadData::LogMetadata(LogMetadataResponse::Available(entries))
                            }
                        }
                        _ => PayloadData::Unknown {
                            attribute: lp.attribute,
                            data: lp.payload.to_vec(),
                        },
                    };

                    payloads.push(payload_data);
                }

                Ok(Packet::DataResponse { payloads })
            }
        }
    }
}

impl Packet {
    /// Convert a high-level packet to a raw packet with the given transaction ID
    pub fn to_raw_packet(self, id: u8) -> Result<RawPacket, KMError> {
        let raw_packet = match self {
            Packet::DataResponse { payloads } => {
                let mut logical_packets = Vec::new();

                for payload in payloads {
                    match payload {
                        PayloadData::Adc(adc) => {
                            let adc_raw = AdcDataRaw::from(adc);
                            logical_packets.push(LogicalPacket {
                                attribute: Attribute::Adc,
                                next: false, // Will be fixed below
                                chunk: 0,
                                size: ADC_DATA_SIZE as u16,
                                payload: adc_raw.as_bytes().to_vec(),
                            });
                        }
                        PayloadData::PdStatus(pd_status) => {
                            let raw = PdStatusRaw::from(pd_status);

                            logical_packets.push(LogicalPacket {
                                attribute: Attribute::PdPacket,
                                next: false, // Will be fixed below
                                chunk: 0,
                                size: PD_STATUS_SIZE as u16,
                                payload: raw.as_bytes().to_vec(),
                            });
                        }
                        PayloadData::AdcQueue(_) => {
                            return Err(KMError::UnsupportedSerialization { packet: "AdcQueue" });
                        }
                        PayloadData::AdcQueueRaw(queue) => {
                            logical_packets.push(LogicalPacket {
                                attribute: Attribute::AdcQueue,
                                next: false,
                                chunk: 0,
                                size: 20,
                                payload: queue.to_bytes(),
                            });
                        }
                        PayloadData::PdEvents(_) => {
                            return Err(KMError::UnsupportedSerialization {
                                packet: "PdEventStream",
                            });
                        }
                        PayloadData::PdTrace(trace) => {
                            logical_packets.push(LogicalPacket {
                                attribute: Attribute::PdTrace,
                                next: false,
                                chunk: 0,
                                size: 0,
                                payload: trace.to_bytes()?,
                            });
                        }
                        PayloadData::Settings(settings) => {
                            logical_packets.push(LogicalPacket {
                                attribute: Attribute::Settings,
                                next: false,
                                chunk: 0,
                                size: crate::settings::SETTINGS_SIZE as u16,
                                payload: settings.to_bytes().to_vec(),
                            });
                        }
                        PayloadData::LogMetadata(response) => {
                            let payload = match response {
                                LogMetadataResponse::Empty => Vec::new(),
                                LogMetadataResponse::Available(metadata) => {
                                    metadata.iter().flat_map(LogMetadata::to_bytes).collect()
                                }
                            };
                            logical_packets.push(LogicalPacket {
                                attribute: Attribute::LogMetadata,
                                next: false,
                                chunk: 0,
                                size: payload.len() as u16,
                                payload,
                            });
                        }
                        PayloadData::Unknown { attribute, data } => {
                            logical_packets.push(LogicalPacket {
                                attribute,
                                next: false, // Will be fixed below
                                chunk: 0,
                                size: data.len() as u16,
                                payload: data,
                            });
                        }
                    }
                }

                // Set `next` flag on all packets except the last
                if let Some((last, rest)) = logical_packets.split_last_mut() {
                    for lp in rest {
                        lp.next = true;
                    }
                    last.next = false;
                }

                // Calculate the serialized size of all extended headers and payloads.
                let total_size: usize = logical_packets.iter().map(|lp| 4 + lp.payload.len()).sum();

                // Recorded PutData responses encode a word count relative to the
                // complete packet: packet_len / 4 - 3. The field is not the raw
                // payload length despite its historical `obj_count_words` name.
                let obj_count_words = if logical_packets.is_empty() {
                    0
                } else {
                    ((4 + total_size) / 4).saturating_sub(3) as u16
                };

                let header = DataHeader::new()
                    .with_packet_type(PacketType::PutData.into())
                    .with_reserved_flag(false)
                    .with_id(id)
                    .with_obj_count_words(obj_count_words);

                RawPacket::Data {
                    header,
                    logical_packets,
                }
            }
            Packet::GetData { attributes } => RawPacket::Ctrl {
                header: CtrlHeader::new()
                    .with_packet_type(PacketType::GetData.into())
                    .with_reserved_flag(false)
                    .with_id(id)
                    .with_attribute(attributes.raw()),
                payload: Vec::new(),
            },
            Packet::StartGraph { rate } => RawPacket::Ctrl {
                header: CtrlHeader::new()
                    .with_packet_type(PacketType::StartGraph.into())
                    .with_reserved_flag(false)
                    .with_id(id)
                    .with_attribute(rate as u16),
                payload: Vec::new(),
            },
            Packet::StopGraph => RawPacket::Ctrl {
                header: CtrlHeader::new()
                    .with_packet_type(PacketType::StopGraph.into())
                    .with_reserved_flag(false)
                    .with_id(id)
                    .with_attribute(0),
                payload: Vec::new(),
            },
            Packet::Accept { id: accept_id } => RawPacket::Ctrl {
                header: CtrlHeader::new()
                    .with_packet_type(PacketType::Accept.into())
                    .with_reserved_flag(false)
                    .with_id(accept_id)
                    .with_attribute(0),
                payload: Vec::new(),
            },
            Packet::Reject { id: reject_id } => RawPacket::Ctrl {
                header: CtrlHeader::new()
                    .with_packet_type(PacketType::Rejected.into())
                    .with_reserved_flag(false)
                    .with_id(reject_id)
                    .with_attribute(0),
                payload: Vec::new(),
            },
            Packet::NotReadable { id: nr_id } => RawPacket::Ctrl {
                header: CtrlHeader::new()
                    .with_packet_type(PacketType::NotReadable.into())
                    .with_reserved_flag(false)
                    .with_id(nr_id)
                    .with_attribute(0),
                payload: Vec::new(),
            },
            Packet::Connect => RawPacket::Ctrl {
                header: CtrlHeader::new()
                    .with_packet_type(PacketType::Connect.into())
                    .with_reserved_flag(false)
                    .with_id(id)
                    .with_attribute(0),
                payload: Vec::new(),
            },
            Packet::Disconnect => RawPacket::Ctrl {
                header: CtrlHeader::new()
                    .with_packet_type(PacketType::Disconnect.into())
                    .with_reserved_flag(false)
                    .with_id(id)
                    .with_attribute(0),
                payload: Vec::new(),
            },
            Packet::EnablePdMonitor => RawPacket::Ctrl {
                header: CtrlHeader::new()
                    .with_packet_type(PacketType::EnablePdMonitor.into())
                    .with_reserved_flag(false)
                    .with_id(id)
                    .with_attribute(PD_MONITOR_ENABLED_PARAMETER),
                payload: Vec::new(),
            },
            Packet::DisablePdMonitor => RawPacket::Ctrl {
                header: CtrlHeader::new()
                    .with_packet_type(PacketType::DisablePdMonitor.into())
                    .with_reserved_flag(false)
                    .with_id(id)
                    .with_attribute(PD_MONITOR_DISABLED_PARAMETER),
                payload: Vec::new(),
            },
            Packet::MemoryRead { address, size } => {
                // Build encrypted MemoryRead payload
                let encrypted_payload = auth::build_memory_read_payload(address, size);
                RawPacket::SimpleData {
                    header: DataHeader::from_bytes([PacketType::MemoryRead.into(), id, 0x01, 0x01]),
                    payload: encrypted_payload.to_vec(),
                }
            }
            Packet::StreamingAuth { credential } => {
                // Build encrypted StreamingAuth payload
                let encrypted_payload = auth::build_streaming_auth_payload(&credential);
                RawPacket::SimpleData {
                    header: DataHeader::from_bytes([PacketType::StreamingAuth.into(), id, 0x00, 0x02]),
                    payload: encrypted_payload.to_vec(),
                }
            }
            Packet::StreamingAuthResponse(result) => {
                // Response packets are typically not sent by client, but support for completeness
                let encrypted_payload = auth::encrypt_streaming_auth_response_payload(&result.decrypted_payload);
                let auth_header = StreamingAuthHeader::new()
                    .with_packet_type(PacketType::StreamingAuth.into())
                    .with_reserved_flag(false)
                    .with_id(auth::STREAMING_AUTH_RESPONSE_ID)
                    .with_attribute(result.attribute);
                RawPacket::Ctrl {
                    header: CtrlHeader::from_bytes(auth_header.into_bytes()),
                    payload: encrypted_payload.to_vec(),
                }
            }
            Packet::Generic(raw_packet) => raw_packet,
        };

        Ok(raw_packet)
    }
}

// Python support for Packet
#[cfg(feature = "python")]
impl<'py> pyo3::IntoPyObject<'py> for Packet {
    type Target = pyo3::PyAny;
    type Output = pyo3::Bound<'py, Self::Target>;
    type Error = pyo3::PyErr;

    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, Self::Error> {
        use pyo3::types::{PyDict, PyDictMethods};

        let dict = PyDict::new(py);
        match self {
            Packet::DataResponse { payloads } => {
                let inner = PyDict::new(py);
                inner.set_item("payloads", payloads.into_pyobject(py)?)?;
                dict.set_item("DataResponse", inner)?;
            }
            Packet::GetData { attributes } => {
                let inner = PyDict::new(py);
                inner.set_item("attribute_mask", attributes.raw())?;
                dict.set_item("GetData", inner)?;
            }
            Packet::StartGraph { rate } => {
                let inner = PyDict::new(py);
                inner.set_item("rate_index", rate as u16)?;
                dict.set_item("StartGraph", inner)?;
            }
            Packet::StopGraph => {
                dict.set_item("StopGraph", py.None())?;
            }
            Packet::Accept { id } => {
                let inner = PyDict::new(py);
                inner.set_item("id", id)?;
                dict.set_item("Accept", inner)?;
            }
            Packet::Reject { id } => {
                let inner = PyDict::new(py);
                inner.set_item("id", id)?;
                dict.set_item("Reject", inner)?;
            }
            Packet::NotReadable { id } => {
                let inner = PyDict::new(py);
                inner.set_item("id", id)?;
                dict.set_item("NotReadable", inner)?;
            }
            Packet::Connect => {
                dict.set_item("Connect", py.None())?;
            }
            Packet::Disconnect => {
                dict.set_item("Disconnect", py.None())?;
            }
            Packet::EnablePdMonitor => {
                dict.set_item("EnablePdMonitor", py.None())?;
            }
            Packet::DisablePdMonitor => {
                dict.set_item("DisablePdMonitor", py.None())?;
            }
            Packet::MemoryRead { address, size } => {
                let inner = PyDict::new(py);
                inner.set_item("address", address)?;
                inner.set_item("size", size)?;
                dict.set_item("MemoryRead", inner)?;
            }
            Packet::StreamingAuth { credential } => {
                let inner = PyDict::new(py);
                inner.set_item("credential", hex::encode(credential.as_bytes()))?;
                dict.set_item("StreamingAuth", inner)?;
            }
            Packet::StreamingAuthResponse(result) => {
                let inner = PyDict::new(py);
                inner.set_item("success", result.success)?;
                inner.set_item("auth_level", result.auth_level)?;
                inner.set_item("attribute", result.attribute)?;
                dict.set_item("StreamingAuthResponse", inner)?;
            }
            Packet::Generic(raw_packet) => {
                dict.set_item("Generic", raw_packet.into_pyobject(py)?)?;
            }
        }
        Ok(dict.into_any())
    }
}