use anchor_lang::prelude::*;
use std::{collections::HashMap, str::FromStr};

#[derive(AnchorDeserialize, AnchorSerialize, Clone)]
pub enum Number {
    U64(u64),
    I64(i64),
    F64(f64),
}

impl ToString for Number {
    fn to_string(&self) -> String {
        match self {
            Self::U64(data) => data.to_string(),
            Self::I64(data) => data.to_string(),
            Self::F64(data) => data.to_string(),
        }
    }
}

impl Into<serde_json::Number> for Number {
    fn into(self) -> serde_json::Number {
        match self {
            Self::U64(value) => serde_json::Number::from(value),
            Self::I64(value) => serde_json::Number::from(value),
            Self::F64(value) => serde_json::Number::from_f64(value).unwrap(),
        }
    }
}

impl From<serde_json::Number> for Number {
    fn from(number: serde_json::Number) -> Self {
        if let Some(number) = number.as_u64() {
            Self::U64(number)
        } else if let Some(number) = number.as_i64() {
            Self::I64(number)
        } else if let Some(number) = number.as_f64() {
            Self::F64(number)
        } else {
            panic!("Serde Number cannot be converted to Number")
        }
    }
}

#[derive(Clone)]
pub enum Schema {
    Null,
    Bool,
    Number,
    String,
    Array(Box<Self>),
    Object(HashMap<String, Self>),
    Pubkey,
    Option(Box<Self>),
    HashMap(Box<Self>, Box<Self>),
    Enum(Vec<(String, Schema)>),
}

impl ToString for Schema {
    fn to_string(&self) -> String {
        let mut schema_str = String::new();
        match self {
            Self::Null => schema_str.push_str("null"),
            Self::Bool => schema_str.push_str("bool"),
            Self::Number => schema_str.push_str("number"),
            Self::String => schema_str.push_str("string"),
            Self::Array(schema) => {
                schema_str.push_str(format!("Array<{}>", schema.to_string()).as_str())
            }
            Self::Object(schema) => {
                schema_str.push_str("{ ");
                schema.iter().for_each(|(key, value)| {
                    schema_str.push_str(format!("{}: {}, ", key, value.to_string()).as_str());
                });
                schema_str.push_str("}");
            }
            Self::Pubkey => schema_str.push_str("pubkey"),
            Self::Option(schema) => {
                schema_str.push_str(format!("Option<{}>", schema.to_string()).as_str())
            }
            Self::HashMap(key, schema) => schema_str
                .push_str(format!("HashMap<{}, {}>", key.to_string(), schema.to_string()).as_str()),
            Self::Enum(variants) => {
                schema_str.push_str("Enum {");
                variants.iter().for_each(|variant| {
                    schema_str.push_str(&format!("{}: {}, ", variant.0, variant.1.to_string()));
                });
                schema_str.push_str("}");
            }
        };
        schema_str
    }
}

impl Schema {
    pub fn validate(&self, value: &mut SchemaValue) -> bool {
        value.validate(self)
    }

    pub fn size(&self) -> usize {
        match self {
            Self::Array(schema) => 1 + schema.size(),
            Self::Object(hashmap) => {
                let mut size = 1;
                for (key, value) in hashmap {
                    size += key.as_bytes().len();
                    size += value.size();
                }
                size
            }
            Self::Option(schema) => 2 + schema.size(),
            Self::HashMap(key_schema, value_schema) => 1 + key_schema.size() + value_schema.size(),
            Self::Enum(variants) => {
                1 + variants.iter().fold(0, |acc, variant| {
                    acc + variant.0.as_bytes().len() + variant.1.size()
                })
            }
            _ => 1,
        }
    }

    pub fn size_for_borsh(&self) -> usize {
        match self.try_to_vec() {
            Ok(vec) => vec.len(),
            Err(_) => 0,
        }
    }
}

