use context_deserialize::{ContextDeserialize, context_deserialize};
use educe::Educe;
use serde::{Deserialize, Deserializer, Serialize};
use ssz::{Decode, Encode};
use ssz_derive::{Decode, Encode};
use superstruct::superstruct;
use tree_hash::Hash256;
use tree_hash_derive::TreeHash;

use crate::{
    block::SignedBlindedBeaconBlock,
    core::{ChainSpec, EthSpec, Slot},
    fork::ForkName,
    light_client::{
        LightClientError, LightClientHeader, LightClientHeaderAltair, LightClientHeaderCapella,
        LightClientHeaderDeneb, LightClientHeaderElectra, LightClientHeaderFulu,
    },
    sync_committee::SyncAggregate,
};

/// A LightClientOptimisticUpdate is the update we send on each slot,
/// it is based off the current unfinalized epoch is verified only against BLS signature.
#[superstruct(
    variants(Altair, Capella, Deneb, Electra, Fulu),
    variant_attributes(
        derive(Debug, Clone, Serialize, Deserialize, Educe, Decode, Encode, TreeHash,),
        educe(PartialEq),
        serde(bound = "E: EthSpec", deny_unknown_fields),
        cfg_attr(
            feature = "arbitrary",
            derive(arbitrary::Arbitrary),
            arbitrary(bound = "E: EthSpec"),
        ),
        context_deserialize(ForkName),
    )
)]
#[cfg_attr(
    feature = "arbitrary",
    derive(arbitrary::Arbitrary),
    arbitrary(bound = "E: EthSpec")
)]
#[derive(Debug, Clone, Serialize, Encode, TreeHash, PartialEq)]
#[serde(untagged)]
#[tree_hash(enum_behaviour = "transparent")]
#[ssz(enum_behaviour = "transparent")]
#[serde(bound = "E: EthSpec", deny_unknown_fields)]
pub struct LightClientOptimisticUpdate<E: EthSpec> {
    /// The last `BeaconBlockHeader` from the last attested block by the sync committee.
    #[superstruct(only(Altair), partial_getter(rename = "attested_header_altair"))]
    pub attested_header: LightClientHeaderAltair<E>,
    #[superstruct(only(Capella), partial_getter(rename = "attested_header_capella"))]
    pub attested_header: LightClientHeaderCapella<E>,
    #[superstruct(only(Deneb), partial_getter(rename = "attested_header_deneb"))]
    pub attested_header: LightClientHeaderDeneb<E>,
    #[superstruct(only(Electra), partial_getter(rename = "attested_header_electra"))]
    pub attested_header: LightClientHeaderElectra<E>,
    #[superstruct(only(Fulu), partial_getter(rename = "attested_header_fulu"))]
    pub attested_header: LightClientHeaderFulu<E>,
    /// current sync aggregate
    pub sync_aggregate: SyncAggregate<E>,
    /// Slot of the sync aggregated signature
    #[superstruct(getter(copy))]
    pub signature_slot: Slot,
}

impl<E: EthSpec> LightClientOptimisticUpdate<E> {
    pub fn new(
        attested_block: &SignedBlindedBeaconBlock<E>,
        sync_aggregate: SyncAggregate<E>,
        signature_slot: Slot,
        chain_spec: &ChainSpec,
    ) -> Result<Self, LightClientError> {
        let optimistic_update = match attested_block
            .fork_name(chain_spec)
            .map_err(|_| LightClientError::InconsistentFork)?
        {
            ForkName::Altair | ForkName::Bellatrix => {
                Self::Altair(LightClientOptimisticUpdateAltair {
                    attested_header: LightClientHeaderAltair::block_to_light_client_header(
                        attested_block,
                    )?,
                    sync_aggregate,
                    signature_slot,
                })
            }
            ForkName::Capella => Self::Capella(LightClientOptimisticUpdateCapella {
                attested_header: LightClientHeaderCapella::block_to_light_client_header(
                    attested_block,
                )?,
                sync_aggregate,
                signature_slot,
            }),
            ForkName::Deneb => Self::Deneb(LightClientOptimisticUpdateDeneb {
                attested_header: LightClientHeaderDeneb::block_to_light_client_header(
                    attested_block,
                )?,
                sync_aggregate,
                signature_slot,
            }),
            ForkName::Electra => Self::Electra(LightClientOptimisticUpdateElectra {
                attested_header: LightClientHeaderElectra::block_to_light_client_header(
                    attested_block,
                )?,
                sync_aggregate,
                signature_slot,
            }),
            ForkName::Fulu => Self::Fulu(LightClientOptimisticUpdateFulu {
                attested_header: LightClientHeaderFulu::block_to_light_client_header(
                    attested_block,
                )?,
                sync_aggregate,
                signature_slot,
            }),
            ForkName::Gloas => return Err(LightClientError::GloasNotImplemented),
            ForkName::Base => return Err(LightClientError::AltairForkNotActive),
        };

        Ok(optimistic_update)
    }

