#[derive(Debug, Clone, Copy)]
pub struct TerminalInfo {
    pub columns: u32,
    pub rows: u32,
}

pub struct ClientMessageProxy;
pub struct ServerMessageProxy;

#[derive(Debug)]
pub enum ClientMessage<'a> {
    OpenSession(TerminalInfo),
    RunCommand(&'a str),
    Bye(&'a str),
    Input(&'a [u8]),
    CloseStdin,
}

#[derive(Debug, Clone, Copy)]
pub enum StreamIndex {
    Stdout,
    Stderr,
}

#[derive(Debug)]
pub enum ServerMessage<'a> {
    SessionOpen,
    CommandStatus(i32),
    Bye(&'a str),
    Output(StreamIndex, &'a [u8]),
}

#[derive(Debug, thiserror::Error)]
pub enum EncodeError {
    #[error("Message too large")]
    MessageTooLarge,
    #[error("Value too long")]
    ValueTooLong,
}

#[derive(Debug, thiserror::Error)]
pub enum DecodeError {
    #[error("Truncated message received")]
    Truncated,
    #[error("Malformed message received: {0}")]
    InvalidMessage(&'static str),
    #[error("Malformed string in the message")]
    InvalidString(core::str::Utf8Error),
}

pub struct Encoder<'a> {
    buffer: &'a mut [u8],
    pos: usize,
}

pub struct Decoder<'a> {
    buffer: &'a [u8],
    pos: usize,
}

pub trait Encode {
    fn encode(&self, buffer: &mut Encoder) -> Result<(), EncodeError>;
}

pub trait Decode<'de>: Sized + 'de {
    fn decode(buffer: &mut Decoder<'de>) -> Result<Self, DecodeError>;
}

pub trait MessageProxy {
    type Type<'de>: Encode + Decode<'de>;
}

impl<'a> Encoder<'a> {
    pub const fn new(buffer: &'a mut [u8]) -> Self {
        Self { buffer, pos: 0 }
    }

    pub fn reset(&mut self) {
        self.pos = 0;
    }

    pub fn write_variable_bytes(&mut self, bytes: &[u8]) -> Result<(), EncodeError> {
        let len: u16 = bytes
            .len()
            .try_into()
            .map_err(|_| EncodeError::ValueTooLong)?;
        self.write(&len.to_le_bytes())?;
        self.write(bytes)
    }

    pub fn write_str(&mut self, s: &str) -> Result<(), EncodeError> {
        self.write_variable_bytes(s.as_bytes())
    }

    pub fn write(&mut self, bytes: &[u8]) -> Result<(), EncodeError> {
        if self.pos + bytes.len() > self.buffer.len() {
            return Err(EncodeError::ValueTooLong);
        }
        self.buffer[self.pos..self.pos + bytes.len()].copy_from_slice(bytes);
        self.pos += bytes.len();
        Ok(())
    }

    pub fn get(&self) -> &[u8] {
        &self.buffer[..self.pos]
    }

    pub fn split_mut(&mut self) -> (&mut [u8], &mut [u8]) {
        self.buffer.split_at_mut(self.pos)
    }
}

impl<'a> Decoder<'a> {
    pub fn new(buffer: &'a [u8]) -> Self {
        Self { buffer, pos: 0 }
    }

    pub fn read_u8(&mut self) -> Result<u8, DecodeError> {
        if self.pos + 1 > self.buffer.len() {
            return Err(DecodeError::Truncated);
        }
        let byte = self.buffer[self.pos];
        self.pos += 1;
        Ok(byte)
    }

    pub fn read_le_u16(&mut self) -> Result<u16, DecodeError> {
        let bytes = self.read_bytes(size_of::<u16>())?;
        Ok(u16::from_le_bytes([bytes[0], bytes[1]]))
    }

    pub fn read_le_u32(&mut self) -> Result<u32, DecodeError> {
        let bytes = self.read_bytes(size_of::<u32>())?;
        Ok(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
    }

    pub fn read_bytes(&mut self, len: usize) -> Result<&'a [u8], DecodeError> {
        if self.pos + len > self.buffer.len() {
            return Err(DecodeError::Truncated);
        }
        let slice = &self.buffer[self.pos..self.pos + len];
        self.pos += len;
        Ok(slice)
    }

    pub fn read_variable_bytes(&mut self) -> Result<&'a [u8], DecodeError> {
        let len = self.read_le_u16()?;
        self.read_bytes(len as usize)
    }

    pub fn read_str(&mut self) -> Result<&'a str, DecodeError> {
        let slice = self.read_variable_bytes()?;
        core::str::from_utf8(slice).map_err(DecodeError::InvalidString)
    }

    pub fn split(&self) -> (&'a [u8], &'a [u8]) {
        self.buffer.split_at(self.pos)
    }
}

impl MessageProxy for ClientMessageProxy {
    type Type<'de> = ClientMessage<'de>;
}

impl MessageProxy for ServerMessageProxy {
    type Type<'de> = ServerMessage<'de>;
}

impl ClientMessage<'_> {
    const TAG_OPEN_SESSION: u8 = 0x80;
    const TAG_RUN_COMMAND: u8 = 0x81;
    const TAG_BYE: u8 = 0x82;
    const TAG_INPUT: u8 = 0x90;
    const TAG_CLOSE_STDIN: u8 = 0x91;
}