impl AnchorSerialize for Schema {
    #[inline]
    fn serialize<W: std::io::prelude::Write>(&self, writer: &mut W) -> std::io::Result<()> {
        match self {
            Self::Null => 0u8.serialize(writer),
            Self::Bool => 1u8.serialize(writer),
            Self::Number => 2u8.serialize(writer),
            Self::String => 3u8.serialize(writer),
            Self::Array(value) => {
                4u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::Object(value) => {
                5u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::Pubkey => 6u8.serialize(writer),
            Self::Option(value) => {
                7u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::HashMap(k, v) => {
                8u8.serialize(writer)?;
                k.serialize(writer)?;
                v.serialize(writer)
            }
            Self::Enum(variants) => {
                9u8.serialize(writer)?;
                variants.serialize(writer)
            }
        }
    }
}

impl AnchorDeserialize for Schema {
    fn deserialize_reader<R: std::io::prelude::Read>(reader: &mut R) -> std::io::Result<Self> {
        let mut buf: [u8; 1] = [0; 1];
        let limit = reader.read(&mut buf)?;
        if limit != 1 {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "Unexpected length of input",
            ));
        }
        match buf[0] {
            0 => Ok(Self::Null),
            1 => Ok(Self::Bool),
            2 => Ok(Self::Number),
            3 => Ok(Self::String),
            4 => Schema::deserialize_reader(reader).and_then(|v| Ok(Self::Array(Box::new(v)))),
            5 => HashMap::deserialize_reader(reader).and_then(|v| Ok(Self::Object(v))),
            6 => Ok(Self::Pubkey),
            7 => Schema::deserialize_reader(reader).and_then(|v| Ok(Self::Option(Box::new(v)))),
            8 => {
                let k = Schema::deserialize_reader(reader)?;
                let v = Schema::deserialize_reader(reader)?;
                Ok(Self::HashMap(Box::new(k), Box::new(v)))
            }
            9 => {
                let variants = Vec::<(String, Schema)>::deserialize_reader(reader)?;
                Ok(Self::Enum(variants))
            }
            _ => {
                let msg = format!(
                    "Invalid Option representation: {}. The first byte must be 0 till 3",
                    buf[0]
                );
                Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, msg))
            }
        }
    }
}

#[derive(Clone)]
pub enum SchemaValue {
    Null,
    Bool(bool),
    Number(Number),
    String(String),
    Array(Vec<Self>),
    Object(HashMap<String, Self>),
    Pubkey(Pubkey),
    HashMap(HashMap<String, Self>),
    Enum(String, Box<Self>),
}

impl SchemaValue {
    pub fn validate(&mut self, schema: &Schema) -> bool {
        match schema {
            Schema::Null => matches!(self, Self::Null),
            Schema::Bool => matches!(self, Self::Bool(_)),
            Schema::Number => matches!(self, Self::Number(_)),
            Schema::String => matches!(self, Self::String(_)),
            Schema::Array(schema) => match self {
                Self::Array(arr) => arr.iter_mut().all(|value| value.validate(schema)),
                _ => false,
            },
            Schema::Object(schema) => match self {
                Self::Object(obj) => schema.iter().all(|(key, schema)| {
                    if let Some(value) = obj.get_mut(key) {
                        return value.validate(schema);
                    } else if matches!(schema, Schema::Option(_)) {
                        return true;
                    } else if matches!(schema, Schema::Null) {
                        obj.insert(key.to_owned(), Self::Null);
                        return true;
                    }
                    false
                }),
                _ => false,
            },
            Schema::Pubkey => matches!(self, Self::Pubkey(_)),
            Schema::Option(schema) => self.validate(&Schema::Null) || self.validate(schema),
            Schema::HashMap(_k, v) => match self {
                Self::HashMap(map) => map.iter_mut().all(|(_, value)| value.validate(v)),
                _ => false,
            },
            Schema::Enum(variants) => match self {
                Self::Enum(variant_kind, variant_value) => variants
                    .iter()
                    .find(|variant| {
                        return variant.0.eq(variant_kind) && variant_value.validate(&variant.1);
                    })
                    .is_some(),
                _ => false,
            },
        }
    }

