use fastwebsockets::{OpCode, WebSocketError};
use tokio::io::{AsyncRead, AsyncReadExt};

use super::Received;

const MAX_MESSAGE_SIZE: usize = 16 * 1024 * 1024;

pub(super) struct Reader<R> {
    pub(super) stream: R,
    message: Option<Received>,
    validated: usize,
    masked: bool,
}

impl<R: AsyncRead + Unpin> Reader<R> {
    pub(super) fn new(stream: R, masked: bool) -> Self {
        Self {
            stream,
            message: None,
            validated: 0,
            masked,
        }
    }

    pub(super) async fn read(&mut self) -> Result<Received, WebSocketError> {
        loop {
            let mut header = [0; 2];
            self.stream.read_exact(&mut header).await?;
            if header[0] & 0x70 != 0 {
                return Err(WebSocketError::ReservedBitsNotZero);
            }
            if (header[1] & 0x80 != 0) != self.masked {
                return Err(WebSocketError::InvalidValue);
            }
            let fin = header[0] & 0x80 != 0;
            let opcode = OpCode::try_from(header[0] & 0x0f)?;
            let control = matches!(opcode, OpCode::Close | OpCode::Ping | OpCode::Pong);
            let length_code = header[1] & 0x7f;
            if control {
                if !fin {
                    return Err(WebSocketError::ControlFrameFragmented);
                }
                if length_code > 125 {
                    return Err(WebSocketError::InvalidValue);
                }
            }
            let length = match length_code {
                126 => {
                    let length = u64::from(self.stream.read_u16().await?);
                    if length < 126 {
                        return Err(WebSocketError::InvalidValue);
                    }
                    length
                }
                127 => {
                    let length = self.stream.read_u64().await?;
                    if !(65536..=i64::MAX as u64).contains(&length) {
                        return Err(WebSocketError::InvalidValue);
                    }
                    length
                }
                length => u64::from(length),
            };
            let mut mask = [0; 4];
            if self.masked {
                self.stream.read_exact(&mut mask).await?;
            }
            if control {
                let mut payload = vec![0; length as usize];
                self.stream.read_exact(&mut payload).await?;
                if self.masked {
                    for (index, byte) in payload.iter_mut().enumerate() {
                        *byte ^= mask[index % 4];
                    }
                }
                return Ok(Received { opcode, payload });
            }
            match opcode {
                OpCode::Text | OpCode::Binary => {
                    if self.message.is_some() {
                        return Err(WebSocketError::InvalidFragment);
                    }
                    self.message = Some(Received { opcode, payload: Vec::new() });
                    self.validated = 0;
                }
                OpCode::Continuation if self.message.is_none() => {
                    return Err(WebSocketError::InvalidContinuationFrame);
                }
                _ => {}
            }
            let message = self.message.as_mut().unwrap();
            let start = message.payload.len();
            // Check the aggregate before allocating or waiting for the payload.
            if length > (MAX_MESSAGE_SIZE - start) as u64 {
                return Err(WebSocketError::FrameTooLarge);
            }
            let read = (&mut self.stream)
                .take(length)
                .read_to_end(&mut message.payload)
                .await?;
            if read as u64 != length {
                return Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof).into());
            }
            if self.masked {
                for (index, byte) in message.payload[start..].iter_mut().enumerate() {
                    *byte ^= mask[index % 4];
                }
            }
            if message.opcode == OpCode::Text {
                match std::str::from_utf8(&message.payload[self.validated..]) {
                    Ok(_) => self.validated = message.payload.len(),
                    Err(error) if !fin && error.error_len().is_none() => {
                        // Recheck only the incomplete suffix when the next fragment arrives.
                        self.validated += error.valid_up_to();
                    }
                    Err(_) => return Err(WebSocketError::InvalidUTF8),
                }
            }
            if fin {
                return Ok(self.message.take().unwrap());
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn truncated_large_frame_is_unexpected_eof() {
        let mut frame = vec![0x82, 0xff];
        frame.extend_from_slice(&(MAX_MESSAGE_SIZE as u64).to_be_bytes());
        frame.extend_from_slice(&[1, 2, 3, 4]);
        frame.extend_from_slice(&[5, 6, 7]);
        let mut reader = Reader::new(frame.as_slice(), true);
        assert!(matches!(
            reader.read().await,
            Err(WebSocketError::IoError(error)) if error.kind() == std::io::ErrorKind::UnexpectedEof
        ));
    }
}