    pub fn map_with_fork_name<F, R>(&self, func: F) -> R
    where
        F: Fn(ForkName) -> R,
    {
        match self {
            Self::Altair(_) => func(ForkName::Altair),
            Self::Capella(_) => func(ForkName::Capella),
            Self::Deneb(_) => func(ForkName::Deneb),
            Self::Electra(_) => func(ForkName::Electra),
            Self::Fulu(_) => func(ForkName::Fulu),
        }
    }

    pub fn get_slot<'a>(&'a self) -> Slot {
        map_light_client_optimistic_update_ref!(&'a _, self.to_ref(), |inner, cons| {
            cons(inner);
            inner.attested_header.beacon.slot
        })
    }

    pub fn get_canonical_root<'a>(&'a self) -> Hash256 {
        map_light_client_optimistic_update_ref!(&'a _, self.to_ref(), |inner, cons| {
            cons(inner);
            inner.attested_header.beacon.canonical_root()
        })
    }

    pub fn get_parent_root<'a>(&'a self) -> Hash256 {
        map_light_client_optimistic_update_ref!(&'a _, self.to_ref(), |inner, cons| {
            cons(inner);
            inner.attested_header.beacon.parent_root
        })
    }

    pub fn from_ssz_bytes(bytes: &[u8], fork_name: ForkName) -> Result<Self, ssz::DecodeError> {
        let optimistic_update = match fork_name {
            ForkName::Altair | ForkName::Bellatrix => {
                Self::Altair(LightClientOptimisticUpdateAltair::from_ssz_bytes(bytes)?)
            }
            ForkName::Capella => {
                Self::Capella(LightClientOptimisticUpdateCapella::from_ssz_bytes(bytes)?)
            }
            ForkName::Deneb => {
                Self::Deneb(LightClientOptimisticUpdateDeneb::from_ssz_bytes(bytes)?)
            }
            ForkName::Electra => {
                Self::Electra(LightClientOptimisticUpdateElectra::from_ssz_bytes(bytes)?)
            }
            ForkName::Fulu => Self::Fulu(LightClientOptimisticUpdateFulu::from_ssz_bytes(bytes)?),
            // TODO(gloas): implement Gloas light client
            ForkName::Base | ForkName::Gloas => {
                return Err(ssz::DecodeError::BytesInvalid(format!(
                    "LightClientOptimisticUpdate decoding for {fork_name} not implemented"
                )));
            }
        };

        Ok(optimistic_update)
    }

    #[allow(clippy::arithmetic_side_effects)]
    pub fn ssz_max_len_for_fork(fork_name: ForkName) -> usize {
        let fixed_len = match fork_name {
            ForkName::Base => 0,
            ForkName::Altair | ForkName::Bellatrix => {
                <LightClientOptimisticUpdateAltair<E> as Encode>::ssz_fixed_len()
            }
            ForkName::Capella => <LightClientOptimisticUpdateCapella<E> as Encode>::ssz_fixed_len(),
            ForkName::Deneb => <LightClientOptimisticUpdateDeneb<E> as Encode>::ssz_fixed_len(),
            ForkName::Electra => <LightClientOptimisticUpdateElectra<E> as Encode>::ssz_fixed_len(),
            ForkName::Fulu => <LightClientOptimisticUpdateFulu<E> as Encode>::ssz_fixed_len(),
            // TODO(gloas): implement Gloas light client
            ForkName::Gloas => 0,
        };
        fixed_len + LightClientHeader::<E>::ssz_max_var_len_for_fork(fork_name)
    }

    // Implements spec prioritization rules:
    // > Full nodes SHOULD provide the LightClientOptimisticUpdate with the highest attested_header.beacon.slot (if multiple, highest signature_slot)
    //
    // ref: https://github.com/ethereum/consensus-specs/blob/113c58f9bf9c08867f6f5f633c4d98e0364d612a/specs/altair/light-client/full-node.md#create_light_client_optimistic_update
    pub fn is_latest(&self, attested_slot: Slot, signature_slot: Slot) -> bool {
        let prev_slot = self.get_slot();
        if attested_slot > prev_slot {
            true
        } else {
            attested_slot == prev_slot && signature_slot > self.signature_slot()
        }
    }
}