    pub fn size(&self) -> usize {
        match self {
            Self::Null => 1,
            Self::Bool(_) => 2,
            Self::Number(_) => 1 + 8,
            Self::String(string) => 1 + string.as_bytes().len(),
            Self::Array(vec) => {
                let mut size = 1;
                for schema in vec {
                    size += schema.size()
                }
                size
            }
            Self::Object(hashmap) => {
                let mut size = 1;
                for (key, value) in hashmap {
                    size += key.as_bytes().len();
                    size += value.size();
                }
                size
            }
            Self::Pubkey(_) => 1 + 32,
            Self::HashMap(hashmap) => {
                let mut size = 1;
                for (key, value) in hashmap {
                    size += key.as_bytes().len();
                    size += value.size();
                }
                size
            }
            Self::Enum(kind, value) => 1 + kind.as_bytes().len() + value.size(),
        }
    }

    pub fn size_for_borsh(&self) -> usize {
        match self.try_to_vec() {
            Ok(vec) => vec.len(),
            Err(_) => 0,
        }
    }
}

impl ToString for SchemaValue {
    fn to_string(&self) -> String {
        let mut schema_value_str = String::new();
        match self {
            Self::Null => schema_value_str.push_str("null"),
            Self::Bool(data) => schema_value_str.push_str(data.to_string().as_str()),
            Self::Number(data) => schema_value_str.push_str(data.to_string().as_str()),
            Self::String(data) => schema_value_str.push_str(format!("\"{}\"", data).as_str()),
            Self::Array(data) => {
                schema_value_str.push_str("[ ");
                data.iter().for_each(|value| {
                    schema_value_str.push_str(format!("{}, ", value.to_string()).as_str());
                });
                schema_value_str.push_str("]");
            }
            Self::Object(data) => {
                schema_value_str.push_str("{ ");
                data.iter().for_each(|(key, value)| {
                    schema_value_str.push_str(format!("{}: {}, ", key, value.to_string()).as_str());
                });
                schema_value_str.push_str("}");
            }
            Self::Pubkey(data) => schema_value_str.push_str(data.to_string().as_str()),
            Self::HashMap(data) => {
                schema_value_str.push_str("Map { ");
                data.iter().for_each(|(key, value)| {
                    schema_value_str.push_str(format!("{}: {}, ", key, value.to_string()).as_str());
                });
                schema_value_str.push_str("}");
            }
            Self::Enum(kind, value) => {
                schema_value_str.push_str(&format!("Enum {{ {}: {} }}", kind, value.to_string()))
            }
        };
        schema_value_str
    }
}

