tekhsi_rs 0.1.1

use crate::data::{
    AnalogSamples, AnalogWaveform, DigitalSamples, DigitalWaveform, IqMetaInfo, IqSamples,
    IqWaveform, Waveform,
};
use crate::errors::DecodeError;
use crate::tekscope::{WaveformHeader, WfmType};
use crate::validation::{expected_payload_len, is_header_valid};
use smol_str::SmolStr;
use tracing::warn;

use super::readers::{
    read_f32_chunks, read_f64_chunks, read_i8_chunks, read_i16_chunks, read_i32_chunks,
};

#[doc(hidden)]
pub fn decode_waveform_chunks(
    header: &WaveformHeader,
    chunks: &[Vec<u8>],
) -> Result<Waveform, DecodeError> {
    if !is_header_valid(header) {
        return Err(DecodeError::InvalidHeader);
    }

    match header.wfmtype() {
        WfmType::WfmtypeAnalog8 | WfmType::WfmtypeAnalog16 | WfmType::WfmtypeAnalogFloat => {
            Ok(Waveform::Analog(decode_analog_chunks(header, chunks)?))
        }
        WfmType::WfmtypeDigital8 | WfmType::WfmtypeDigital16 => {
            Ok(Waveform::Digital(decode_digital_chunks(header, chunks)?))
        }
        WfmType::WfmtypeAnalog16Iq | WfmType::WfmtypeAnalog32Iq => {
            Ok(Waveform::Iq(decode_iq_chunks(header, chunks)?))
        }
        WfmType::WfmtypeUnspecified => Err(DecodeError::UnsupportedWaveformType {
            waveform_type: header.wfmtype.to_string(),
        }),
    }
}

fn decode_analog_chunks(
    header: &WaveformHeader,
    chunks: &[Vec<u8>],
) -> Result<AnalogWaveform, DecodeError> {
    log_short_data(
        header,
        chunk_total_len(chunks),
        expected_payload_len(header),
    );

    let sample_count = header.noofsamples as usize;
    let values = match header.sourcewidth {
        1 => AnalogSamples::I8(read_i8_chunks(chunks, sample_count)),
        2 => AnalogSamples::I16(read_i16_chunks(chunks, sample_count)),
        4 => AnalogSamples::F32(read_f32_chunks(chunks, sample_count)),
        8 => AnalogSamples::F64(read_f64_chunks(chunks, sample_count)),
        _ => {
            return Err(DecodeError::UnsupportedSourceWidth {
                sourcewidth: header.sourcewidth,
                expected: &[1, 2, 4, 8],
            });
        }
    };

    Ok(AnalogWaveform {
        source_name: SmolStr::new(&header.sourcename),
        y_axis_values: values,
        y_axis_spacing: header.verticalspacing,
        y_axis_offset: header.verticaloffset,
        y_axis_units: SmolStr::new(&header.verticalunits),
        x_axis_spacing: header.horizontalspacing,
        x_axis_units: SmolStr::new(&header.horizontal_units),
        trigger_index: header.horizontalzeroindex,
    })
}

fn decode_digital_chunks(
    header: &WaveformHeader,
    chunks: &[Vec<u8>],
) -> Result<DigitalWaveform, DecodeError> {
    if header.sourcewidth != 1 && header.sourcewidth != 2 {
        return Err(DecodeError::UnsupportedSourceWidth {
            sourcewidth: header.sourcewidth,
            expected: &[1, 2],
        });
    }

    log_short_data(
        header,
        chunk_total_len(chunks),
        expected_payload_len(header),
    );

    let sample_count = header.noofsamples as usize;
    let values = match header.sourcewidth {
        // Observed on MSO 5 series with an 8-channel logic probe: sourcewidth=2,
        // with every other bit unused. We keep i16 support to handle this device.
        // It is unknown whether sourcewidth=1 is used for 4-channel probes or a
        // denser 8-channel encoding, but we keep i8 support for parity.
        1 => DigitalSamples::I8(read_i8_chunks(chunks, sample_count)),
        2 => DigitalSamples::I16(read_i16_chunks(chunks, sample_count)),
        _ => unreachable!("source width validated above"),
    };

    Ok(DigitalWaveform {
        source_name: SmolStr::new(&header.sourcename),
        y_axis_bytes: values,
        y_axis_units: SmolStr::new(&header.verticalunits),
        x_axis_spacing: header.horizontalspacing,
        x_axis_units: SmolStr::new(&header.horizontal_units),
        trigger_index: header.horizontalzeroindex,
    })
}

fn decode_iq_chunks(
    header: &WaveformHeader,
    chunks: &[Vec<u8>],
) -> Result<IqWaveform, DecodeError> {
    log_short_data(
        header,
        chunk_total_len(chunks),
        expected_payload_len(header),
    );

    let sample_count = header.noofsamples as usize;
    let values = match header.sourcewidth {
        1 => IqSamples::I8(read_i8_chunks(chunks, sample_count)),
        2 => IqSamples::I16(read_i16_chunks(chunks, sample_count)),
        4 => IqSamples::I32(read_i32_chunks(chunks, sample_count)),
        _ => {
            return Err(DecodeError::UnsupportedSourceWidth {
                sourcewidth: header.sourcewidth,
                expected: &[1, 2, 4],
            });
        }
    };

    let sample_rate = iq_sample_rate(header);

    Ok(IqWaveform {
        source_name: SmolStr::new(&header.sourcename),
        interleaved_iq: values,
        iq_axis_spacing: header.verticalspacing,
        iq_axis_offset: header.verticaloffset,
        iq_axis_units: SmolStr::new(&header.verticalunits),
        x_axis_spacing: header.horizontalspacing,
        x_axis_units: SmolStr::new(&header.horizontal_units),
        trigger_index: header.horizontalzeroindex,
        meta_info: IqMetaInfo {
            iq_center_frequency: header.iq_center_frequency,
            iq_fft_length: header.iq_fft_length,
            iq_resolution_bandwidth: header.iq_rbw,
            iq_span: header.iq_span,
            iq_window_type: SmolStr::new(&header.iq_window_type),
            iq_sample_rate: sample_rate,
        },
    })
}

#[inline]
fn log_short_data(header: &WaveformHeader, data_len: usize, expected_len: usize) {
    if data_len < expected_len {
        warn!(
            "short waveform data for {}: expected {} bytes, got {}",
            header.sourcename, expected_len, data_len
        );
    }
}

#[inline]
fn chunk_total_len(chunks: &[Vec<u8>]) -> usize {
    chunks.iter().map(|chunk| chunk.len()).sum()
}

fn iq_sample_rate(header: &WaveformHeader) -> f64 {
    match header.iq_window_type.as_str() {
        "Blackharris" => (header.iq_fft_length * header.iq_rbw) / 1.9,
        "Flattop2" => (header.iq_fft_length * header.iq_rbw) / 3.77,
        "Hanning" => (header.iq_fft_length * header.iq_rbw) / 1.44,
        "Hamming" => (header.iq_fft_length * header.iq_rbw) / 1.3,
        "Rectangle" => (header.iq_fft_length * header.iq_rbw) / 0.89,
        "Kaiserbessel" => (header.iq_fft_length * header.iq_rbw) / 2.23,
        _ => header.iq_span,
    }
}