//! mmx150-probe — reverse-engineering tool for the MMX 150 dongle protocol.
//!
//! Phase 1 workflow: send candidate output reports to the dongle and observe
//! what input reports come back, correlating fields with headset on/off,
//! battery charge level and button presses.

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();
    match args.first().map(String::as_str) {
        Some("list") => cmd_list(&args[1..]),
        Some("scan") => cmd_scan(),
        Some("send") => cmd_send(&args[1..]),
        Some("sweep") => cmd_sweep(&args[1..]),
        Some("status") => cmd_status_dump(),
        _ => {
            eprintln!("usage: mmx150-probe <list|scan|send|sweep|status>");
            ExitCode::FAILURE
        }
    }
}

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 Ok(report_id) = u8::from_str_radix(rest[0].trim_start_matches("0x"), 16) else {
        eprintln!("bad report id");
        return ExitCode::FAILURE;
    };
    let mut data = vec![report_id];
    for part in &rest[1..] {
        let bytes = parse_hex(part);
        if bytes.is_empty() {
            eprintln!("bad hex: {part}");
            return ExitCode::FAILURE;
        }
        data.extend(bytes);
    }

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

    println!(
        ">> 0x{report_id:02x} [{}B] {}",
        data.len() - 1,
        hex(&data[1..])
    );
    match send_and_read(&mut hid, &data) {
        Some(resp) => println!("<< [{}B] {}", resp.len(), hex(&resp)),
        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;
        }
    }
    read_one_response(hid, Duration::from_millis(400))
}

/// Wait up to `timeout` for a single input report and return it.
fn read_one_response(hid: &mut Hidraw, timeout: Duration) -> Option<Vec<u8>> {
    let deadline = std::time::Instant::now() + timeout;
    loop {
        let now = std::time::Instant::now();
        if now >= deadline {
            return None;
        }
        let mut set = unsafe { std::mem::zeroed::<libc::fd_set>() };
        unsafe { libc::FD_ZERO(&mut set) };
        unsafe { libc::FD_SET(hid.as_raw_fd(), &mut set) };
        let remaining = deadline.saturating_duration_since(now).as_millis();
        let mut to = libc::timeval {
            tv_sec: (remaining / 1000) as i64,
            tv_usec: ((remaining % 1000) * 1000) as i64,
        };
        let n = unsafe {
            libc::select(
                hid.as_raw_fd() + 1,
                &mut set,
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                &mut to,
            )
        };
        if n <= 0 {
            return None;
        }
        let mut buf = [0u8; 256];
        match hid.read(&mut buf) {
            Ok(len) => return Some(buf[..len].to_vec()),
            Err(_) => return None,
        }
    }
}

/// Sweep a set of candidate commands against the dongle, printing responses
/// that differ from the all-zero baseline. Reports only non-trivial responses.
///
/// usage: mmx150-probe sweep <report_id> <B> ...
///   <report_id> is the output report id to use,
///   <B> is the report payload length (bytes, excluding the id byte).
fn cmd_status_dump() -> ExitCode {
    let Some(dev) = mmx150::hid::find_mmx150().unwrap_or(None) else {
        eprintln!("MMX 150 dongle not found");
        return ExitCode::FAILURE;
    };
    let hid = match Hidraw::open(&dev.devnode) {
        Ok(h) => h,
        Err(e) => {
            eprintln!("open failed: {e}");
            return ExitCode::FAILURE;
        }
    };
    let mut hid = hid;

    // Each row: label, report id, command payload.
    let cmds: Vec<(&str, u8, Vec<u8>)> = vec![
        ("0x2a echo/sel=0x0b", 0x2a, {
            let mut c = vec![0u8; 60];
            c[0] = 0x09;
            c[1] = 0x22;
            c[2] = 0x0b;
            c
        }),
        ("0x2a printf family 09 22", 0x2a, {
            let mut c = vec![0u8; 60];
            c[0] = 0x09;
            c[1] = 0x22;
            c[2] = 0x01;
            c
        }),
        ("0x2a eq family 09 00", 0x2a, {
            let mut c = vec![0u8; 60];
            c[0] = 0x09;
            c[1] = 0x00;
            c[2] = 0x0b;
            c
        }),
        ("0x3a sel=0x0b", 0x3a, {
            let mut c = vec![0u8; 60];
            c[0] = 0x09;
            c[1] = 0x22;
            c[2] = 0x0b;
            c
        }),
        ("0x3a sel=0x05", 0x3a, {
            let mut c = vec![0u8; 60];
            c[0] = 0x09;
            c[1] = 0x22;
            c[2] = 0x05;
            c
        }),
        ("0x3a sel=0x01", 0x3a, {
            let mut c = vec![0u8; 60];
            c[0] = 0x09;
            c[1] = 0x22;
            c[2] = 0x01;
            c
        }),
        ("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 resp = send_and_read(&mut hid, &frame);
        println!(
            "{label:<26} >>0x{rid:02x} << {}",
            resp.map(|b| hex(&b)).unwrap_or_else(|| "<none>".into())
        );
        std::thread::sleep(Duration::from_millis(120));
    }
    ExitCode::SUCCESS
}

fn cmd_sweep(rest: &[String]) -> ExitCode {
    if rest.len() < 1 {
        eprintln!("usage: mmx150-probe sweep <report_id> <payload_len>");
        return ExitCode::FAILURE;
    }
    let Ok(report_id) = u8::from_str_radix(rest[0].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(dev) = mmx150::hid::find_mmx150().unwrap_or(None) else {
        eprintln!("MMX 150 dongle not found");
        return ExitCode::FAILURE;
    };
    let hid = match Hidraw::open(&dev.devnode) {
        Ok(h) => h,
        Err(e) => {
            eprintln!("open failed: {e}");
            return ExitCode::FAILURE;
        }
    };
    let mut hid = hid;

    // Baseline: all-zero command.
    let base = {
        let mut cmd = vec![report_id];
        cmd.extend(std::iter::repeat(0u8).take(plen));
        send_and_read(&mut hid, &cmd)
    };
    println!(
        "baseline (all zero): {}",
        base.as_ref()
            .map(|b| hex(b.as_slice()))
            .unwrap_or_else(|| "<none>".into())
    );

    // 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(std::iter::repeat(0u8).take(plen));
            cmd[idx] = val;
            match send_and_read(&mut hid, &cmd) {
                Some(resp) if Some(&resp) != base.as_ref() => {
                    println!("cmd[{idx:02}]=0x{val:02x} -> {}", hex(&resp));
                }
                _ => {}
            }
        }
    }
    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() % 2 != 0 {
        return Vec::new();
    }
    s.as_bytes()
        .chunks(2)
        .map(|c| {
            let hi = (c[0] as char).to_digit(16)?;
            let lo = (c[1] as char).to_digit(16)?;
            Some((hi * 16 + lo) as u8)
        })
        .collect::<Option<Vec<u8>>>()
        .unwrap_or_default()
}