TwoFootBot-rust/src/pingEsp32.rs

use std::io::{BufRead, BufReader, Read, Write};
use std::net::TcpStream;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{self, Receiver, Sender, RecvTimeoutError};
use std::sync::Arc;
use std::thread::{self, JoinHandle};
use std::time::Duration;

use std::f32::consts::PI;

pub const ESP32_IP: &str = "192.168.4.1";
pub const ESP32_PORT: u16 = 3333;
pub const NUM_SERVOS: usize = 8;

const TELEMETRY_MAGIC: [u8; 2] = [0xA5, 0x5A];
const TELEMETRY_FRAME_LEN: usize = 16;

/// 从 ESP32 收到的一条消息（可包含 ASCII 文本，也可包含 36 维 IMU 数据）
#[derive(Debug, Clone)]
pub struct Esp32Message {
    text: String,
    telemetry: Option<[f32; 17]>,
}

impl Esp32Message {
    pub fn new_text(text: impl Into<String>) -> Self {
        Self {
            text: text.into(),
            telemetry: None,
        }
    }

    pub fn new_telemetry(data: [f32; 17]) -> Self {
        Self {
            text: "IMU".to_string(),
            telemetry: Some(data),
        }
    }

    /// 是否是数据推送（二进制 IMU 帧）
    pub fn is_data(&self) -> bool { self.telemetry.is_some() }

    /// 是否是应答，比如 "ESP32: OK ..."
    pub fn is_reply(&self) -> bool {
        self.telemetry.is_none() && self.text.starts_with("ESP32:")
    }

    pub fn text(&self) -> &str { &self.text }

    /// 返回 36 维 f32 数组；其余位置按 0.0 填充。
    pub fn telemetry(&self) -> Option<[f32; 17]> { self.telemetry }
}

// ==================== 双向连接 ====================

pub struct Link {
    writer: TcpStream,
    rx: Receiver<Esp32Message>,
    stop: Arc<AtomicBool>,
    reader_job: Option<JoinHandle<()>>,
}

impl Link {
    /// 建立连接并启动后台接收线程
    pub fn open(ip: &str, port: u16) -> std::io::Result<Self> {
        let addr = format!("{}:{}", ip, port);
        println!("[*] 连接 {} ...", addr);
        println!("请连接上esp32的wifi热点后再运行本程序【ESP32_AP】");
        let stream = TcpStream::connect(&addr)?;
        stream.set_nodelay(true)?;
        stream.set_read_timeout(Some(Duration::from_millis(200)))?;
        println!("[+] 已连接");

        let reader_stream = stream.try_clone()?;
        let (tx, rx): (Sender<Esp32Message>, Receiver<Esp32Message>) = mpsc::channel();
        let stop = Arc::new(AtomicBool::new(false));

        let reader_job = spawn_reader(reader_stream, tx, Arc::clone(&stop));

        Ok(Self {
            writer: stream,
            rx,
            stop,
            reader_job: Some(reader_job),
        })
    }

    pub fn open_default() -> std::io::Result<Self> {
        Self::open(ESP32_IP, ESP32_PORT)
    }

    // ---------- 发送 ----------

    /// 发送弧度列表
    pub fn send_angles(&mut self, rads: &[f32]) -> std::io::Result<()> {
        let n = rads.len().min(NUM_SERVOS);
        let payload = rads[..n]
            .iter()
            .map(|r| format!("{:.4}", r.clamp(0.0, PI)))
            .collect::<Vec<_>>()
            .join(",");
        self.send_line(&format!("RAD:{}", payload))
    }

    /// 发送角度列表
    pub fn send_degrees(&mut self, degs: &[f32]) -> std::io::Result<()> {
        let n = degs.len().min(NUM_SERVOS);
        let payload = degs[..n]
            .iter()
            .map(|d| format!("{:.2}", d.clamp(0.0, 180.0)))
            .collect::<Vec<_>>()
            .join(",");
        self.send_line(&format!("DEG:{}", payload))
    }

    /// 发送任意一行（自动补 '\n'）
    pub fn send_line(&mut self, line: &str) -> std::io::Result<()> {
        self.writer.write_all(line.as_bytes())?;
        self.writer.write_all(b"\n")?;
        self.writer.flush()
    }

    // ---------- 接收 ----------

    /// 阻塞等待一条消息（最多等 timeout）
    pub fn recv_timeout(&self, timeout: Duration) -> Option<Esp32Message> {
        match self.rx.recv_timeout(timeout) {
            Ok(m) => Some(m),
            Err(RecvTimeoutError::Timeout) => None,
            Err(RecvTimeoutError::Disconnected) => None,
        }
    }

    /// 非阻塞：立刻返回当前队列里有没有消息
    pub fn try_recv(&self) -> Option<Esp32Message> {
        self.rx.try_recv().ok()
    }
}

