use crate::{
bluetooth::{self, Connection, RecordingError},
cube::{Direction, Face, Move},
};
use byteorder::{LittleEndian, ReadBytesExt};
use futures::{pin_mut, StreamExt};
use std::time::Duration;
use tokio::{
fs::File,
io::AsyncReadExt,
io::AsyncWriteExt,
spawn,
sync::mpsc::channel,
sync::mpsc::{Receiver, Sender},
time::sleep,
};
use bluer::{gatt::remote::Characteristic, AdapterEvent, Address, Device};
use async_trait::async_trait;
const SERVICE_UUID: &str = "00001000-0000-1000-8000-00805f9b34fb";
const CHAR_UUID_CMD_READ: &str = "00001002-0000-1000-8000-00805f9b34fb";
const CHAR_UUID_TURNS: &str = "00001003-0000-1000-8000-00805f9b34fb";
const CHAR_UUID_GYRO: &str = "00001004-0000-1000-8000-00805f9b34fb";
#[derive(Debug)]
pub struct FaceMessage {
pub increments: u8,
pub seconds: u16,
pub subseconds: u16,
pub direction: Direction,
pub face: Face,
}
impl FaceMessage {
const FACE_ORDER: [Face; 6] = [Face::D, Face::L, Face::B, Face::R, Face::F, Face::U];
fn decode(encoded: &[u8]) -> Option<Self> {
let increments = *encoded.get(0)?;
let direction = match encoded.get(6)? {
36 => Direction::Clockwise,
220 => Direction::CounterClockwise,
_ => return None,
};
let seconds = ReadBytesExt::read_u16::<LittleEndian>(&mut encoded.get(1..3)?).ok()?;
let subseconds = ReadBytesExt::read_u16::<LittleEndian>(&mut encoded.get(3..5)?).ok()?;
let face = *FaceMessage::FACE_ORDER.get(*encoded.get(5)? as usize)?;
Some(Self {
increments,
seconds,
subseconds,
direction,
face,
})
}
fn compress(raw: &[u8]) -> &[u8] {
&raw[..7]
}
fn signed_increments(&self) -> i32 {
match self.direction {
Direction::Clockwise => self.increments as i32,
Direction::CounterClockwise => -(self.increments as i32),
}
}
pub async fn load_dump(path: &str) -> Option<Vec<FaceMessage>> {
let mut messages = vec![];
let mut file = File::open(path).await.ok()?;
let mut buffer = [0; 7];
while file.read(&mut buffer).await.ok() == Some(7) {
messages.push(FaceMessage::decode(&buffer).unwrap());
}
Some(messages)
}
}
pub fn reconstruct_qtm(face_messages: &Vec<FaceMessage>) -> Vec<Move> {
struct SimplifiedFaceMessage {
increments: i32,
face: Face,
}
let mut condensed_msgs: Vec<SimplifiedFaceMessage> = vec![];
for msg in face_messages {
match &mut condensed_msgs[..] {
[.., a] if a.face == msg.face => {
if a.increments.signum() == msg.signed_increments().signum() {
a.increments += msg.signed_increments();
} else {
let excess = a.increments % 9;
if excess.abs() > 4 {
condensed_msgs.push(SimplifiedFaceMessage {
increments: (-9 * excess.signum()) + excess,
face: msg.face,
});
} else {
a.increments += msg.signed_increments();
}
}
}
[.., b, a] if a.face.opposite() == msg.face && b.face == msg.face => {
if b.increments.signum() == msg.signed_increments().signum() {
b.increments += msg.signed_increments();
} else {
let excess = b.increments % 9;
if excess.abs() > 4 {
condensed_msgs.push(SimplifiedFaceMessage {
increments: (-9 * excess.signum()) + excess,
face: msg.face,
});
} else {
b.increments += msg.signed_increments();
}
}
}
_ => {
condensed_msgs.push(SimplifiedFaceMessage {
face: msg.face,
increments: msg.signed_increments(),
});
}
}
}
condensed_msgs
.iter()
.map(|a| {
let mut turns = a.increments / 9;
let excess = a.increments % 9;
if excess.abs() > 4 {
turns += excess.signum();
}
Move {
face: a.face,
amount: turns,
}
})
.flat_map(|m| {
std::iter::repeat(Move {
face: m.face,
amount: m.amount.signum(),
})
.take(m.amount.abs() as usize)
})
.collect::<Vec<Move>>()
}
enum Signal {
GetBuffer,
StopRecording,
}
struct Characteristics {
cmd_read: Characteristic,
turns: Characteristic,
gyro: Characteristic,
}
pub struct MoyuWeilongAI {
device: Device,
characteristics: Characteristics,
receiver: Option<Receiver<Vec<Move>>>,
sender: Option<Sender<Signal>>,
}
impl MoyuWeilongAI {
async fn get_characteristics(device: &Device) -> bluer::Result<Characteristics> {
let service_uuid = uuid::Uuid::parse_str(SERVICE_UUID).unwrap();
let cmd_read_uuid = uuid::Uuid::parse_str(CHAR_UUID_CMD_READ).unwrap();
let turns_uuid = uuid::Uuid::parse_str(CHAR_UUID_TURNS).unwrap();
let gyro_uuid = uuid::Uuid::parse_str(CHAR_UUID_GYRO).unwrap();
let mut char_cmd_read = None;
let mut char_turns = None;
let mut char_gyro = None;
let uuids = device.uuids().await?.unwrap_or_default();
if uuids.contains(&service_uuid) {
sleep(Duration::from_secs(2)).await;
if !device.is_connected().await? {
let mut retries = 2;
loop {
match device.connect().await {
Ok(()) => break,
Err(_) if retries > 0 => {
retries -= 1;
}
Err(err) => return Err(err),
}
}
}
for service in device.services().await? {
let uuid = service.uuid().await?;
if uuid == service_uuid {
for char_ in service.characteristics().await? {
let uuid = char_.uuid().await?;
if uuid == cmd_read_uuid {
char_cmd_read = Some(char_);
} else if uuid == turns_uuid {
char_turns = Some(char_);
} else if uuid == gyro_uuid {
char_gyro = Some(char_);
}
}
}
}
} else {
return Err(bluer::Error {
kind: bluer::ErrorKind::NotFound,
message: "Weilong AI service not found.".to_string(),
});
}
if let (Some(cmd_read), Some(turns), Some(gyro)) = (char_cmd_read, char_turns, char_gyro) {
return Ok(Characteristics {
cmd_read,
turns,
gyro,
});
} else {
device.disconnect().await?;
}
Err(bluer::Error {
kind: bluer::ErrorKind::NotFound,
message: "Service does not contain Weilong AI characteristics.".to_string(),
})
}
pub async fn connect(addr: &Address) -> bluer::Result<Self> {
let session = bluer::Session::new().await?;
let adapter = session.default_adapter().await?;
adapter.set_powered(true).await?;
let device = adapter.device(*addr)?;
Ok(Self {
characteristics: Self::get_characteristics(&device).await?,
device,
receiver: None,
sender: None,
})
}
pub async fn auto_connect() -> bluer::Result<Self> {
let session = bluer::Session::new().await?;
let adapter = session.default_adapter().await?;
adapter.set_powered(true).await?;
let discover = adapter.discover_devices().await?;
pin_mut!(discover);
while let Some(evt) = discover.next().await {
match evt {
AdapterEvent::DeviceAdded(device_addr) => {
let device = adapter.device(device_addr)?;
match Self::get_characteristics(&device).await {
Ok(characteristics) => {
return Ok(Self {
device,
characteristics,
receiver: None,
sender: None,
})
}
_ => {}
}
}
_ => {}
}
}
Err(bluer::Error {
kind: bluer::ErrorKind::NotFound,
message: "Could not find any Moyu Weilong AIs.".to_string(),
})
}
}
#[async_trait]
impl Connection for MoyuWeilongAI {
async fn is_connected(&self) -> bluer::Result<bool> {
self.device.is_connected().await
}
async fn get_connected(&self) -> Address {
self.device.address()
}
async fn stop_recording(&mut self) {
if let Some(sender) = &self.sender {
let _ = sender.send(Signal::StopRecording).await;
}
self.sender = None;
self.receiver = None;
}
async fn get_buffer(&mut self) -> Option<Vec<Move>> {
if let Some(sender) = &self.sender {
let _ = sender.send(Signal::GetBuffer).await;
if let Some(receiver) = &mut self.receiver {
return receiver.recv().await;
}
}
None
}
async fn record_to_files(&mut self, stem: &str) -> bluetooth::Result<()> {
if self.sender.is_some() {
self.stop_recording().await;
}
let (tx, mut rx) = channel::<Signal>(32);
let _cmd_read_notify = self
.characteristics
.cmd_read
.notify()
.await
.map_err(|e| RecordingError::BlueR(e))?;
let turns_notify = self
.characteristics
.turns
.notify()
.await
.map_err(|e| RecordingError::BlueR(e))?;
let gyro_notify = self
.characteristics
.gyro
.notify()
.await
.map_err(|e| RecordingError::BlueR(e))?;
let mut turns_file = File::create(format!("{}.weilong.turns", stem))
.await
.map_err(|e| RecordingError::IO(e))?;
let mut gyro_file = File::create(format!("{}.weilong.gyro", stem))
.await
.map_err(|e| RecordingError::IO(e))?;
spawn(async move {
pin_mut!(turns_notify);
pin_mut!(gyro_notify);
pin_mut!(_cmd_read_notify);
'writing: loop {
tokio::select! {
Some(value) = turns_notify.next() => {
turns_file.write_all(FaceMessage::compress(&value)).await.unwrap();
},
Some(value) = gyro_notify.next() => {
gyro_file.write_all(&value).await.unwrap();
},
Some(signal) = rx.recv() => {
match signal {
Signal::StopRecording => break 'writing,
_ => { },
}
},
else => break,
}
}
});
self.sender = Some(tx);
Ok(())
}
async fn record_to_buffer(&mut self) -> bluetooth::Result<()> {
if self.sender.is_some() {
self.stop_recording().await;
}
let (tx, mut rx) = channel::<Signal>(32);
let (tx2, rx2) = channel::<Vec<Move>>(32);
let _cmd_read_notify = self
.characteristics
.cmd_read
.notify()
.await
.map_err(|e| RecordingError::BlueR(e))?;
let _gyro_notify = self
.characteristics
.gyro
.notify()
.await
.map_err(|e| RecordingError::BlueR(e))?;
let turns_notify = self
.characteristics
.turns
.notify()
.await
.map_err(|e| RecordingError::BlueR(e))?;
spawn(async move {
pin_mut!(turns_notify);
pin_mut!(_gyro_notify);
pin_mut!(_cmd_read_notify);
let mut buffer: Vec<FaceMessage> = vec![];
'writing_buffer: loop {
tokio::select! {
biased;
Some(value) = turns_notify.next() => {
let msg = FaceMessage::decode(&value).unwrap();
buffer.push(msg);
},
Some(signal) = rx.recv() => {
match signal {
Signal::StopRecording => break 'writing_buffer,
Signal::GetBuffer =>
tx2.send(reconstruct_qtm(&buffer)).await.unwrap(),
}
},
else => break,
}
}
});
self.sender = Some(tx);
self.receiver = Some(rx2);
return Ok(());
}
}