impl StreamIndex {
    const TAG_STDOUT: u8 = 0x01;
    const TAG_STDERR: u8 = 0x02;
}

impl Encode for StreamIndex {
    fn encode(&self, buffer: &mut Encoder) -> Result<(), EncodeError> {
        match self {
            Self::Stdout => buffer.write(&[Self::TAG_STDOUT]),
            Self::Stderr => buffer.write(&[Self::TAG_STDERR]),
        }
    }
}

impl<'de> Decode<'de> for StreamIndex {
    fn decode(buffer: &mut Decoder<'de>) -> Result<Self, DecodeError> {
        let tag = buffer.read_u8()?;
        match tag {
            Self::TAG_STDOUT => Ok(Self::Stdout),
            Self::TAG_STDERR => Ok(Self::Stderr),
            _ => Err(DecodeError::InvalidMessage("Invalid StreamIndex tag")),
        }
    }
}

impl Encode for TerminalInfo {
    fn encode(&self, buffer: &mut Encoder) -> Result<(), EncodeError> {
        buffer.write(&self.columns.to_le_bytes())?;
        buffer.write(&self.rows.to_le_bytes())?;
        Ok(())
    }
}

impl<'de> Decode<'de> for TerminalInfo {
    fn decode(buffer: &mut Decoder<'de>) -> Result<Self, DecodeError> {
        let columns = buffer.read_le_u32()?;
        let rows = buffer.read_le_u32()?;
        Ok(Self { columns, rows })
    }
}

impl Encode for ClientMessage<'_> {
    fn encode(&self, buffer: &mut Encoder) -> Result<(), EncodeError> {
        match self {
            Self::OpenSession(info) => {
                buffer.write(&[Self::TAG_OPEN_SESSION])?;
                info.encode(buffer)
            }
            Self::RunCommand(command) => {
                buffer.write(&[Self::TAG_RUN_COMMAND])?;
                buffer.write_str(command)
            }
            Self::Bye(reason) => {
                buffer.write(&[Self::TAG_BYE])?;
                buffer.write_str(reason)
            }
            Self::Input(data) => {
                buffer.write(&[Self::TAG_INPUT])?;
                buffer.write_variable_bytes(data)
            }
            Self::CloseStdin => buffer.write(&[Self::TAG_CLOSE_STDIN]),
        }
    }
}

impl<'de> Decode<'de> for ClientMessage<'de> {
    fn decode(buffer: &mut Decoder<'de>) -> Result<Self, DecodeError> {
        let tag = buffer.read_u8()?;
        match tag {
            Self::TAG_OPEN_SESSION => TerminalInfo::decode(buffer).map(Self::OpenSession),
            Self::TAG_RUN_COMMAND => buffer.read_str().map(Self::RunCommand),
            Self::TAG_BYE => buffer.read_str().map(Self::Bye),
            Self::TAG_INPUT => buffer.read_variable_bytes().map(Self::Input),
            Self::TAG_CLOSE_STDIN => Ok(Self::CloseStdin),
            _ => Err(DecodeError::InvalidMessage("Invalid ClientMessage tag")),
        }
    }
}

impl ServerMessage<'_> {
    const TAG_SESSION_OPEN: u8 = 0x10;
    const TAG_COMMAND_STATUS: u8 = 0x11;
    const TAG_BYE: u8 = 0x12;
    const TAG_OUTPUT: u8 = 0x20;
}

impl Encode for ServerMessage<'_> {
    fn encode(&self, buffer: &mut Encoder) -> Result<(), EncodeError> {
        match self {
            Self::SessionOpen => buffer.write(&[Self::TAG_SESSION_OPEN]),
            Self::CommandStatus(status) => buffer.write(&status.to_le_bytes()),
            Self::Bye(reason) => {
                buffer.write(&[Self::TAG_BYE])?;
                buffer.write_str(reason)
            }
            Self::Output(index, data) => {
                buffer.write(&[Self::TAG_OUTPUT])?;
                index.encode(buffer)?;
                buffer.write_variable_bytes(data)
            }
        }
    }
}

impl<'de> Decode<'de> for ServerMessage<'de> {
    fn decode(buffer: &mut Decoder<'de>) -> Result<Self, DecodeError> {
        let tag = buffer.read_u8()?;
        match tag {
            Self::TAG_SESSION_OPEN => Ok(Self::SessionOpen),
            Self::TAG_COMMAND_STATUS => {
                let mut status = [0; size_of::<u32>()];
                let bytes = buffer.read_bytes(size_of::<i32>())?;
                status.copy_from_slice(bytes);
                Ok(Self::CommandStatus(i32::from_le_bytes(status)))
            }
            Self::TAG_BYE => buffer.read_str().map(Self::Bye),
            Self::TAG_OUTPUT => {
                let index = StreamIndex::decode(buffer)?;
                let data = buffer.read_variable_bytes()?;
                Ok(Self::Output(index, data))
            }
            _ => Err(DecodeError::InvalidMessage("Invalid ServerMessage tag")),
        }
    }
}
