#![no_std]

use core::str::from_utf8;

use thiserror::Error;

#[allow(async_fn_in_trait)]
pub trait PacketSocket {
    async fn write_all(&mut self, buf: &[u8]) -> bool;
    async fn read_exact(&mut self, buf: &mut [u8]) -> bool;
    async fn flush(&mut self) -> bool;
}

#[derive(Debug)]
pub enum Packet<'a> {
    Register { ip: [u8; 4], port: u16, id: &'a str },
    Ok,
    TimeSync(i64),
    State(State),
}

#[derive(Debug, Clone)]
pub struct Color(pub u8, pub u8, pub u8);

#[derive(Debug, Error)]
pub enum ColorError {
    #[error("invalid hex")]
    InvalidHex,
}

impl Color {
    pub fn from_hex(hex: &str) -> Result<Color, ColorError> {
        let r = u8::from_str_radix(&hex[1..3], 16).map_err(|_| ColorError::InvalidHex)?;
        let g = u8::from_str_radix(&hex[3..5], 16).map_err(|_| ColorError::InvalidHex)?;
        let b = u8::from_str_radix(&hex[5..7], 16).map_err(|_| ColorError::InvalidHex)?;

        Ok(Color(r, g, b))
    }
}

#[derive(Debug, Clone)]
pub enum State {
    Off,
    Color(Color),
    Sun(u8),
}

#[derive(Debug, Error)]
pub enum PacketError {
    #[error("unknown type")]
    UnknownType,
    #[error("buffer too small")]
    BufferTooSmall,
    #[error("no ok received")]
    NoOkRecieved,
    #[error("failed to write")]
    FailedToWrite,
    #[error("failed to read")]
    FailedToRead,
    #[error("data too large")]
    DataTooLarge,
    #[error("id too large")]
    IdTooLarge,
}

