use bytes::{Buf, BytesMut};
use tokio_util::codec::{Decoder, Encoder};

use crate::error::WireError;
use crate::parser::{parse_message, ParsedCommand};

/// Tokio codec for framing wire protocol messages on a TCP stream.
///
/// The wire protocol is naturally length-prefixed:
/// the first 4 bytes of each message contain the total message length.
pub struct WireCodec;

impl Decoder for WireCodec {
    type Item = ParsedCommand;
    type Error = WireError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        if src.len() < 4 {
            return Ok(None);
        }

        // Peek at message length
        let msg_len = i32::from_le_bytes([src[0], src[1], src[2], src[3]]) as usize;

        if src.len() < msg_len {
            // Reserve space for the rest of the message
            src.reserve(msg_len - src.len());
            return Ok(None);
        }

        match parse_message(src)? {
            Some((cmd, bytes_consumed)) => {
                src.advance(bytes_consumed);
                Ok(Some(cmd))
            }
            None => Ok(None),
        }
    }
}

/// Encoder for raw byte responses (already serialized by the command handlers).
impl Encoder<Vec<u8>> for WireCodec {
    type Error = WireError;

    fn encode(&mut self, item: Vec<u8>, dst: &mut BytesMut) -> Result<(), Self::Error> {
        dst.extend_from_slice(&item);
        Ok(())
    }
}