sparslog 0.1.5

#![allow(clippy::missing_panics_doc)]
use std::collections::VecDeque;
use std::io::Write;
use std::net::SocketAddr;

use anyhow::anyhow;
use log::debug;

use rustradio::block::{Block, BlockRet};
use rustradio::blocks::{
    AddConst, BinarySlicer, FftFilter, FileSource, QuadratureDemod, RationalResampler,
    RtlSdrDecode, RtlSdrSource, TcpSource, ZeroCrossing,
};
use rustradio::graph::GraphRunner;
use rustradio::stream::ReadStream;
use rustradio::window::WindowType;
use rustradio::{Complex, Result, blockchain};

use std::sync::LazyLock;

pub const VERSION: &str = env!("CARGO_PKG_VERSION");

static REGISTRY: LazyLock<prometheus::Registry> = LazyLock::new(prometheus::Registry::new);

static WATTS: LazyLock<prometheus::Gauge> = LazyLock::new(|| {
    let metric =
        prometheus::Gauge::new("electricity_watts", "The instantenous watts used.").unwrap();
    REGISTRY.register(Box::new(metric.clone())).unwrap();
    metric
});

static BATTERY: LazyLock<prometheus::IntGauge> = LazyLock::new(|| {
    let metric = prometheus::IntGauge::new(
        "electricity_meter_battery",
        "Battery status of electricity meter.",
    )
    .unwrap();
    REGISTRY.register(Box::new(metric.clone())).unwrap();
    metric
});

static KWH: LazyLock<prometheus::Counter> = LazyLock::new(|| {
    let metric = prometheus::Counter::new("electricity_kwh", "Kwh counter.").unwrap();
    REGISTRY.register(Box::new(metric.clone())).unwrap();
    metric
});

static DECODES: LazyLock<prometheus::IntCounterVec> = LazyLock::new(|| {
    let metric = prometheus::IntCounterVec::new(
        prometheus::Opts::new("electricity_sparsnas_decodes", "Number of decodes."),
        &["status"],
    )
    .unwrap();
    REGISTRY.register(Box::new(metric.clone())).unwrap();
    metric
});

#[derive(clap::Parser, Debug)]
#[command(version=concat!(
        env!("CARGO_PKG_VERSION"),
        "\nCommit: ",
        env!("GIT_VERSION"),
        // This makes the build non-reproducible.
        // "\nBuilt at: ",
        // env!("BUILD_TIMESTAMP"),
        ), about)]
pub struct Opt {
    /// Serial number of the sensor.
    #[arg(short, long = "serial")]
    sensor_id: u32,

    /// Output file that will be appended to.
    #[arg(short, long = "output", default_value = "sparslog.csv")]
    output: String,

    /// Read 32bit Complex float stream by connecting with TCP.
    #[arg(short, long = "connect")]
    connect: Option<String>,

    /// Read I/Q from file. Can be combined with --rtlsdr.
    #[arg(short, long = "read")]
    read: Option<String>,

    /// Read from RTLSDR dongle, or in its format from a file when used with
    /// --read.
    #[arg(long = "rtlsdr")]
    rtlsdr: bool,

    /// Verbosity level.
    #[arg(short, default_value = "0")]
    pub verbose: usize,

    /// Input gain. Used with --rtlsdr.
    #[arg(long = "gain", default_value = "30")]
    gain: f32,

    /// Sample rate in file or with dongle.
    #[arg(long = "sample_rate", default_value_t = 1_024_000)]
    sample_rate: u32,

    /// Frequency to tune to, in Hz.
    #[arg(long = "freq", default_value_t = 868_000_000)]
    freq: u64,

    /// FSK offset value.
    #[arg(long = "offset", default_value = "0.4")]
    offset: f32,

    /// Run multithreaded.
    #[arg(long)]
    pub multithread: bool,

    /// Prometheus gateway server to push metrics to.
    #[arg(long, requires = "where_")]
    prometheus: Option<String>,

    /// Name for location we're measuring.
    #[arg(long = "where")]
    where_: Option<String>,
}

#[derive(rustradio::rustradio_macros::Block)]
#[rustradio(new, custom_name)]
struct Decode {
    #[rustradio(in)]
    src: ReadStream<u8>,
    sensor_id: u32,
    output: String,

    #[rustradio(default)]
    history: VecDeque<u8>,
}

impl Decode {
    fn custom_name(&self) -> &'static str {
        let _ = self;
        "Sparsnäs decoder"
    }
}