impl AnchorSerialize for SchemaValue {
    #[inline]
    fn serialize<W: std::io::prelude::Write>(&self, writer: &mut W) -> std::io::Result<()> {
        match self {
            Self::Null => 0u8.serialize(writer),
            Self::Bool(value) => {
                1u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::Number(value) => {
                2u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::String(value) => {
                3u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::Array(value) => {
                4u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::Object(value) => {
                5u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::Pubkey(value) => {
                6u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::HashMap(value) => {
                7u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::Enum(kind, value) => {
                8u8.serialize(writer)?;
                kind.serialize(writer)?;
                value.serialize(writer)
            }
        }
    }
}

impl AnchorDeserialize for SchemaValue {
    #[inline]
    fn deserialize_reader<R: std::io::prelude::Read>(reader: &mut R) -> std::io::Result<Self> {
        let mut buf: [u8; 1] = [0; 1];
        let limit = reader.read(&mut buf)?;
        if limit != 1 {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "Unexpected length of input",
            ));
        }
        match buf[0] {
            0 => Ok(Self::Null),
            1 => bool::deserialize_reader(reader).and_then(|v| Ok(Self::Bool(v))),
            2 => Number::deserialize_reader(reader).and_then(|v| Ok(Self::Number(v))),
            3 => String::deserialize_reader(reader).and_then(|v| Ok(Self::String(v))),
            4 => Vec::deserialize_reader(reader).and_then(|v| Ok(Self::Array(v))),
            5 => HashMap::deserialize_reader(reader).and_then(|v| Ok(Self::Object(v))),
            6 => Pubkey::deserialize_reader(reader).and_then(|v| Ok(Self::Pubkey(v))),
            7 => HashMap::deserialize_reader(reader).and_then(|v| Ok(Self::HashMap(v))),
            8 => Ok(Self::Enum(
                String::deserialize_reader(reader)?,
                Box::new(Self::deserialize_reader(reader)?),
            )),
            _ => {
                let msg = format!(
                    "Invalid Option representation: {}. The first byte must be 0 till 6",
                    buf[0]
                );
                Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, msg))
            }
        }
    }
}
const PUBKEY_PREFIX: &'static str = "pubkey:";
const HASHMAP_PREFIX: &'static str = "_is_a_hashmap_hc";
impl Into<serde_json::Value> for SchemaValue {
    fn into(self) -> serde_json::Value {
        match self {
            Self::Null => serde_json::Value::Null,
            Self::Bool(value) => serde_json::Value::Bool(value),
            Self::Number(value) => serde_json::Value::Number(value.into()),
            Self::String(value) => serde_json::Value::String(value),
            Self::Array(value) => {
                serde_json::Value::Array(value.into_iter().map(|i| i.into()).collect())
            }
            Self::Object(value) => {
                serde_json::Value::Object(value.into_iter().map(|(k, v)| (k, v.into())).collect())
            }
            Self::Pubkey(value) => {
                let string = format!("{}{}", PUBKEY_PREFIX, value.to_string());
                serde_json::Value::String(string)
            }
            Self::HashMap(value) => {
                let mut obj: serde_json::Map<String, serde_json::Value> =
                    value.into_iter().map(|(k, v)| (k, v.into())).collect();
                obj.insert(HASHMAP_PREFIX.to_string(), serde_json::Value::Bool(true));
                serde_json::Value::Object(obj)
            }
            Self::Enum(kind, value) => {
                let mut obj = serde_json::Map::<String, serde_json::Value>::new();
                obj.insert(String::from("__kind"), kind.into());

                let schema_value = value.as_ref().to_owned();
                obj.insert(String::from("params"), schema_value.into());

                serde_json::Value::Object(obj)
            }
        }
    }
}

impl From<serde_json::Value> for SchemaValue {
    fn from(value: serde_json::Value) -> Self {
        match value {
            serde_json::Value::Null => Self::Null,
            serde_json::Value::Bool(value) => Self::Bool(value),
            serde_json::Value::Number(value) => Self::Number(value.into()),
            serde_json::Value::String(value) => {
                if value.starts_with(PUBKEY_PREFIX) && value.len() == (PUBKEY_PREFIX.len() + 32) {
                    let pubkey = Pubkey::from_str(value.split_at(PUBKEY_PREFIX.len()).1).unwrap();
                    return Self::Pubkey(pubkey);
                }
                Self::String(value)
            }
            serde_json::Value::Array(value) => {
                Self::Array(value.into_iter().map(|i| i.into()).collect())
            }
            serde_json::Value::Object(mut value) => {
                if let Some(prefix) = value.get(HASHMAP_PREFIX) {
                    if let Some(prefix) = prefix.as_bool() {
                        if prefix {
                            value.remove(HASHMAP_PREFIX);
                            return Self::HashMap(
                                value.into_iter().map(|(k, v)| (k, v.into())).collect(),
                            );
                        }
                    }
                } else if let Some(kind) = value.get("__kind") {
                    let params = if let Some(params) = value.get("params") {
                        params.to_owned().into()
                    } else {
                        Self::Null
                    };
                    return Self::Enum(kind.to_string(), Box::new(params));
                }
                Self::Object(value.into_iter().map(|(k, v)| (k, v.into())).collect())
            }
        }
    }
}