use std::{error::Error, fmt::Debug};

use crate::{binary::checksum, hash};

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

#[derive(Debug)]
pub enum DeserializationError {
    InvalidSignature,
    InvalidSignatureLength,
    InvalidChecksum,
    InvalidChecksumLength,
    InvalidDataLength,
    #[allow(dead_code)]
    InvalidData(&'static str),
    InvalidDataLengthMarker,
}

impl Error for DeserializationError {}

impl std::fmt::Display for DeserializationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self)
    }
}

pub trait Serialize {
    const SIGNATURE: &'static str;
    const CHECKSUM_LENGTH: usize = 1;
    const LENGTH: usize = 0;

    fn signature() -> [u8; 8] {
        hash::Hash::from(Self::SIGNATURE).as_bytes_capped::<8>()
    }

    fn serialize_unchecked(&self) -> Vec<u8>;
    fn deserialize_unchecked<B>(data: B) -> Result<Self>
    where
        B: AsRef<[u8]>,
        Self: Sized;

    fn serialize_checked(&self) -> Vec<u8> {
        let mut serialized = Vec::new();
        let inner = self.serialize_unchecked();

        let checksum = checksum(&inner);

        serialized.extend_from_slice(&Self::signature());
        serialized.extend_from_slice(&[checksum]);
        serialized.extend_from_slice(&inner);

        serialized
    }

    fn deserialize_checked<B>(data: B) -> Result<Self>
    where
        B: AsRef<[u8]>,
        Self: Sized,
    {
        let data = data.as_ref();
        let signature = Self::signature();

        if data.len() < signature.len() {
            return Err(DeserializationError::InvalidSignatureLength);
        }

        if data.len() < signature.len() + Self::CHECKSUM_LENGTH {
            return Err(DeserializationError::InvalidChecksumLength);
        }

        let encoded_sig = &data[..signature.len()];
        if encoded_sig != signature {
            return Err(DeserializationError::InvalidSignature);
        }

        let expected_checksum = data[signature.len()];

        let inner_data = &data[signature.len() + Self::CHECKSUM_LENGTH..];
        if Self::LENGTH > 0 && inner_data.len() != Self::LENGTH as usize {
            return Err(DeserializationError::InvalidDataLength);
        }

        if checksum(&inner_data) != expected_checksum {
            return Err(DeserializationError::InvalidChecksum);
        }

        Self::deserialize_unchecked(inner_data)
    }
}

#[cfg(test)]
fn basic_test<T>(value: T)
where
    T: Serialize + PartialEq + Debug,
{
    let serialized = value.serialize_unchecked();
    let deserialized =
        T::deserialize_unchecked(&serialized).expect("deserialize_unchecked failed.");

    let ser_checked = value.serialize_checked();
    let deserialized_checked =
        T::deserialize_checked(&ser_checked).expect("deserialize_checked failed.");

    assert_eq!(value, deserialized);
    assert_eq!(value, deserialized_checked);
}

pub type SetOfSerializedObjects = Vec<Vec<u8>>;
type SetOfSerializedObjectsLengthMarker = u16; // NOTE: 65534 bytes is the max length of a serialized object
const SET_OF_SERIALIZED_OBJECTS_LENGTH_MARKER_BYTES: usize =
    (SetOfSerializedObjectsLengthMarker::BITS / 8) as usize;

impl Serialize for SetOfSerializedObjects {
    const SIGNATURE: &'static str = "SetOfSerializedObjects";

    fn serialize_unchecked(&self) -> Vec<u8> {
        let mut lengths: Vec<SetOfSerializedObjectsLengthMarker> = Vec::new();
        for item in self {
            // TODO (low): Account for wrapping
            lengths.push(item.len() as SetOfSerializedObjectsLengthMarker);
        }

        let mut serialized = Vec::new();
        for length in &lengths {
            serialized.extend_from_slice(&length.to_be_bytes());
        }

        serialized.extend_from_slice(&SetOfSerializedObjectsLengthMarker::MAX.to_be_bytes()[..]);

        for item in self {
            serialized.extend_from_slice(item);
        }

        serialized
    }