#[allow(clippy::cast_precision_loss)]
fn f32_to_i32(value: f32) -> anyhow::Result<i32> {
    let value = f64::from(value);
    if value.is_finite() && value >= f64::from(i32::MIN) && value <= f64::from(i32::MAX) {
        #[allow(clippy::cast_possible_truncation)]
        Ok(value as i32)
    } else {
        Err(anyhow!("invalid conversion from {value} to i32"))
    }
}

#[allow(clippy::cast_precision_loss)]
fn f32_to_usize(value: f32) -> anyhow::Result<usize> {
    let value = f64::from(value);
    if value.is_finite() && value >= 0.0 && value <= usize::MAX as f64 {
        #[allow(clippy::cast_possible_truncation)]
        #[allow(clippy::cast_sign_loss)]
        Ok(value as usize)
    } else {
        Err(anyhow!("invalid conversion from {value} to usize"))
    }
}

fn bits2byte(data: &[u8]) -> u8 {
    assert_eq!(data.len(), 8);
    (data[0] << 7)
        | (data[1] << 6)
        | (data[2] << 5)
        | (data[3] << 4)
        | (data[4] << 3)
        | (data[5] << 2)
        | (data[6] << 1)
        | data[7]
}
fn calc_crc(mut s: u8, mut reg: u16) -> u16 {
    let poly: u16 = 0x8005;
    for _i in 0..8 {
        let regbit = reg & 0x8000 != 0;
        let databit = s & 0x80 != 0;
        if regbit ^ databit {
            reg = (reg << 1) ^ poly;
        } else {
            reg <<= 1;
        }
        s <<= 1;
    }
    reg
}

fn crc16(input: &[u8], expected: u16) -> bool {
    let mut checksum = 0xffffu16;
    for i in input {
        checksum = calc_crc(*i, checksum);
    }
    //eprintln!("Got checksum {:04x}, want {:04x}", checksum, expected);
    checksum == expected
}

// packet: from length to and including the CRC.
fn fix_packet(packet: &[u8]) -> Vec<u8> {
    let crc = (u16::from(packet[packet.len() - 2]) << 8) | u16::from(packet[packet.len() - 1]);
    if crc16(&packet[..packet.len() - 2], crc) {
        return packet.to_vec();
    }
    for i in 0..(packet.len() * 8) {
        let mut test = packet.to_vec();
        let bit = 1 << (i % 8);
        test[i / 8] ^= bit;
        if crc16(&test[..packet.len() - 2], crc) {
            return test.clone();
        }
    }
    packet.to_vec()
}

fn parsepacket(packet: &[u8], sensor_id: u32) -> String {
    assert_eq!(packet.len(), 20);
    //let sensor = packet[0];
    //let app = packet[1];
    let packet = fix_packet(packet);

    // This is the correct packet.
    println!("Packet: {packet:02x?}");

    let sensor_id_sub = {
        let magic = 0x5D38_E8CB;
        if sensor_id >= magic {
            sensor_id - magic
        } else {
            4_294_967_295 - (magic - sensor_id - 1)
        }
    };
    let enc_key = [
        ((sensor_id_sub >> 24) & 0xff) as u8,
        (sensor_id_sub & 0xff) as u8,
        ((sensor_id_sub >> 8) & 0xff) as u8,
        0x47u8,
        ((sensor_id_sub >> 16) & 0xff) as u8,
    ];
    let mut dec = Vec::new();
    for i in 0..13 {
        dec.push(packet[i + 5] ^ enc_key[i % 5]);
    }
    //println!("Decoded: {:02x?}", dec);
    //let mut prep = vec![0x11];
    //prep.extend(&packet[..packet.len()-2]);
    let crc = (u16::from(packet[packet.len() - 2]) << 8) | u16::from(packet[packet.len() - 1]);
    let crc_ok = crc16(&packet[..packet.len() - 2], crc);

    let seq = (u16::from(dec[4]) << 8) | u16::from(dec[5]);
    let effect = (u16::from(dec[6]) << 8) | u16::from(dec[7]);
    let wh = (u32::from(dec[8]) << 24)
        | (u32::from(dec[9]) << 16)
        | (u32::from(dec[10]) << 8)
        | u32::from(dec[11]);
    let kwh = format!("{}.{:03}", wh / 1000, wh % 1000);
    let battery = dec[12];

    let watt = 3600.0 * 1024.0 / f32::from(effect);

    let now = std::time::SystemTime::now()
        .duration_since(std::time::SystemTime::UNIX_EPOCH)
        .expect("Time went backwards")
        .as_secs();

    if crc_ok {
        WATTS.set(watt.into());
        BATTERY.set(battery.into());
        if let Ok(v) = kwh.parse() {
            let old = KWH.get();
            if old > v {
                KWH.reset();
                KWH.inc_by(v);
            } else {
                KWH.inc_by(v - old);
            }
        }
    }
    let status = if crc_ok { "OK" } else { "BAD" };
    DECODES.with_label_values(&[status]).inc();
    format!("{now},{seq},{watt:.3},{kwh},{battery},{status}")
}

