//! mmx150-probe — reverse-engineering tool for the MMX 150 dongle protocol.
//!
//! Sends candidate output reports to the dongle and prints what input
//! reports come back. Kept for future protocol work: battery status, for
//! example, is not exposed by any command found so far.

use std::process::ExitCode;
use std::time::Duration;

use mmx150::hid::{find_devices_by_vendor, Hidraw, BEYERDYNAMIC_VID};

fn main() -> ExitCode {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let rest: &[String] = args.get(1..).unwrap_or(&[]);
    match args.first().map(String::as_str) {
        Some("list") => cmd_list(rest),
        Some("scan") => cmd_scan(),
        Some("send") => cmd_send(rest),
        Some("sweep") => cmd_sweep(rest),
        Some("status") => cmd_status_dump(),
        Some("selenum") => cmd_sel_enum(rest),
        _ => {
            eprintln!("usage: mmx150-probe <list|scan|send|sweep|status|selenum>");
            ExitCode::FAILURE
        }
    }
}

/// Find and open the dongle, printing an error and returning `None` if it
/// is missing or cannot be opened.
fn open_dongle() -> Option<Hidraw> {
    let Some(dev) = mmx150::hid::find_mmx150().unwrap_or(None) else {
        eprintln!("MMX 150 dongle not found");
        return None;
    };
    match Hidraw::open(&dev.devnode) {
        Ok(h) => Some(h),
        Err(e) => {
            eprintln!("open failed: {e}");
            None
        }
    }
}

fn cmd_list(_rest: &[String]) -> ExitCode {
    match find_devices_by_vendor(BEYERDYNAMIC_VID) {
        Ok(devices) => {
            for d in &devices {
                println!(
                    "{:04x}:{:04x} {} {}",
                    d.vid,
                    d.pid,
                    d.devnode.display(),
                    d.name
                );
            }
            if devices.is_empty() {
                eprintln!("no beyerdynamic hidraw devices found");
                return ExitCode::FAILURE;
            }
            ExitCode::SUCCESS
        }
        Err(e) => {
            eprintln!("error: {e}");
            ExitCode::FAILURE
        }
    }
}

fn cmd_scan() -> ExitCode {
    let Some(dev) = mmx150::hid::find_mmx150().unwrap_or(None) else {
        eprintln!("MMX 150 dongle not found");
        return ExitCode::FAILURE;
    };
    println!("device: {} at {}", dev.name, dev.devnode.display());
    let hid = match Hidraw::open(&dev.devnode) {
        Ok(h) => h,
        Err(e) => {
            eprintln!("open failed: {e}");
            return ExitCode::FAILURE;
        }
    };

    // Known report sizes from the report descriptor:
    // 62-byte reports: 0x24/0x27, 0x80/0x81, 0x34/0x37
    // 60-byte reports: 0x2a/0x2b, 0x3a/0x3b
    // 16-byte reports: 0x2c/0x2d, 0x3c/0x3d
    for (report_id, len) in [
        (0x24, 63usize),
        (0x34, 63),
        (0x80, 63),
        (0x3a, 61),
        (0x3b, 61),
        (0x2a, 61),
        (0x2b, 61),
        (0x3d, 17),
        (0x2d, 17),
    ] {
        match hid.get_feature(report_id, len) {
            Ok(buf) => println!("FEATURE 0x{report_id:02x}: {}", hex(&buf)),
            Err(e) => println!("FEATURE 0x{report_id:02x}: {e}"),
        }
    }
    for (report_id, len) in [
        (0x27, 63usize),
        (0x37, 63),
        (0x81, 63),
        (0x3c, 17),
        (0x3d, 17),
        (0x2c, 17),
        (0x2d, 17),
    ] {
        match hid.get_input(report_id, len) {
            Ok(buf) => println!("GINPUT 0x{report_id:02x}: {}", hex(&buf)),
            Err(e) => println!("GINPUT 0x{report_id:02x}: {e}"),
        }
    }
    ExitCode::SUCCESS
}

