use std::io::{self, Read, Write};

use thiserror::Error;

#[derive(Debug, Error)]
pub enum ProtoError {
    #[error("i/o error: {0}")]
    Io(#[from] io::Error),
    #[error("unexpected end of stream")]
    UnexpectedEof,
    #[error("invalid modified-UTF-8: {0}")]
    InvalidUtf8(String),
}

pub type Result<T> = std::result::Result<T, ProtoError>;

// ---------------------------------------------------------------------------
// Read primitives — mirrors h.java
// ---------------------------------------------------------------------------

/// Read 1 byte as unsigned (h.new, line 106).
pub fn read_u8(r: &mut impl Read) -> Result<u8> {
    let mut buf = [0u8; 1];
    r.read_exact(&mut buf)?;
    Ok(buf[0])
}

/// Read 1 byte as signed (h.try, line 91).
pub fn read_i8(r: &mut impl Read) -> Result<i8> {
    Ok(read_u8(r)? as i8)
}

/// Read 2 bytes big-endian as unsigned (h.int, line 138).
pub fn read_u16_be(r: &mut impl Read) -> Result<u16> {
    let mut buf = [0u8; 2];
    r.read_exact(&mut buf)?;
    Ok(u16::from_be_bytes(buf))
}

/// Read 2 bytes big-endian as signed (h.char, line 121).
pub fn read_i16_be(r: &mut impl Read) -> Result<i16> {
    let mut buf = [0u8; 2];
    r.read_exact(&mut buf)?;
    Ok(i16::from_be_bytes(buf))
}

/// Read 4 bytes big-endian as signed (h.do, line 172).
pub fn read_i32_be(r: &mut impl Read) -> Result<i32> {
    let mut buf = [0u8; 4];
    r.read_exact(&mut buf)?;
    Ok(i32::from_be_bytes(buf))
}

/// Read exactly `n` bytes into a new vec.
pub fn read_exact(r: &mut impl Read, n: usize) -> Result<Vec<u8>> {
    let mut buf = vec![0u8; n];
    r.read_exact(&mut buf)?;
    Ok(buf)
}

// ---------------------------------------------------------------------------
// Varint — aw.int(), line 484
//
// 1–3 byte variable-length integer, top bit = continuation.
//   Byte 0: value bits [6:0]
//   Byte 1 (if bit 7 of byte 0 set): value bits [13:7]
//   Byte 2 (if bit 7 of byte 1 set): value bits [21:14] (all 8 bits used)
//
// Maximum representable value: (0x7F) | (0x7F << 7) | (0xFF << 14)
//   = 127 + 16256 + 4177920 = 4194303 = 0x3FFFFF
// ---------------------------------------------------------------------------

/// Read a 1–3 byte varint (aw.int, line 484). Used for Tight compressed-stream
/// lengths, NOT for rectangle header coords.
pub fn read_varint(r: &mut impl Read) -> Result<u32> {
    let b0 = read_u8(r)? as u32;
    let mut val = b0 & 0x7F;
    if b0 & 0x80 != 0 {
        let b1 = read_u8(r)? as u32;
        val |= (b1 & 0x7F) << 7;
        if b1 & 0x80 != 0 {
            let b2 = read_u8(r)? as u32;
            val |= (b2 & 0xFF) << 14;
        }
    }
    Ok(val)
}

/// Write a value as a 1–3 byte varint. Panics if `val > 0x3FFFFF`.
pub fn write_varint(w: &mut impl Write, val: u32) -> Result<()> {
    assert!(val <= 0x3F_FFFF, "varint overflow: {val}");
    if val < 0x80 {
        w.write_all(&[val as u8])?;
    } else if val < 0x4000 {
        w.write_all(&[(val & 0x7F) as u8 | 0x80, ((val >> 7) & 0x7F) as u8])?;
    } else {
        w.write_all(&[
            (val & 0x7F) as u8 | 0x80,
            ((val >> 7) & 0x7F) as u8 | 0x80,
            ((val >> 14) & 0xFF) as u8,
        ])?;
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Modified-UTF-8 string — h.byte(), line 188
//
// Java's modified-UTF-8: u16 byte-length prefix, then encoded bytes.
// Encoding matches standard UTF-8 for BMP codepoints except:
//   - U+0000 is encoded as [0xC0, 0x80] (2 bytes, not 1)
//   - Supplementary characters use surrogate pairs
// We accept standard UTF-8 as well for robustness.
// ---------------------------------------------------------------------------

/// Read a modified-UTF-8 string (h.byte, line 188).
/// Format: u16-BE length prefix + `length` bytes of Java modified-UTF-8.
pub fn read_modified_utf8(r: &mut impl Read) -> Result<String> {
    let len = read_u16_be(r)? as usize;
    let data = read_exact(r, len)?;
    decode_modified_utf8(&data)
}

fn decode_modified_utf8(data: &[u8]) -> Result<String> {
    let mut out = String::with_capacity(data.len());
    let mut i = 0;
    while i < data.len() {
        let b = data[i];
        match b >> 4 {
            // Single byte: 0x01..=0x7F (standard ASCII, but NOT 0x00)
            0..=7 => {
                out.push(b as char);
                i += 1;
            }
            // Two-byte sequence: 110xxxxx 10xxxxxx
            12 | 13 => {
                if i + 1 >= data.len() {
                    return Err(ProtoError::InvalidUtf8(
                        "truncated 2-byte sequence".into(),
                    ));
                }
                let b2 = data[i + 1];
                if b2 & 0xC0 != 0x80 {
                    return Err(ProtoError::InvalidUtf8(
                        "invalid continuation byte".into(),
                    ));
                }
                let cp = ((b as u32 & 0x1F) << 6) | (b2 as u32 & 0x3F);
                out.push(char::from_u32(cp).unwrap_or('\u{FFFD}'));
                i += 2;
            }
            // Three-byte sequence: 1110xxxx 10xxxxxx 10xxxxxx
            14 => {
                if i + 2 >= data.len() {
                    return Err(ProtoError::InvalidUtf8(
                        "truncated 3-byte sequence".into(),
                    ));
                }
                let b2 = data[i + 1];
                let b3 = data[i + 2];
                if (b2 & 0xC0 != 0x80) || (b3 & 0xC0 != 0x80) {
                    return Err(ProtoError::InvalidUtf8(
                        "invalid continuation byte".into(),
                    ));
                }
                let cp =
                    ((b as u32 & 0x0F) << 12) | ((b2 as u32 & 0x3F) << 6) | (b3 as u32 & 0x3F);
                out.push(char::from_u32(cp).unwrap_or('\u{FFFD}'));
                i += 3;
            }
            _ => {
                return Err(ProtoError::InvalidUtf8(format!(
                    "invalid leading byte: 0x{b:02X}"
                )));
            }
        }
    }
    Ok(out)
}

// ---------------------------------------------------------------------------
// Rectangle header — aw.f(), line 468
//
// 4 × u16-BE coords + 1 × i32-BE encoding = 12 bytes fixed.
// Encoding is i32 because standard RFB uses negative pseudo-encoding IDs
// (e.g. -250).
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RectHeader {
    pub x: u16,
    pub y: u16,
    pub w: u16,
    pub h: u16,
    pub encoding: i32,
}

impl RectHeader {
    pub fn read_from(r: &mut impl Read) -> Result<Self> {
        Ok(Self {
            x: read_u16_be(r)?,
            y: read_u16_be(r)?,
            w: read_u16_be(r)?,
            h: read_u16_be(r)?,
            encoding: read_i32_be(r)?,
        })
    }
}

// ---------------------------------------------------------------------------
// Write helpers — for client-to-server messages
// ---------------------------------------------------------------------------

pub fn write_u8(w: &mut impl Write, v: u8) -> Result<()> {
    w.write_all(&[v])?;
    Ok(())
}

pub fn write_u16_be(w: &mut impl Write, v: u16) -> Result<()> {
    w.write_all(&v.to_be_bytes())?;
    Ok(())
}

pub fn write_i32_be(w: &mut impl Write, v: i32) -> Result<()> {
    w.write_all(&v.to_be_bytes())?;
    Ok(())
}

// ---------------------------------------------------------------------------
// RGB332 lookup table — ByteColorRFBRenderer constructor, lines 57-66
//
// DirectColorModel(8, 7, 56, 192):
//   red   mask = 0b_0000_0111 (bits 0-2), shift to 8-bit: v * 255 / 7
//   green mask = 0b_0011_1000 (bits 3-5), shift to 8-bit: v * 255 / 7
//   blue  mask = 0b_1100_0000 (bits 6-7), shift to 8-bit: v * 255 / 3
// ---------------------------------------------------------------------------

/// Expand an 8bpp RGB332 index to an RGBA pixel (alpha = 0xFF).
pub const fn rgb332_to_rgba(idx: u8) -> [u8; 4] {
    let r3 = idx & 0x07;
    let g3 = (idx >> 3) & 0x07;
    let b2 = (idx >> 6) & 0x03;
    [
        (r3 as u16 * 255 / 7) as u8,
        (g3 as u16 * 255 / 7) as u8,
        (b2 as u16 * 255 / 3) as u8,
        0xFF,
    ]
}

/// Precomputed RGB332 → RGBA lookup table (256 entries, 4 bytes each).
pub const RGB332_LUT: [[u8; 4]; 256] = {
    let mut lut = [[0u8; 4]; 256];
    let mut i = 0u16;
    while i < 256 {
        lut[i as usize] = rgb332_to_rgba(i as u8);
        i += 1;
    }
    lut
};

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // -- read primitives --

    #[test]
    fn test_read_u8() {
        let mut c = Cursor::new([0x42u8]);
        assert_eq!(read_u8(&mut c).unwrap(), 0x42);
    }

    #[test]
    fn test_read_i8() {
        let mut c = Cursor::new([0xFEu8]); // -2 as i8
        assert_eq!(read_i8(&mut c).unwrap(), -2);
    }

    #[test]
    fn test_read_u16_be() {
        let mut c = Cursor::new([0x01u8, 0x00]);
        assert_eq!(read_u16_be(&mut c).unwrap(), 256);
    }

    #[test]
    fn test_read_i16_be() {
        let mut c = Cursor::new([0xFF, 0xFE]); // -2 as i16
        assert_eq!(read_i16_be(&mut c).unwrap(), -2);
    }

    #[test]
    fn test_read_i32_be() {
        let mut c = Cursor::new([0xFF, 0xFF, 0xFF, 0x06]); // -250 as i32
        assert_eq!(read_i32_be(&mut c).unwrap(), -250);
    }

    // -- varint --

    #[test]
    fn test_varint_single_byte() {
        // Values 0..=127 are encoded as a single byte
        let mut c = Cursor::new([0x00]);
        assert_eq!(read_varint(&mut c).unwrap(), 0);

        let mut c = Cursor::new([0x7F]);
        assert_eq!(read_varint(&mut c).unwrap(), 127);
    }

    #[test]
    fn test_varint_two_bytes() {
        // 128 = 0x80: byte0 = (128 & 0x7F) | 0x80 = 0x80, byte1 = 128 >> 7 = 1
        let mut c = Cursor::new([0x80, 0x01]);
        assert_eq!(read_varint(&mut c).unwrap(), 128);

        // 16383 = 0x3FFF: byte0 = 0xFF, byte1 = 0x7F
        let mut c = Cursor::new([0xFF, 0x7F]);
        assert_eq!(read_varint(&mut c).unwrap(), 16383);
    }

    #[test]
    fn test_varint_three_bytes() {
        // 16384 = 0x4000: byte0 = 0x80, byte1 = 0x80, byte2 = 0x01
        let mut c = Cursor::new([0x80, 0x80, 0x01]);
        assert_eq!(read_varint(&mut c).unwrap(), 16384);

        // max: 0x3FFFFF = 4194303
        let mut c = Cursor::new([0xFF, 0xFF, 0xFF]);
        assert_eq!(read_varint(&mut c).unwrap(), 0x3F_FFFF);
    }

    #[test]
    fn test_varint_roundtrip() {
        for val in [0, 1, 127, 128, 255, 16383, 16384, 100_000, 0x3F_FFFF] {
            let mut buf = Vec::new();
            write_varint(&mut buf, val).unwrap();
            let mut c = Cursor::new(&buf);
            assert_eq!(read_varint(&mut c).unwrap(), val, "roundtrip failed for {val}");
        }
    }

    // -- modified UTF-8 --

    #[test]
    fn test_read_modified_utf8_ascii() {
        // u16 length = 5, then "hello"
        let data = [0x00, 0x05, b'h', b'e', b'l', b'l', b'o'];
        let mut c = Cursor::new(&data[..]);
        assert_eq!(read_modified_utf8(&mut c).unwrap(), "hello");
    }

    #[test]
    fn test_read_modified_utf8_null() {
        // Java modified-UTF-8 encodes U+0000 as [0xC0, 0x80]
        let data = [0x00, 0x02, 0xC0, 0x80];
        let mut c = Cursor::new(&data[..]);
        assert_eq!(read_modified_utf8(&mut c).unwrap(), "\0");
    }

    #[test]
    fn test_read_modified_utf8_multibyte() {
        // U+00E9 (é) = [0xC3, 0xA9] in UTF-8
        let data = [0x00, 0x02, 0xC3, 0xA9];
        let mut c = Cursor::new(&data[..]);
        assert_eq!(read_modified_utf8(&mut c).unwrap(), "é");
    }

    // -- rect header --

    #[test]
    fn test_rect_header() {
        // x=10, y=20, w=640, h=480, encoding=7 (Tight)
        let mut data = Vec::new();
        data.extend_from_slice(&10u16.to_be_bytes());
        data.extend_from_slice(&20u16.to_be_bytes());
        data.extend_from_slice(&640u16.to_be_bytes());
        data.extend_from_slice(&480u16.to_be_bytes());
        data.extend_from_slice(&7i32.to_be_bytes());

        let mut c = Cursor::new(&data[..]);
        let hdr = RectHeader::read_from(&mut c).unwrap();
        assert_eq!(hdr, RectHeader { x: 10, y: 20, w: 640, h: 480, encoding: 7 });
    }

    #[test]
    fn test_rect_header_negative_encoding() {
        // encoding = -250 (pseudo-encoding)
        let mut data = Vec::new();
        data.extend_from_slice(&0u16.to_be_bytes());
        data.extend_from_slice(&0u16.to_be_bytes());
        data.extend_from_slice(&0u16.to_be_bytes());
        data.extend_from_slice(&0u16.to_be_bytes());
        data.extend_from_slice(&(-250i32).to_be_bytes());

        let mut c = Cursor::new(&data[..]);
        let hdr = RectHeader::read_from(&mut c).unwrap();
        assert_eq!(hdr.encoding, -250);
    }

    // -- proptest --

    use proptest::prelude::*;

    proptest! {
        #[test]
        fn prop_varint_roundtrip(val in 0u32..=0x3F_FFFF) {
            let mut buf = Vec::new();
            write_varint(&mut buf, val).unwrap();
            let mut c = Cursor::new(&buf);
            let decoded = read_varint(&mut c).unwrap();
            prop_assert_eq!(decoded, val);
        }
    }

    // -- RGB332 LUT --

    #[test]
    fn test_rgb332_black() {
        assert_eq!(RGB332_LUT[0x00], [0, 0, 0, 255]);
    }

    #[test]
    fn test_rgb332_white() {
        // 0xFF = r=7, g=7, b=3 → [255, 255, 255, 255]
        assert_eq!(RGB332_LUT[0xFF], [255, 255, 255, 255]);
    }

    #[test]
    fn test_rgb332_pure_red() {
        // pure red = r=7, g=0, b=0 → 0x07
        assert_eq!(RGB332_LUT[0x07], [255, 0, 0, 255]);
    }

    #[test]
    fn test_rgb332_pure_green() {
        // pure green = r=0, g=7, b=0 → 0x38
        assert_eq!(RGB332_LUT[0x38], [0, 255, 0, 255]);
    }

    #[test]
    fn test_rgb332_pure_blue() {
        // pure blue = r=0, g=0, b=3 → 0xC0
        assert_eq!(RGB332_LUT[0xC0], [0, 0, 255, 255]);
    }
}