use std::{
    fmt::Display,
    str::FromStr,
};

use serde::{
    Deserializer,
    Serializer,
};

pub struct HexDisplayFromStr;

impl<T> serde_with::SerializeAs<T> for HexDisplayFromStr
where
    T: Display,
{
    fn serialize_as<S>(value: &T, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let s = format!("0x{value}");
        serializer.serialize_str(&s)
    }
}

impl<'de, T> serde_with::DeserializeAs<'de, T> for HexDisplayFromStr
where
    T: FromStr,
    T::Err: Display,
{
    fn deserialize_as<D>(deserializer: D) -> Result<T, D::Error>
    where
        D: Deserializer<'de>,
    {
        use serde::Deserialize;
        let s = String::deserialize(deserializer)?;
        let trimmed = s.strip_prefix("0x").unwrap_or(&s);
        T::from_str(trimmed).map_err(serde::de::Error::custom)
    }
}

pub fn serialize_hex<S, A>(value: A, serializer: S) -> Result<S::Ok, S::Error>
where
    S: Serializer,
    A: Display,
{
    serializer.serialize_str(&format!("0x{value}"))
}

pub struct HexVecFromStr;

impl serde_with::SerializeAs<Vec<u8>> for HexVecFromStr {
    fn serialize_as<S>(value: &Vec<u8>, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let hex_str = value.iter().map(|b| format!("{b:02x}")).collect::<String>();
        let hex_str = format!("0x{hex_str}");
        serializer.serialize_str(&hex_str)
    }
}

impl<'de> serde_with::DeserializeAs<'de, Vec<u8>> for HexVecFromStr {
    fn deserialize_as<D>(deserializer: D) -> Result<Vec<u8>, D::Error>
    where
        D: Deserializer<'de>,
    {
        use serde::{
            Deserialize,
            de,
        };
        use serde_json::Value;

        let value = Value::deserialize(deserializer)?;

        match value {
            Value::String(hex_str) => {
                let trimmed = hex_str.strip_prefix("0x").unwrap_or(&hex_str);

                (0..trimmed.len())
                    .step_by(2)
                    .map(|i| u8::from_str_radix(&trimmed[i..i + 2], 16))
                    .collect::<Result<Vec<u8>, _>>()
                    .map_err(de::Error::custom)
            }
            Value::Array(arr) => arr
                .into_iter()
                .map(|v| match v {
                    Value::Number(n) => {
                        if let Some(u) = n.as_u64() {
                            if u <= 255 {
                                Ok(u as u8)
                            } else {
                                Err(de::Error::custom(format!(
                                    "Number {u} is out of range for u8"
                                )))
                            }
                        } else {
                            Err(de::Error::custom("Invalid number in array"))
                        }
                    }
                    _ => Err(de::Error::custom("Array must contain only numbers")),
                })
                .collect::<Result<Vec<u8>, _>>(),
            _ => Err(de::Error::custom("Expected hex string or array of numbers")),
        }
    }
}