fn cmd_send(rest: &[String]) -> ExitCode {
    if rest.len() < 2 {
        eprintln!("usage: mmx150-probe send <report_id> <hex bytes...>");
        return ExitCode::FAILURE;
    }
    let Some(id_text) = rest.first() else {
        return ExitCode::FAILURE;
    };
    let Ok(report_id) = u8::from_str_radix(id_text.trim_start_matches("0x"), 16) else {
        eprintln!("bad report id");
        return ExitCode::FAILURE;
    };
    let mut data = vec![report_id];
    for part in rest.get(1..).unwrap_or(&[]) {
        let bytes = parse_hex(part);
        if bytes.is_empty() {
            eprintln!("bad hex: {part}");
            return ExitCode::FAILURE;
        }
        data.extend(bytes);
    }

    let Some(mut hid) = open_dongle() else {
        return ExitCode::FAILURE;
    };

    println!(
        ">> 0x{report_id:02x} [{}B] {}",
        data.len().saturating_sub(1),
        hex(data.get(1..).unwrap_or(&[]))
    );
    match send_and_read(&mut hid, &data) {
        Some(response) => println!("<< [{}B] {}", response.len(), hex(&response)),
        None => println!("<< no response"),
    }
    ExitCode::SUCCESS
}

/// Send one full output report (including report-id byte as data[0]) via the
/// interrupt OUT endpoint and return the first input report received, if any.
fn send_and_read(hid: &mut Hidraw, data: &[u8]) -> Option<Vec<u8>> {
    match hid.write(data) {
        Ok(_) => {}
        Err(e) => {
            eprintln!("write() failed: {e}");
            return None;
        }
    }
    hid.read_with_timeout(Duration::from_millis(400))
        .ok()
        .flatten()
}

/// Send a fixed set of diagnostic commands (one per known report family)
/// and print their responses side by side.
fn cmd_status_dump() -> ExitCode {
    let Some(mut hid) = open_dongle() else {
        return ExitCode::FAILURE;
    };

    // Each row: label, report id, command payload.
    let cmds: Vec<(&str, u8, Vec<u8>)> = vec![
        ("0x2a echo/sel=0x0b", 0x2a, cmd_payload(&[0x09, 0x22, 0x0b])),
        (
            "0x2a printf family 09 22",
            0x2a,
            cmd_payload(&[0x09, 0x22, 0x01]),
        ),
        (
            "0x2a eq family 09 00",
            0x2a,
            cmd_payload(&[0x09, 0x00, 0x0b]),
        ),
        ("0x3a sel=0x0b", 0x3a, cmd_payload(&[0x09, 0x22, 0x0b])),
        ("0x3a sel=0x05", 0x3a, cmd_payload(&[0x09, 0x22, 0x05])),
        ("0x3a sel=0x01", 0x3a, cmd_payload(&[0x09, 0x22, 0x01])),
        ("0x2d signature", 0x2d, vec![0u8; 16]),
        ("0x3d signature", 0x3d, vec![0u8; 16]),
    ];

    for (label, rid, payload) in cmds {
        let mut frame = vec![rid];
        frame.extend(payload);
        let response = send_and_read(&mut hid, &frame);
        println!(
            "{label:<26} >>0x{rid:02x} << {}",
            response
                .as_deref()
                .map_or_else(|| "<none>".to_string(), hex)
        );
        std::thread::sleep(Duration::from_millis(120));
    }
    ExitCode::SUCCESS
}

/// A 60-byte command payload with `bytes` at the front and zero padding.
fn cmd_payload(bytes: &[u8]) -> Vec<u8> {
    let mut payload = vec![0u8; 60];
    let n = bytes.len().min(60);
    if let (Some(slot), Some(source)) = (payload.get_mut(..n), bytes.get(..n)) {
        slot.copy_from_slice(source);
    }
    payload
}

