use tokio::io::{AsyncReadExt, AsyncWriteExt};

use crate::errors::{ErrorCode, TetrixError};

const FRAME_HEADER_SIZE: usize = 4;
const DEFAULT_MAX_FRAME_SIZE: u32 = 16 << 20; // 16 MiB

/// Write a length-prefixed JSON frame to the writer.
pub(crate) async fn write_frame<W: AsyncWriteExt + Unpin>(
    writer: &mut W,
    payload: &impl serde::Serialize,
) -> Result<(), TetrixError> {
    let body = serde_json::to_vec(payload)?;
    let len = body.len() as u32;
    let mut header = [0u8; FRAME_HEADER_SIZE];
    header.copy_from_slice(&len.to_be_bytes());
    writer.write_all(&header).await?;
    writer.write_all(&body).await?;
    writer.flush().await?;
    Ok(())
}

/// Read a length-prefixed JSON frame from the reader.
pub(crate) async fn read_frame<R: AsyncReadExt + Unpin>(
    reader: &mut R,
    max_size: Option<u32>,
) -> Result<serde_json::Value, TetrixError> {
    let max = max_size.unwrap_or(DEFAULT_MAX_FRAME_SIZE);
    let mut header = [0u8; FRAME_HEADER_SIZE];
    reader.read_exact(&mut header).await?;
    let len = u32::from_be_bytes(header);
    if len > max {
        return Err(TetrixError::new(
            ErrorCode::InvalidArgument,
            format!("frame size {len} exceeds maximum {max}"),
        ));
    }
    let mut body = vec![0u8; len as usize];
    reader.read_exact(&mut body).await?;
    let value: serde_json::Value = serde_json::from_slice(&body)?;
    Ok(value)
}

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

    #[tokio::test]
    async fn test_write_read_frame() {
        let payload = serde_json::json!({"hello": "world"});
        let mut buf = Vec::new();
        write_frame(&mut buf, &payload).await.unwrap();

        let mut reader = BufReader::new(Cursor::new(buf));
        let result = read_frame(&mut reader, None).await.unwrap();
        assert_eq!(result, payload);
    }

    #[tokio::test]
    async fn test_frame_too_large() {
        let payload = serde_json::json!({"data": "x".repeat(100)});
        let mut buf = Vec::new();
        write_frame(&mut buf, &payload).await.unwrap();

        let mut reader = BufReader::new(Cursor::new(buf));
        let result = read_frame(&mut reader, Some(10)).await;
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert_eq!(err.code, ErrorCode::InvalidArgument);
    }

    #[tokio::test]
    async fn test_empty_object() {
        let payload = serde_json::json!({});
        let mut buf = Vec::new();
        write_frame(&mut buf, &payload).await.unwrap();

        let mut reader = BufReader::new(Cursor::new(buf));
        let result = read_frame(&mut reader, None).await.unwrap();
        assert_eq!(result, serde_json::json!({}));
    }

    #[tokio::test]
    async fn test_multiple_frames() {
        let mut buf = Vec::new();
        let p1 = serde_json::json!({"op": "AUTH"});
        let p2 = serde_json::json!({"op": "INDEX"});
        write_frame(&mut buf, &p1).await.unwrap();
        write_frame(&mut buf, &p2).await.unwrap();

        let mut reader = BufReader::new(Cursor::new(buf));
        let r1 = read_frame(&mut reader, None).await.unwrap();
        let r2 = read_frame(&mut reader, None).await.unwrap();
        assert_eq!(r1["op"], "AUTH");
        assert_eq!(r2["op"], "INDEX");
    }

    #[tokio::test]
    async fn test_truncated_header() {
        let buf = vec![0u8; 2]; // only 2 bytes, need 4
        let mut reader = BufReader::new(Cursor::new(buf));
        let result = read_frame(&mut reader, None).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_truncated_body() {
        let mut buf = Vec::new();
        buf.extend_from_slice(&100u32.to_be_bytes()); // claims 100 bytes
        buf.extend_from_slice(b"short"); // only 5 bytes
        let mut reader = BufReader::new(Cursor::new(buf));
        let result = read_frame(&mut reader, None).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_large_payload() {
        let big_text = "a".repeat(50_000);
        let payload = serde_json::json!({"text": big_text});
        let mut buf = Vec::new();
        write_frame(&mut buf, &payload).await.unwrap();

        let mut reader = BufReader::new(Cursor::new(buf));
        let result = read_frame(&mut reader, None).await.unwrap();
        assert_eq!(result["text"].as_str().unwrap().len(), 50_000);
    }

    #[tokio::test]
    async fn test_binary_frame_encoding() {
        let payload = serde_json::json!({"key": "value"});
        let mut buf = Vec::new();
        write_frame(&mut buf, &payload).await.unwrap();

        let body = serde_json::to_vec(&payload).unwrap();
        let expected_len = body.len() as u32;
        let header_bytes = &buf[..4];
        let actual_len = u32::from_be_bytes(header_bytes.try_into().unwrap());
        assert_eq!(actual_len, expected_len);
        assert_eq!(&buf[4..], &body);
    }
}
