use std::collections::HashMap;

use anchor_lang::prelude::*;
// use anchor_lang::{AnchorDeserialize, AnchorSerialize};
use spl_account_compression::{wrap_application_data_v1, Node, Noop};

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

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")
        }
    }
}

pub enum Schema {
    Null,
    Bool,
    Number,
    String,
    Pubkey,
    Array(Box<Schema>),
    Object(HashMap<String, Schema>),
}

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

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::Pubkey => 4u8.serialize(writer),
            Self::Array(value) => {
                5u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::Object(value) => {
                6u8.serialize(writer)?;
                value.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 => Ok(Self::Pubkey),
            5 => Schema::deserialize_reader(reader).and_then(|v| Ok(Self::Array(Box::new(v)))),
            6 => HashMap::deserialize_reader(reader).and_then(|v| Ok(Self::Object(v))),
            _ => {
                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),
    Pubkey(Pubkey),
    Array(Vec<SchemaValue>),
    Object(HashMap<String, SchemaValue>),
}

impl SchemaValue {
    pub fn validate(&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::Pubkey => matches!(self, Self::Pubkey(_)),
            Schema::Array(schema) => match self {
                Self::Array(arr) => arr.iter().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(key) {
                        return value.validate(schema);
                    }
                    false
                }),
                _ => false,
            },
        }
    }
}

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::Pubkey(value) => {
                4u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::Array(value) => {
                5u8.serialize(writer)?;
                value.serialize(writer)
            }
            Self::Object(value) => {
                6u8.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 => Pubkey::deserialize_reader(reader).and_then(|v| Ok(Self::Pubkey(v))),
            5 => Vec::deserialize_reader(reader).and_then(|v| Ok(Self::Array(v))),
            6 => HashMap::deserialize_reader(reader).and_then(|v| Ok(Self::Object(v))),
            _ => {
                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))
            }
        }
    }
}

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::Pubkey(value) => {
                let bytes = &mut &[&[0, 0], &value.to_bytes()[..]].concat()[..];
                serde_json::Value::String(String::deserialize(bytes).unwrap())
            }
            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())
            }
        }
    }
}

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) => {
                let bytes = value.try_to_vec().unwrap();
                if bytes.len() == 34 && bytes[0..2] == [0, 0] {
                    let mut pubkey_bytes: [u8; 32] = [0; 32];
                    pubkey_bytes.copy_from_slice(&bytes[2..]);
                    return Self::Pubkey(Pubkey::from(pubkey_bytes));
                }
                Self::String(value)
            }
            serde_json::Value::Array(value) => {
                Self::Array(value.into_iter().map(|i| i.into()).collect())
            }
            serde_json::Value::Object(value) => {
                Self::Object(value.into_iter().map(|(k, v)| (k, v.into())).collect())
            }
        }
    }
}

pub trait Compress {
    fn to_bytes(&self) -> Vec<u8>
    where
        Self: AnchorSerialize,
    {
        self.try_to_vec().unwrap()
    }

    fn size(&self) -> usize
    where
        Self: AnchorSerialize,
    {
        self.to_bytes().len()
    }

    fn compress(&self) -> Node;

    fn schema() -> Schema;

    fn schema_value(&self) -> SchemaValue;
}

pub trait CompressedData {
    fn event_stream(&self) -> CompressedDataEventStream
    where
        Self: AnchorSerialize + Compress,
    {
        CompressedDataEventStream::Full {
            data: self.schema_value(),
        }
    }
}

pub trait CompressedDataChunk {
    const KEY: &'static str;

    fn event_stream(&self) -> CompressedDataEventStream
    where
        Self: AnchorSerialize + Compress,
    {
        CompressedDataEventStream::PatchChunk {
            key: Self::KEY.to_string(),
            data: self.schema_value(),
        }
    }
}

#[derive(AnchorDeserialize, AnchorSerialize)]
pub enum CompressedDataEventStream {
    Full { data: SchemaValue },
    PatchChunk { key: String, data: SchemaValue },
}

#[derive(AnchorDeserialize, AnchorSerialize)]
pub enum CompressedDataEvent {
    Leaf {
        slot: u64,
        tree_id: [u8; 32],
        leaf_idx: u32,
        seq: u64,
        stream_type: CompressedDataEventStream,
    },
    TreeSchemaValue {
        tree_id: [u8; 32],
        schema: Schema,
    },
}
impl CompressedDataEvent {
    pub fn wrap<'info>(&self, noop: &Program<'info, Noop>) -> Result<()> {
        wrap_application_data_v1(self.try_to_vec().unwrap(), noop)
    }
}