vncrs 0.1.8

use crate::error::{Result, VncError};
use std::io::Read;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClientMessage {
    SetPixelFormat {
        format: [u8; 16],
    },
    SetEncodings {
        encodings: Vec<i32>,
    },
    FramebufferUpdateRequest {
        incremental: bool,
        x: u16,
        y: u16,
        width: u16,
        height: u16,
    },
    KeyEvent {
        down: bool,
        key: u32,
    },
    PointerEvent {
        buttons: u8,
        x: u16,
        y: u16,
    },
    ClientCutText {
        text: String,
    },
    EnableContinuousUpdates {
        enable: bool,
        x: u16,
        y: u16,
        width: u16,
        height: u16,
    },
    ClientFence {
        flags: u32,
        payload: Vec<u8>,
    },
}

pub const ENCODING_RAW: i32 = 0;
pub const ENCODING_COPYRECT: i32 = 1;
pub const ENCODING_HEXTILE: i32 = 5;
pub const ENCODING_ZLIB: i32 = 6;
pub const ENCODING_TIGHT: i32 = 7;
pub const ENCODING_ZRLE: i32 = 16;
pub const PSEUDO_TIGHT_QUALITY_BASE: i32 = -32; // -32..=-23 -> Quality 0..9
pub const PSEUDO_TIGHT_COMPRESSION_BASE: i32 = -260; // -260..=-251 -> Level 0..9
pub const PSEUDO_FENCE: i32 = -312;
pub const PSEUDO_CONTINUOUS_UPDATES: i32 = -313;

const MAX_CUT_TEXT_LEN: usize = 1024 * 1024; // 1 MB limit
const MAX_ENCODINGS_COUNT: usize = 256;

