#![feature(async_closure)]

use std::{sync::Arc, time::Duration};

use anyhow::{Context, Result};
use common::{Command, Response, SensorData, BAUDRATE};
use framed_codec::FramedCodec;
use futures::{SinkExt, StreamExt};
use nalgebra::{SimdPartialOrd, SimdValue, Vector3};
use tokio::{io::{AsyncReadExt, AsyncWriteExt}, net::{TcpListener, TcpSocket}, sync::{broadcast::{self, error::RecvError, Receiver}, watch::{self, Sender}, RwLock}, task::JoinHandle, time::{error::Elapsed, timeout}};
use tokio_serial::SerialPortBuilderExt;
use tokio_util::codec::Framed;

mod framed_codec;

#[derive(Default)]
struct RobotState {
    bus_voltage: f32,
}

#[tokio::main]
async fn main() -> Result<()> {
    let mut serial = tokio_serial::new("/dev/ttyAMA0", BAUDRATE).open_native_async()?;

    let (sensor_sender, sensor_data) = broadcast::channel(5);

    serial.set_exclusive(false)?;

    let (mut write, mut read) = Framed::new(serial, FramedCodec::new()).split();

    let (send, commands) = watch::channel(Command::Stop);

    let state = Arc::new(RwLock::new(RobotState {
        bus_voltage: f32::MAX,
        ..Default::default()
    }));

    tokio::spawn(update_telem(state.clone(), sensor_data.resubscribe()));

    let control_telem = sensor_data.resubscribe();
    tokio::spawn(async move {
        loop {
            let _ = send.send(Command::Stop);
            if let Err(e) = control(send.clone(), control_telem.resubscribe()).await {
                println!("controller exited: {e}");
            }
        }
    });

    println!("starting");
    write.send(Command::Enable).await?;
    write.flush().await?;
    println!("enable");

    loop {
        let Some(Ok(data)) = read.next().await else {
            continue;
        };
        if data.enabled {
            write.send(Command::FeedWatchdog).await?;
            write.flush().await?;

        } else {
            println!("enabled southbridge");
            write.send(Command::Enable).await?;
        }

        write.send(Command::Enable).await?;
        write.flush().await?;

        if let Some(data) = data.sensor_data {
            let _ = sensor_sender.send(data);
        }

        let voltage = state.read().await.bus_voltage;

        // 100% effort at 9.5v, 0% effort at 7.5v
        let brownout = ((voltage-7.5)/2.0).clamp(0., 1.);
        
        write.send(commands.borrow().clone().brownout(brownout)).await?;
        write.flush().await?;
    }
}

async fn control(sender: Sender<Command>, mut telem: Receiver<SensorData>) -> Result<()> {
    let listener = TcpListener::bind("0.0.0.0:3322").await?;

    let (stream, addr) = listener.accept().await?;

    println!("connected to {addr:?}");

    let (mut read, mut write) = stream.into_split();

    let _: JoinHandle<Result<()>> = tokio::spawn(async move {
        loop {
            let data = match telem.recv().await {
                Ok(it) => it,
                Err(RecvError::Lagged(_)) => {continue;},
                Err(e) => {Err(e)?},
            };
            let data = postcard::to_stdvec(&data)?;
            write.write_u32(data.len() as u32).await?;
            write.write_all(&data).await?;
        }
    });

    loop {
        let len = loop {
            match timeout(Duration::from_millis(150), read.read_u32()).await {
                Ok(it) => { break it?; },
                Err(Elapsed) => {
                    sender.send(Command::Stop)?;
                },
            };
        };
        let mut buf = vec![0; len as usize];
        timeout(Duration::from_millis(100), read.read_exact(&mut buf)).await??;

        let cmd: Command = postcard::from_bytes(&buf)?;

        sender.send(cmd)?;
    }
}

async fn update_telem(state: Arc<RwLock<RobotState>>, mut telem: Receiver<SensorData>) -> Result<()> {
    loop {
        let telem = telem.recv().await?;
        println!("sensor {telem:?}");

        match telem {
            SensorData::BusVoltage(voltage) => {
                state.write().await.bus_voltage = voltage;
            }
            _ => {}
        }

    }
}

trait Brownout {
    /// limit exertion to factor of full output (0 to 1)
    fn brownout(self, factor: f32) -> Self;
}

impl Brownout for Command {
    fn brownout(self, factor: f32) -> Self {
        // nothing is preserved
        match self {
            Self::Twist(fwd, rot) => Self::Twist(
                fwd.clamp(-factor, factor), 
                rot.clamp(-factor, factor)
            ),
            Self::SetLed(rgb) => {
                Self::SetLed(rgb.simd_min(Vector3::repeat((factor * u8::MAX as f32) as u8)))
            }
            other => other,
        }
    }
}