/// Enumerate many sub-command selectors across a report family, dumping any
/// response whose payload carries a real data byte (not just an echo of the
/// selector / a constant). This surfaces battery and state registers.
///
/// usage: mmx150-probe selenum `[report_id]`  (default: 0x3a)
fn cmd_sel_enum(rest: &[String]) -> ExitCode {
    let report_id = rest
        .first()
        .and_then(|s| u8::from_str_radix(s.trim_start_matches("0x"), 16).ok())
        .unwrap_or(0x3a);

    let Some(mut hid) = open_dongle() else {
        return ExitCode::FAILURE;
    };

    println!("sweeping report 0x{report_id:02x}, selector range 0x00..0xff");
    for sel in 0x00u8..=0xffu8 {
        let mut frame = vec![report_id];
        frame.extend(cmd_payload(&[0x09, 0x22, sel]));
        let response = send_and_read(&mut hid, &frame);

        // Only report responses containing a "real" data byte beyond the
        // id/echo/constant region that could be battery/state.
        if let Some(bytes) = response {
            if bytes
                .get(6..)
                .is_some_and(|tail| tail.iter().any(|&x| x != 0x00))
            {
                println!("sel={sel:02x} resp={}", hex(&bytes));
            }
            std::thread::sleep(Duration::from_millis(8));
        }
    }
    ExitCode::SUCCESS
}

fn cmd_sweep(rest: &[String]) -> ExitCode {
    if rest.is_empty() {
        eprintln!("usage: mmx150-probe sweep <report_id> <payload_len>");
        return ExitCode::FAILURE;
    }
    let Some(id_text) = rest.first() else {
        return ExitCode::FAILURE;
    };
    let Ok(report_id) = u8::from_str_radix(id_text.trim_start_matches("0x"), 16) else {
        eprintln!("bad report id");
        return ExitCode::FAILURE;
    };
    let plen: usize = rest.get(1).and_then(|s| s.parse().ok()).unwrap_or(60);

    let Some(mut hid) = open_dongle() else {
        return ExitCode::FAILURE;
    };

    // Baseline: all-zero command.
    let baseline = {
        let mut cmd = vec![report_id];
        cmd.extend(vec![0u8; plen]);
        send_and_read(&mut hid, &cmd)
    };
    println!(
        "baseline (all zero): {}",
        baseline
            .as_deref()
            .map_or_else(|| "<none>".to_string(), hex)
    );

    // Try setting individual low nibbles / bytes across the header to probe
    // which command byte selects a status register.
    for idx in 1..plen.clamp(1, 12) {
        for val in [0x01u8, 0x22, 0x2c, 0x2f, 0x80, 0x5a, 0xff] {
            let mut cmd = vec![report_id];
            cmd.extend(vec![0u8; plen]);
            if let Some(slot) = cmd.get_mut(idx) {
                *slot = val;
            }
            match send_and_read(&mut hid, &cmd) {
                Some(response) if Some(&response) != baseline.as_ref() => {
                    println!("cmd[{idx:02}]=0x{val:02x} -> {}", hex(&response));
                }
                _ => {}
            }
        }
    }
    ExitCode::SUCCESS
}

fn hex(bytes: &[u8]) -> String {
    bytes
        .iter()
        .map(|b| format!("{b:02x}"))
        .collect::<Vec<_>>()
        .join(" ")
}

fn parse_hex(s: &str) -> Vec<u8> {
    let s = s.trim_start_matches("0x");
    if !s.len().is_multiple_of(2) {
        return Vec::new();
    }
    s.as_bytes()
        .chunks_exact(2)
        .map(|pair| {
            let [hi, lo] = pair else { return None };
            let hi = char::from(*hi).to_digit(16)?;
            let lo = char::from(*lo).to_digit(16)?;
            u8::try_from((hi << 4) | lo).ok()
        })
        .collect::<Option<Vec<u8>>>()
        .unwrap_or_default()
}

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

    #[test]
    fn parses_hex_pairs() {
        assert_eq!(parse_hex("0922"), vec![0x09, 0x22]);
        assert_eq!(parse_hex("0x0922"), vec![0x09, 0x22]);
        assert_eq!(parse_hex("ff"), vec![0xff]);
        assert_eq!(parse_hex("FF"), vec![0xff]);
        assert_eq!(parse_hex("0A0b"), vec![0x0a, 0x0b]);
    }

    #[test]
    fn rejects_odd_length() {
        assert!(parse_hex("092").is_empty());
    }

    #[test]
    fn rejects_invalid_digits() {
        assert!(parse_hex("zz").is_empty());
        assert!(parse_hex("0g").is_empty());
    }

    #[test]
    fn empty_input_yields_empty() {
        assert!(parse_hex("").is_empty());
        assert!(parse_hex("0x").is_empty());
    }
}