impl ClientMessage {
    pub fn read_from<R: Read>(reader: &mut R) -> Result<Self> {
        let mut msg_type = [0u8; 1];
        reader.read_exact(&mut msg_type)?;

        match msg_type[0] {
            0 => {
                let mut buf = [0u8; 19]; // 3 padding + 16 pixel format
                reader.read_exact(&mut buf)?;
                let mut format = [0u8; 16];
                format.copy_from_slice(&buf[3..19]);
                Ok(ClientMessage::SetPixelFormat { format })
            }
            2 => {
                let mut buf = [0u8; 3];
                reader.read_exact(&mut buf)?;
                let num = u16::from_be_bytes([buf[1], buf[2]]) as usize;
                if num > MAX_ENCODINGS_COUNT {
                    return Err(VncError::MessageTooLarge(num, MAX_ENCODINGS_COUNT));
                }
                let mut encodings = Vec::with_capacity(num);
                for _ in 0..num {
                    let mut enc = [0u8; 4];
                    reader.read_exact(&mut enc)?;
                    encodings.push(i32::from_be_bytes(enc));
                }
                Ok(ClientMessage::SetEncodings { encodings })
            }
            3 => {
                let mut buf = [0u8; 9];
                reader.read_exact(&mut buf)?;
                Ok(ClientMessage::FramebufferUpdateRequest {
                    incremental: buf[0] != 0,
                    x: u16::from_be_bytes([buf[1], buf[2]]),
                    y: u16::from_be_bytes([buf[3], buf[4]]),
                    width: u16::from_be_bytes([buf[5], buf[6]]),
                    height: u16::from_be_bytes([buf[7], buf[8]]),
                })
            }
            4 => {
                let mut buf = [0u8; 7];
                reader.read_exact(&mut buf)?;
                Ok(ClientMessage::KeyEvent {
                    down: buf[0] != 0,
                    key: u32::from_be_bytes([buf[3], buf[4], buf[5], buf[6]]),
                })
            }
            5 => {
                let mut buf = [0u8; 5];
                reader.read_exact(&mut buf)?;
                Ok(ClientMessage::PointerEvent {
                    buttons: buf[0],
                    x: u16::from_be_bytes([buf[1], buf[2]]),
                    y: u16::from_be_bytes([buf[3], buf[4]]),
                })
            }
            6 => {
                let mut buf = [0u8; 7];
                reader.read_exact(&mut buf)?;
                let len = u32::from_be_bytes([buf[3], buf[4], buf[5], buf[6]]) as usize;
                if len > MAX_CUT_TEXT_LEN {
                    return Err(VncError::MessageTooLarge(len, MAX_CUT_TEXT_LEN));
                }
                let mut text_buf = vec![0u8; len];
                reader.read_exact(&mut text_buf)?;
                let text = String::from_utf8_lossy(&text_buf).to_string();
                Ok(ClientMessage::ClientCutText { text })
            }
            150 => {
                let mut buf = [0u8; 9];
                reader.read_exact(&mut buf)?;
                Ok(ClientMessage::EnableContinuousUpdates {
                    enable: buf[0] != 0,
                    x: u16::from_be_bytes([buf[1], buf[2]]),
                    y: u16::from_be_bytes([buf[3], buf[4]]),
                    width: u16::from_be_bytes([buf[5], buf[6]]),
                    height: u16::from_be_bytes([buf[7], buf[8]]),
                })
            }
            248 => {
                let mut buf = [0u8; 8];
                reader.read_exact(&mut buf)?;
                let flags = u32::from_be_bytes([buf[3], buf[4], buf[5], buf[6]]);
                let len = buf[7] as usize;
                let mut payload = vec![0u8; len];
                if len > 0 {
                    reader.read_exact(&mut payload)?;
                }
                Ok(ClientMessage::ClientFence { flags, payload })
            }
            other => Err(VncError::Protocol(format!(
                "Unknown message type: {}",
                other
            ))),
        }
    }
}

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

    #[test]
    fn test_parse_key_event() {
        let data: Vec<u8> = vec![
            4, // message type
            1, // down-flag
            0, 0, // padding
            0x00, 0x00, 0xFF, 0x0D, // key (Return)
        ];
        let mut cursor = Cursor::new(&data[1..]);

        let mut buf = [0u8; 7];
        cursor.read_exact(&mut buf).unwrap();
        let down = buf[0] != 0;
        let key = u32::from_be_bytes([buf[3], buf[4], buf[5], buf[6]]);

        assert!(down);
        assert_eq!(key, 0xFF0D);
    }

    #[test]
    fn test_parse_pointer_event() {
        let data: Vec<u8> = vec![
            5, // message type
            1, // buttons
            0, 100, // x
            0, 200, // y
        ];
        let mut cursor = Cursor::new(data);
        let msg = ClientMessage::read_from(&mut cursor).unwrap();

        match msg {
            ClientMessage::PointerEvent { buttons, x, y } => {
                assert_eq!(buttons, 1);
                assert_eq!(x, 100);
                assert_eq!(y, 200);
            }
            _ => panic!("Expected PointerEvent"),
        }
    }

    #[test]
    fn test_parse_framebuffer_update_request() {
        let data: Vec<u8> = vec![
            3, // message type
            1, // incremental
            0, 0, // x
            0, 0, // y
            0x07, 0x80, // width = 1920
            0x04, 0x38, // height = 1080
        ];
        let mut cursor = Cursor::new(data);
        let msg = ClientMessage::read_from(&mut cursor).unwrap();

        match msg {
            ClientMessage::FramebufferUpdateRequest {
                incremental,
                x,
                y,
                width,
                height,
            } => {
                assert!(incremental);
                assert_eq!(x, 0);
                assert_eq!(y, 0);
                assert_eq!(width, 1920);
                assert_eq!(height, 1080);
            }
            _ => panic!("Expected FramebufferUpdateRequest"),
        }
    }

    #[test]
    fn test_parse_set_encodings() {
        let data: Vec<u8> = vec![
            2, // message type
            0, // padding
            0, 3, // number of encodings = 3
            0, 0, 0, 5, // Hextile
            0, 0, 0, 6, // Zlib
            0, 0, 0, 0, // Raw
        ];
        let mut cursor = Cursor::new(data);
        let msg = ClientMessage::read_from(&mut cursor).unwrap();

        match msg {
            ClientMessage::SetEncodings { encodings } => {
                assert_eq!(encodings, vec![5, 6, 0]);
            }
            _ => panic!("Expected SetEncodings"),
        }
    }

    #[test]
    fn test_parse_client_cut_text() {
        let text = b"Hello VNC!";
        let mut data: Vec<u8> = vec![
            6, // message type
            0, 0, 0, // padding
        ];
        data.extend_from_slice(&(text.len() as u32).to_be_bytes());
        data.extend_from_slice(text);

        let mut cursor = Cursor::new(data);
        let msg = ClientMessage::read_from(&mut cursor).unwrap();

        match msg {
            ClientMessage::ClientCutText { text: t } => {
                assert_eq!(t, "Hello VNC!");
            }
            _ => panic!("Expected ClientCutText"),
        }
    }

    #[test]
    fn test_parse_client_cut_text_too_large() {
        let mut data: Vec<u8> = vec![
            6, // message type
            0, 0, 0, // padding
        ];
        // 2 MB exceeds 1 MB limit
        data.extend_from_slice(&(2 * 1024 * 1024u32).to_be_bytes());

        let mut cursor = Cursor::new(data);
        let res = ClientMessage::read_from(&mut cursor);
        assert!(matches!(res, Err(VncError::MessageTooLarge(_, _))));
    }

    #[test]
    fn test_parse_enable_continuous_updates() {
        let mut data: Vec<u8> = vec![150, 1]; // enable = true
        data.extend_from_slice(&10u16.to_be_bytes()); // x
        data.extend_from_slice(&20u16.to_be_bytes()); // y
        data.extend_from_slice(&1920u16.to_be_bytes()); // width
        data.extend_from_slice(&1080u16.to_be_bytes()); // height

        let mut cursor = Cursor::new(data);
        let msg = ClientMessage::read_from(&mut cursor).unwrap();
        match msg {
            ClientMessage::EnableContinuousUpdates {
                enable,
                x,
                y,
                width,
                height,
            } => {
                assert!(enable);
                assert_eq!(x, 10);
                assert_eq!(y, 20);
                assert_eq!(width, 1920);
                assert_eq!(height, 1080);
            }
            _ => panic!("Expected EnableContinuousUpdates"),
        }
    }

    #[test]
    fn test_parse_client_fence() {
        let mut data: Vec<u8> = vec![248, 0, 0, 0]; // type + 3 pad
        data.extend_from_slice(&0x12345678u32.to_be_bytes()); // flags
        data.push(4); // length
        data.extend_from_slice(b"SYNC"); // payload

        let mut cursor = Cursor::new(data);
        let msg = ClientMessage::read_from(&mut cursor).unwrap();
        match msg {
            ClientMessage::ClientFence { flags, payload } => {
                assert_eq!(flags, 0x12345678);
                assert_eq!(payload, b"SYNC");
            }
            _ => panic!("Expected ClientFence"),
        }
    }
}