    fn deserialize_unchecked<B>(data: B) -> Result<Self>
    where
        B: AsRef<[u8]>,
    {
        let data = data.as_ref();

        let mut lengths = Vec::new();

        for len_bytes in data.chunks_exact(SET_OF_SERIALIZED_OBJECTS_LENGTH_MARKER_BYTES) {
            let length = SetOfSerializedObjectsLengthMarker::from_be_bytes(
                len_bytes
                    .try_into()
                    .map_err(|_| DeserializationError::InvalidDataLengthMarker)?,
            );
            if length == SetOfSerializedObjectsLengthMarker::MAX {
                break;
            }
            lengths.push(length);
        }

        let mut after_lengths_data = data
            .iter()
            .skip((lengths.len() + 1) * SET_OF_SERIALIZED_OBJECTS_LENGTH_MARKER_BYTES)
            .cloned();

        let mut items = Vec::new();

        for length in lengths {
            let mut item_data = Vec::new();
            for _ in 0..length {
                let Some(item_byte) = after_lengths_data.next() else {
                    return Err(DeserializationError::InvalidDataLength);
                };
                item_data.push(item_byte);
            }

            items.push(item_data);
        }

        Ok(items)
    }
}

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

    fn sample_data() -> SetOfSerializedObjects {
        vec![
            vec![1, 2, 3],
            vec![4, 5, 6, 7, 8],
            vec![9, 10, 11, 12, 13, 14],
        ]
    }

    #[test]
    fn test_serialize_deserialize() {
        let data = sample_data();

        let serialized = data.serialize_unchecked();
        let deserialized = SetOfSerializedObjects::deserialize_unchecked(&serialized).unwrap();

        assert_eq!(data, deserialized);
    }

    #[test]
    fn test_serialize_checked() {
        let data = sample_data();

        let serialized = data.serialize_checked();
        let deserialized = SetOfSerializedObjects::deserialize_checked(&serialized).unwrap();

        assert_eq!(data, deserialized);
    }
}

pub type SetOfSerializableObjects<S> = Vec<S>;

impl<S> Serialize for SetOfSerializableObjects<S>
where
    S: Serialize,
{
    const SIGNATURE: &'static str = "SetOfSerializableObjects";

    fn serialize_unchecked(&self) -> Vec<u8> {
        let mut serialized = Vec::new();

        for item in self {
            serialized.push(item.serialize_checked());
        }

        serialized.serialize_checked()
    }

    fn deserialize_unchecked<B>(data: B) -> Result<Self>
    where
        B: AsRef<[u8]>,
        Self: Sized,
    {
        let data = data.as_ref();
        let serialized_items = SetOfSerializedObjects::deserialize_checked(data)?;

        let mut items = Vec::new();
        for serialized_item in serialized_items {
            let item = S::deserialize_checked(serialized_item)?;
            items.push(item);
        }

        Ok(items)
    }
}

impl Serialize for String {
    const SIGNATURE: &'static str = "String";

    fn serialize_unchecked(&self) -> Vec<u8> {
        self.as_bytes().to_vec()
    }

    fn deserialize_unchecked<B>(data: B) -> Result<Self>
    where
        B: AsRef<[u8]>,
        Self: Sized,
    {
        let data = data.as_ref();

        let string = String::from_utf8(data.to_vec())
            .map_err(|_| DeserializationError::InvalidData("Invalid UTF-8"))?;

        Ok(string)
    }
}

#[test]
fn test_string_serialization() {
    basic_test(String::from("Hello, world!"));
}

impl Serialize for char {
    const SIGNATURE: &'static str = "char";
    const LENGTH: usize = 4;

    fn serialize_unchecked(&self) -> Vec<u8> {
        (*self as u32).to_be_bytes().to_vec()
    }

    fn deserialize_unchecked<B>(data: B) -> Result<Self>
    where
        B: AsRef<[u8]>,
        Self: Sized,
    {
        let data = data.as_ref();

        let mut buf = [0u8; 4];
        buf.copy_from_slice(data);

        let i = u32::from_be_bytes(buf);
        let c = char::from_u32(i).ok_or(DeserializationError::InvalidData("Not a valid char."))?;

        Ok(c)
    }
}

#[test]
fn test_char_serialization() {
    basic_test('A');
}

impl Serialize for u64 {
    const SIGNATURE: &'static str = "u64";
    const LENGTH: usize = u64::BITS as usize / 8;

    fn serialize_unchecked(&self) -> Vec<u8> {
        let b = self.to_be_bytes();
        b.to_vec()
    }

    fn deserialize_unchecked<B>(data: B) -> Result<Self>
    where
        B: AsRef<[u8]>,
        Self: Sized,
    {
        let data = data.as_ref();

        let mut buf = [0u8; Self::LENGTH];
        buf.copy_from_slice(data);

        let i = u64::from_be_bytes(buf);

        Ok(i)
    }
}

#[test]
fn test_u64_serialization() {
    basic_test(1234567890u128);
}

impl Serialize for u128 {
    const SIGNATURE: &'static str = "u128";
    const LENGTH: usize = u128::BITS as usize / 8;

    fn serialize_unchecked(&self) -> Vec<u8> {
        let b = self.to_be_bytes();
        b.to_vec()
    }

    fn deserialize_unchecked<B>(data: B) -> Result<Self>
    where
        B: AsRef<[u8]>,
        Self: Sized,
    {
        let data = data.as_ref();

        let mut buf = [0u8; Self::LENGTH];
        buf.copy_from_slice(data);

        let i = u128::from_be_bytes(buf);

        Ok(i)
    }
}