impl<'de, E: EthSpec> ContextDeserialize<'de, ForkName> for LightClientOptimisticUpdate<E> {
    fn context_deserialize<D>(deserializer: D, context: ForkName) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let convert_err = |e| {
            serde::de::Error::custom(format!(
                "LightClientOptimisticUpdate failed to deserialize: {:?}",
                e
            ))
        };
        Ok(match context {
            ForkName::Base => {
                return Err(serde::de::Error::custom(format!(
                    "LightClientOptimisticUpdate failed to deserialize: unsupported fork '{}'",
                    context
                )));
            }
            ForkName::Altair | ForkName::Bellatrix => {
                Self::Altair(Deserialize::deserialize(deserializer).map_err(convert_err)?)
            }
            ForkName::Capella => {
                Self::Capella(Deserialize::deserialize(deserializer).map_err(convert_err)?)
            }
            ForkName::Deneb => {
                Self::Deneb(Deserialize::deserialize(deserializer).map_err(convert_err)?)
            }
            ForkName::Electra => {
                Self::Electra(Deserialize::deserialize(deserializer).map_err(convert_err)?)
            }
            ForkName::Fulu => {
                Self::Fulu(Deserialize::deserialize(deserializer).map_err(convert_err)?)
            }
            ForkName::Gloas => {
                // TODO(EIP-7732): check if this is correct
                return Err(serde::de::Error::custom(format!(
                    "LightClientBootstrap failed to deserialize: unsupported fork '{}'",
                    context
                )));
            }
        })
    }
}

#[cfg(test)]
mod tests {
    // `ssz_tests!` can only be defined once per namespace
    #[cfg(test)]
    mod altair {
        use crate::{LightClientOptimisticUpdateAltair, MainnetEthSpec};
        ssz_tests!(LightClientOptimisticUpdateAltair<MainnetEthSpec>);
    }

    #[cfg(test)]
    mod capella {
        use crate::{LightClientOptimisticUpdateCapella, MainnetEthSpec};
        ssz_tests!(LightClientOptimisticUpdateCapella<MainnetEthSpec>);
    }

    #[cfg(test)]
    mod deneb {
        use crate::{LightClientOptimisticUpdateDeneb, MainnetEthSpec};
        ssz_tests!(LightClientOptimisticUpdateDeneb<MainnetEthSpec>);
    }

    #[cfg(test)]
    mod electra {
        use crate::{LightClientOptimisticUpdateElectra, MainnetEthSpec};
        ssz_tests!(LightClientOptimisticUpdateElectra<MainnetEthSpec>);
    }

    #[cfg(test)]
    mod fulu {
        use crate::{LightClientOptimisticUpdateFulu, MainnetEthSpec};
        ssz_tests!(LightClientOptimisticUpdateFulu<MainnetEthSpec>);
    }
}