impl Block for Decode {
    fn work(&mut self) -> Result<BlockRet<'_>> {
        let cac = [
            1u8, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0,
            0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1,
        ];
        let (input, _) = self.src.read_buf()?;
        if input.is_empty() {
            return Ok(BlockRet::WaitForStream(&self.src, 1));
        }
        //eprintln!("Decode got {}", input.available());
        self.history.extend(input.iter());
        {
            let n = input.len();
            input.consume(n);
        }

        let packet_bits_len = cac.len() + 19 * 8;
        //let cac = vec![1,0,1,0,1,0,1,0,1,0,1];
        let n = self.history.len();
        //println!("Called with {n}");
        if n < packet_bits_len {
            //debug!("{} < {} len, sleeping", n, cac.len());
            return Ok(BlockRet::WaitForStream(&self.src, packet_bits_len - n));
        }
        //println!("Running on data size {n}");
        let input = &self.history;
        for i in 0..(n - packet_bits_len) {
            let equal = cac
                .iter()
                .zip(input.range(i..(i + cac.len())))
                .all(|(a, b)| a == b);
            //if &cac == input.range(i..(i + cac.len())) {
            if equal {
                debug!("Found CAC");
                let bits = &input
                    .range(i..(i + cac.len() + 19 * 8))
                    .copied()
                    .collect::<Vec<u8>>();
                let mut bytes = Vec::new();
                for j in (0..bits.len()).step_by(8) {
                    bytes.push(bits2byte(&bits[j..j + 8]));
                }
                //println!("bytes: {:02x?}", bytes);
                let packet = &bytes[4..];
                //println!("packet: {:02x?}", packet);
                let parsed = parsepacket(packet, self.sensor_id);
                std::fs::OpenOptions::new()
                    .append(true)
                    .create(true)
                    .open(&self.output)
                    .map_err(|e| -> rustradio::Error { e.into() })?
                    .write_all(format!("{parsed}\n").as_bytes())
                    .map_err(|e| -> rustradio::Error { e.into() })?;
                println!("{parsed}");
            }
        }
        self.history
            .drain(0..(self.history.len() - packet_bits_len));
        Ok(BlockRet::Again)
    }
}

static HOSTNAME: LazyLock<String> = LazyLock::new(|| {
    std::env::var("HOSTNAME")
        .ok()
        .filter(|s| !s.is_empty())
        .unwrap_or_else(|| {
            std::fs::read_to_string("/etc/hostname")
                .map_or_else(|_| "unknown".to_owned(), |s| s.trim_end().to_owned())
        })
});

fn push_metrics(gw: &str, wh: &str, serial: u32) -> anyhow::Result<()> {
    debug!("Pushing metrics");
    let grouping = std::collections::HashMap::from([
        ("instance".to_string(), (*HOSTNAME).clone()),
        ("where".to_string(), wh.to_string()),
        ("serial".to_string(), serial.to_string()),
    ]);

    prometheus::push_metrics(
        "sparslog", // job name
        grouping,   // grouping labels
        gw,
        REGISTRY.gather(),
        None, // optional basic auth
    )?;
    Ok(())
}