impl Drop for Link {
    fn drop(&mut self) {
        self.stop.store(true, Ordering::SeqCst);
        let _ = self.writer.shutdown(std::net::Shutdown::Both);
        if let Some(job) = self.reader_job.take() {
            let _ = job.join();
        }
    }
}

// ==================== 后台接收线程 ====================

fn spawn_reader(
    stream: TcpStream,
    tx: Sender<Esp32Message>,
    stop: Arc<AtomicBool>,
) -> JoinHandle<()> {
    thread::spawn(move || {
        let mut stream = stream;
        let mut pending = Vec::new();
        let mut buffer = [0u8; 1024];

        loop {
            if stop.load(Ordering::SeqCst) {
                break;
            }

            match stream.read(&mut buffer) {
                Ok(0) => break,
                Ok(n) => {
                    pending.extend_from_slice(&buffer[..n]);

                    loop {
                        if pending.len() >= TELEMETRY_FRAME_LEN
                            && pending[..2] == TELEMETRY_MAGIC
                        {
                            let frame = &pending[..TELEMETRY_FRAME_LEN];
                            if let Some(data) = parse_telemetry_frame(frame) {
                                if tx.send(Esp32Message::new_telemetry(data)).is_err() {
                                    return;
                                }
                                pending.drain(..TELEMETRY_FRAME_LEN);
                                continue;
                            }

                            // CRC 错误，丢掉首字节继续找下一帧
                            pending.drain(..1);
                            continue;
                        }

                        if let Some(newline_pos) = pending.iter().position(|&b| b == b'\n') {
                            let line_bytes = &pending[..newline_pos];
                            // 只处理纯 ASCII 文本; 非法 UTF-8 也不阻断连接
                            let line = String::from_utf8_lossy(line_bytes).trim().to_string();
                            if !line.is_empty() {
                                if tx.send(Esp32Message::new_text(line)).is_err() {
                                    return;
                                }
                            }
                            pending.drain(..=newline_pos);
                            continue;
                        }

                        // 还没凑齐完整帧/整行，等下一批数据
                        break;
                    }

                    if pending.len() > 65536 {
                        pending.clear();
                    }
                }
                Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock || e.kind() == std::io::ErrorKind::TimedOut => {
                    continue;
                }
                Err(_) => break,
            }
        }

        println!("[*] 接收线程退出");
    })
}

fn parse_telemetry_frame(frame: &[u8]) -> Option<[f32; 17]> {
    // Frame layout (from ESP32):
    // 0..1   magic 0xA5 0x5A
    // 2..3   gravity.x (int16 big-endian)  = projection * 10000
    // 4..5   gravity.y
    // 6..7   gravity.z
    // 8..9   gyro_x encoded (int16 BE) = rad/s * 1000
    // 10..11 gyro_y
    // 12..13 gyro_z
    // 14..15 CRC16 over bytes 0..13 (big-endian)
    if frame.len() != TELEMETRY_FRAME_LEN {
        return None;
    }
    if frame[..2] != TELEMETRY_MAGIC {
        return None;
    }

    let expected_crc = u16::from_be_bytes([frame[14], frame[15]]);
    let actual_crc = calculate_crc16_ccitt(&frame[..14]);
    if actual_crc != expected_crc {
        return None;
    }

    // Decode gravity projections (stored as projection * 10000)
    let gx_proj = i16::from_be_bytes([frame[2], frame[3]]) as f32 / 10000.0;
    let gy_proj = i16::from_be_bytes([frame[4], frame[5]]) as f32 / 10000.0;
    let gz_proj = i16::from_be_bytes([frame[6], frame[7]]) as f32 / 10000.0;

    // Decode base angular velocities (stored as rad/s * 1000)
    let avx = i16::from_be_bytes([frame[8], frame[9]]) as f32 / 1000.0;
    let avy = i16::from_be_bytes([frame[10], frame[11]]) as f32 / 1000.0;
    let avz = i16::from_be_bytes([frame[12], frame[13]]) as f32 / 1000.0;

    let mut values = [0.0f32; 17];
    // fill first 6 entries: gravity x,y,z then angular vel x,y,z
    values[0] = gx_proj;
    values[1] = gy_proj;
    values[2] = gz_proj;
    values[3] = avx;
    values[4] = avy;
    values[5] = avz;

    Some(values)
}

fn calculate_crc16_ccitt(data: &[u8]) -> u16 {
    let mut crc = 0xFFFFu16;
    for &byte in data {
        crc ^= (byte as u16) << 8;
        for _ in 0..8 {
            if crc & 0x8000 != 0 {
                crc = ((crc << 1) ^ 0x1021) & 0xFFFF;
            } else {
                crc = (crc << 1) & 0xFFFF;
            }
        }
    }
    crc
}