impl<'a> Packet<'a> {
    pub async fn send<'b>(
        socket: &mut impl PacketSocket,
        buf: &'b mut [u8],
        packet: Packet<'b>,
    ) -> Result<(), PacketError> {
        let length = packet.encode(&mut buf[4..])?;
        Self::write_u32(
            buf,
            0,
            u32::try_from(length + 4).map_err(|_| PacketError::DataTooLarge)?,
        );

        if !socket.write_all(&buf[..length + 4]).await {
            return Err(PacketError::FailedToWrite);
        }
        if !socket.flush().await {
            return Err(PacketError::FailedToWrite);
        }

        Ok(())
    }

    pub async fn save_send(
        socket: &mut impl PacketSocket,
        buf: &'a mut [u8],
        packet: Packet<'a>,
    ) -> Result<(), PacketError> {
        Self::send(socket, buf, packet).await?;

        let Some(Packet::Ok) = Packet::read(socket, buf).await? else {
            return Err(PacketError::NoOkRecieved);
        };

        Ok(())
    }

    pub async fn read(
        socket: &mut impl PacketSocket,
        buf: &'a mut [u8],
    ) -> Result<Option<Packet<'a>>, PacketError> {
        if !socket.read_exact(&mut buf[..4]).await {
            return Ok(None);
        }

        let target_length = Self::read_u32(buf, 0) as usize;
        if target_length > buf.len() - 4 {
            return Err(PacketError::DataTooLarge);
        }

        if !socket.read_exact(&mut buf[4..target_length]).await {
            return Err(PacketError::FailedToRead);
        }

        Ok(Some(Packet::decode(&buf[4..target_length])?))
    }

    fn encode(&self, buf: &mut [u8]) -> Result<usize, PacketError> {
        let len = match self {
            Packet::Ok => {
                Self::write_u8(buf, 0, 1);

                1
            }
            Packet::Register { ip, port, id } => {
                Self::write_u8(buf, 0, 2);

                Self::write_u8(buf, 1, ip[0]);
                Self::write_u8(buf, 2, ip[1]);
                Self::write_u8(buf, 3, ip[2]);
                Self::write_u8(buf, 4, ip[3]);

                Self::write_u16(buf, 5, *port);

                let id_bytes = id.as_bytes();
                Self::write_u8(
                    buf,
                    7,
                    u8::try_from(id_bytes.len()).map_err(|_| PacketError::IdTooLarge)?,
                );
                Self::write(buf, 8, id_bytes);

                8 + id_bytes.len()
            }
            Packet::TimeSync(timestamp) => {
                Self::write_u8(buf, 0, 3);
                Self::write_i64(buf, 1, *timestamp);

                9
            }
            Packet::State(State::Off) => {
                Self::write_u8(buf, 0, 4);
                1
            }
            Packet::State(State::Color(Color(r, g, b))) => {
                Self::write_u8(buf, 0, 5);

                Self::write_u8(buf, 1, *r);
                Self::write_u8(buf, 2, *g);
                Self::write_u8(buf, 3, *b);

                4
            }
            Packet::State(State::Sun(brightness)) => {
                Self::write_u8(buf, 0, 6);
                Self::write_u8(buf, 1, *brightness);

                2
            }
        };

        Ok(len)
    }

    fn write(buf: &mut [u8], offset: usize, value: &[u8]) {
        buf[offset..offset + value.len()].copy_from_slice(value);
    }

    fn write_i64(buf: &mut [u8], offset: usize, value: i64) {
        Self::write(buf, offset, &value.to_be_bytes());
    }

    fn write_u32(buf: &mut [u8], offset: usize, value: u32) {
        Self::write(buf, offset, &value.to_be_bytes());
    }

    fn write_u16(buf: &mut [u8], offset: usize, value: u16) {
        Self::write(buf, offset, &value.to_be_bytes());
    }

    fn write_u8(buf: &mut [u8], offset: usize, value: u8) {
        buf[offset] = value;
    }

    fn decode(buf: &'a [u8]) -> Result<Self, PacketError> {
        match Self::read_u8(buf, 0) {
            1 => Ok(Packet::Ok),
            2 => {
                let ip = [
                    Self::read_u8(buf, 1),
                    Self::read_u8(buf, 2),
                    Self::read_u8(buf, 3),
                    Self::read_u8(buf, 4),
                ];
                let port = Self::read_u16(buf, 5);

                let id_len = Self::read_u8(buf, 7) as usize;
                let id = from_utf8(&buf[8..8 + id_len]).map_err(|_| PacketError::FailedToRead)?;

                Ok(Packet::Register { ip, port, id })
            }
            3 => {
                let timestamp = Self::read_i64(buf, 1);

                Ok(Packet::TimeSync(timestamp))
            }
            4 => Ok(Packet::State(State::Off)),
            5 => {
                let red = Self::read_u8(buf, 1);
                let green = Self::read_u8(buf, 2);
                let blue = Self::read_u8(buf, 3);

                Ok(Packet::State(State::Color(Color(red, green, blue))))
            }
            6 => {
                let brightness = Self::read_u8(buf, 1);

                Ok(Packet::State(State::Sun(brightness)))
            }
            _ => Err(PacketError::UnknownType),
        }
    }

    fn read_i64(buf: &[u8], offset: usize) -> i64 {
        i64::from_be_bytes([
            buf[offset],
            buf[offset + 1],
            buf[offset + 2],
            buf[offset + 3],
            buf[offset + 4],
            buf[offset + 5],
            buf[offset + 6],
            buf[offset + 7],
        ])
    }

    fn read_u32(buf: &[u8], offset: usize) -> u32 {
        u32::from_be_bytes([
            buf[offset],
            buf[offset + 1],
            buf[offset + 2],
            buf[offset + 3],
        ])
    }

    fn read_u16(buf: &[u8], offset: usize) -> u16 {
        u16::from_be_bytes([buf[offset], buf[offset + 1]])
    }

    fn read_u8(buf: &[u8], offset: usize) -> u8 {
        buf[offset]
    }
}