/// Create the graph to decode sparsnäs.
///
/// # Errors
///
/// If given incompatible cmdline options.
pub fn create_graph(graph: &mut (impl GraphRunner + ?Sized), opt: &Opt) -> anyhow::Result<()> {
    if let Some(prom) = &opt.prometheus {
        let prom = prom.clone();
        let wh = opt.where_.clone().expect("Can't happen: clap promised!");
        let serial = opt.sensor_id;
        std::thread::Builder::new()
            .name("prometheus-pusher".to_string())
            .spawn(move || {
                loop {
                    if let Err(err) = push_metrics(&prom, &wh, serial) {
                        eprintln!("Failed to push prometheus metrics: {err}");
                    }
                    std::thread::sleep(std::time::Duration::from_mins(1));
                }
            })
            .expect("spawn prometheus pusher thread");
    }
    // Source.
    let src = {
        if let Some(connect) = &opt.connect {
            if opt.read.is_some() {
                return Err(anyhow::Error::msg("-c and -r can't be combined"));
            }
            let sa: SocketAddr = connect.parse()?;
            let host = format!("{}", sa.ip());
            let port = sa.port();
            println!("Connecting to host {host} port {port}");
            blockchain![graph, prev, TcpSource::<Complex>::new(&host, port)?]
        } else if let Some(read) = &opt.read {
            if opt.rtlsdr {
                blockchain![
                    graph,
                    prev,
                    FileSource::<u8>::new(read)?,
                    RtlSdrDecode::new(prev),
                ]
            } else {
                blockchain![graph, prev, FileSource::<Complex>::new(read)?]
            }
        } else if opt.rtlsdr {
            blockchain![
                graph,
                prev,
                RtlSdrSource::new(opt.freq, opt.sample_rate, f32_to_i32(opt.gain)?)?,
                RtlSdrDecode::new(prev),
            ]
        } else {
            return Err(anyhow::Error::msg(
                "Need to provide either -r, -c, or --rtlsdr",
            ));
        }
    };

    #[allow(clippy::cast_precision_loss)]
    let samp_rate = opt.sample_rate as f32;
    let samp_rate_2 = 200_000.0;
    let baud = 38383.5;

    let prev = src;
    // Resample.
    let prev = blockchain![
        graph,
        prev,
        // TODO: doing filtering in multiple steps, with a decimating FIR filter, would
        // probably be more CPU efficient.
        FftFilter::new(
            prev,
            rustradio::fir::low_pass_complex(samp_rate, 50000.0, 10000.0, &WindowType::Hamming)
        ),
        RationalResampler::new(prev, f32_to_usize(samp_rate_2)?, f32_to_usize(samp_rate)?)?,
        QuadratureDemod::new(prev, 1.0),
        AddConst::new(prev, opt.offset),
        ZeroCrossing::new(prev, samp_rate_2 / baud, 0.1),
        BinarySlicer::new(prev),
    ];

    // Decode.
    let decode = Box::new(Decode::new(prev, opt.sensor_id, opt.output.clone()));
    graph.add(decode);
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn convert_i32() -> anyhow::Result<()> {
        for (i, o) in [
            (0.0, 0),
            (-1.0, -1),
            (-1.1, -1),
            (-1.9, -1),
            (1.0, 1),
            (1.1, 1),
            (1.9, 1),
        ] {
            assert_eq!(f32_to_i32(i)?, o);
        }
        Ok(())
    }

    #[test]
    fn convert_usize() -> anyhow::Result<()> {
        for (i, o) in [
            (0.0, Some(0)),
            (-1.0, None),
            (-1.1, None),
            (1.0, Some(1)),
            (1.1, Some(1)),
            (1.9, Some(1)),
        ] {
            match o {
                None => {
                    assert!(f32_to_usize(i).is_err());
                }
                Some(v) => {
                    assert_eq!(f32_to_usize(i)?, v);
                }
            }
        }
        Ok(())
    }

    #[test]
    fn decode() {
        let packet = vec![
            0x11, 0xa1, 0x38, 0x07, 0x0e, 0xa2, 0xde, 0x29, 0xe6, 0x8b, 0x1a, 0xfd, 0x74, 0x47,
            0xcf, 0xf2, 0x14, 0x80, 0x23, 0x7b,
        ];
        let got = parsepacket(&packet, 576_929);
        let want = ",17592,330.560,20.674,100,OK";
        assert!(got.ends_with(want), "got: {got}, want {want}");

        // With one bitflip.
        let packet = vec![
            0x11, 0xa1, 0x38, 0x07, 0x0e, 0xa2, 0xde, 0x29, 0xe7, 0x8b, 0x1a, 0xfd, 0x74, 0x47,
            0xcf, 0xf2, 0x14, 0x80, 0x23, 0x7b,
        ];
        let got = parsepacket(&packet, 576_929);
        let want = ",17592,330.560,20.674,100,OK";
        assert!(got.ends_with(want), "got: {got}, want {want}");

        // With two bitflips.
        let packet = vec![
            0x11, 0xa1, 0x38, 0x07, 0x0e, 0xa2, 0xdf, 0x29, 0xe6, 0x8b, 0x1a, 0xfd, 0x74, 0x47,
            0xcf, 0xf2, 0x14, 0x80, 0x23, 0x7a,
        ];
        let got = parsepacket(&packet, 576_929);
        let want = ",17592,330.560,20.674,100,BAD";
        assert!(got.ends_with(want), "got: {got}, want {want}");
    }
}