use bytes::BufMut;

use super::limits::{validate_body_size, validate_raw_size};
use super::message::{MessageName, MessageSchema, WireCodecError, WireFrame, WireLimits};

const MAGIC: &[u8; 4] = b"EFDN";
const VERSION: u8 = 1;
const HEADER_LEN: usize = 24;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CompressionTag {
    Raw,
    Lz4,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FrameBody {
    bytes: Vec<u8>,
    compression: CompressionTag,
}

impl FrameBody {
    pub fn new(bytes: Vec<u8>, compression: CompressionTag) -> Self {
        Self { bytes, compression }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct FrameHeader {
    schema: MessageSchema,
    raw_len: usize,
    body_len: usize,
    compression: CompressionTag,
}

impl FrameHeader {
    pub fn write(
        schema: MessageSchema,
        raw_len: usize,
        body: FrameBody,
        limits: WireLimits,
    ) -> Result<WireFrame, WireCodecError> {
        validate_raw_size(raw_len, limits.max_raw)?;
        validate_body_size(body.bytes.len(), limits.max_body)?;
        let mut bytes = bytes::BytesMut::with_capacity(HEADER_LEN + body.bytes.len());
        write_header(&mut bytes, schema, raw_len, &body)?;
        bytes.extend_from_slice(&body.bytes);
        Ok(WireFrame::from_bytes(bytes.freeze()))
    }

    pub fn read(bytes: &[u8], limits: WireLimits) -> Result<(Self, &[u8]), WireCodecError> {
        if bytes.len() < HEADER_LEN {
            return Err(short_header());
        }
        validate_magic(bytes)?;
        let header = read_header(bytes)?;
        validate_body_size(header.body_len, limits.max_body)?;
        validate_raw_size(header.raw_len, limits.max_raw)?;
        let body = body_slice(bytes, header.body_len)?;
        Ok((header, body))
    }

    pub fn expect_schema(
        self,
        expected: MessageSchema,
        name: MessageName,
    ) -> Result<(), WireCodecError> {
        if self.schema == expected {
            return Ok(());
        }
        Err(WireCodecError::new(
            format!("wire schema does not match {}", name.as_str()),
            "decode the frame as the message type used to encode it",
        ))
    }

    pub const fn raw_len(self) -> usize {
        self.raw_len
    }

    pub const fn schema(self) -> MessageSchema {
        self.schema
    }

    pub const fn compression(self) -> CompressionTag {
        self.compression
    }
}

fn write_header(
    bytes: &mut bytes::BytesMut,
    schema: MessageSchema,
    raw_len: usize,
    body: &FrameBody,
) -> Result<(), WireCodecError> {
    let raw_len = checked_u32(raw_len, "raw")?;
    let body_len = checked_u32(body.bytes.len(), "body")?;
    bytes.extend_from_slice(MAGIC);
    bytes.put_u8(VERSION);
    bytes.put_u8(tag_byte(body.compression));
    bytes.put_u16(0);
    bytes.put_u64(schema.get());
    bytes.put_u32(raw_len);
    bytes.put_u32(body_len);
    Ok(())
}

fn read_header(bytes: &[u8]) -> Result<FrameHeader, WireCodecError> {
    if bytes[4] != VERSION {
        return Err(unsupported_version(bytes[4]));
    }
    Ok(FrameHeader {
        schema: MessageSchema::from_u64(read_u64(bytes, 8)),
        raw_len: read_u32(bytes, 16) as usize,
        body_len: read_u32(bytes, 20) as usize,
        compression: read_compression(bytes[5])?,
    })
}

fn body_slice(bytes: &[u8], len: usize) -> Result<&[u8], WireCodecError> {
    let expected = HEADER_LEN + len;
    if bytes.len() == expected {
        return Ok(&bytes[HEADER_LEN..]);
    }
    Err(WireCodecError::new(
        format!(
            "wire frame length is {}, header declares {expected}",
            bytes.len()
        ),
        "send a complete frame without extra trailing bytes",
    ))
}

fn validate_magic(bytes: &[u8]) -> Result<(), WireCodecError> {
    if &bytes[..4] == MAGIC {
        return Ok(());
    }
    Err(WireCodecError::new(
        "wire frame has invalid magic bytes",
        "send bytes produced by efdn wire",
    ))
}

fn read_compression(value: u8) -> Result<CompressionTag, WireCodecError> {
    match value {
        0 => Ok(CompressionTag::Raw),
        1 => Ok(CompressionTag::Lz4),
        _ => Err(WireCodecError::new(
            format!("unsupported wire compression tag {value}"),
            "upgrade efdn or reject the packet",
        )),
    }
}

fn tag_byte(tag: CompressionTag) -> u8 {
    match tag {
        CompressionTag::Raw => 0,
        CompressionTag::Lz4 => 1,
    }
}

fn checked_u32(value: usize, label: &'static str) -> Result<u32, WireCodecError> {
    u32::try_from(value).map_err(|_| {
        WireCodecError::new(
            format!("wire {label} length exceeds u32"),
            "lower WireLimits so one frame cannot exceed u32::MAX bytes",
        )
    })
}

fn read_u64(bytes: &[u8], start: usize) -> u64 {
    let mut out = [0; 8];
    out.copy_from_slice(&bytes[start..start + 8]);
    u64::from_be_bytes(out)
}

fn read_u32(bytes: &[u8], start: usize) -> u32 {
    let mut out = [0; 4];
    out.copy_from_slice(&bytes[start..start + 4]);
    u32::from_be_bytes(out)
}

fn short_header() -> WireCodecError {
    WireCodecError::new(
        "wire frame is shorter than the fixed header",
        "send a complete efdn wire frame",
    )
}

fn unsupported_version(version: u8) -> WireCodecError {
    WireCodecError::new(
        format!("unsupported wire frame version {version}"),
        "upgrade efdn or reject the packet",
    )
}
