use crate::duties_service::DutiesService;
use beacon_node_fallback::{ApiTopic, BeaconNodeFallback};
use bls::PublicKeyBytes;
use eth2::types::BlockId;
use futures::StreamExt;
use futures::future::FutureExt;
use logging::crit;
use slot_clock::SlotClock;
use std::collections::HashMap;
use std::ops::Deref;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use task_executor::TaskExecutor;
use tokio::time::{Duration, Instant, sleep, sleep_until};
use tracing::{Instrument, debug, error, info, info_span, instrument, trace, warn};
use types::{
    ChainSpec, EthSpec, Hash256, Slot, SyncCommitteeSubscription, SyncContributionData, SyncDuty,
    SyncSelectionProof, SyncSubnetId,
};
use validator_store::{ContributionToSign, SyncMessageToSign, ValidatorStore};

pub const SUBSCRIPTION_LOOKAHEAD_EPOCHS: u64 = 4;

pub struct SyncCommitteeService<S: ValidatorStore, T: SlotClock + 'static> {
    inner: Arc<Inner<S, T>>,
}

impl<S: ValidatorStore, T: SlotClock + 'static> Clone for SyncCommitteeService<S, T> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

impl<S: ValidatorStore, T: SlotClock + 'static> Deref for SyncCommitteeService<S, T> {
    type Target = Inner<S, T>;

    fn deref(&self) -> &Self::Target {
        self.inner.deref()
    }
}

pub struct Inner<S: ValidatorStore, T: SlotClock + 'static> {
    duties_service: Arc<DutiesService<S, T>>,
    validator_store: Arc<S>,
    slot_clock: T,
    beacon_nodes: Arc<BeaconNodeFallback<T>>,
    executor: TaskExecutor,
    /// Boolean to track whether the service has posted subscriptions to the BN at least once.
    ///
    /// This acts as a latch that fires once upon start-up, and then never again.
    first_subscription_done: AtomicBool,
}

impl<S: ValidatorStore + 'static, T: SlotClock + 'static> SyncCommitteeService<S, T> {
    pub fn new(
        duties_service: Arc<DutiesService<S, T>>,
        validator_store: Arc<S>,
        slot_clock: T,
        beacon_nodes: Arc<BeaconNodeFallback<T>>,
        executor: TaskExecutor,
    ) -> Self {
        Self {
            inner: Arc::new(Inner {
                duties_service,
                validator_store,
                slot_clock,
                beacon_nodes,
                executor,
                first_subscription_done: AtomicBool::new(false),
            }),
        }
    }

    /// Check if the Altair fork has been activated and therefore sync duties should be performed.
    ///
    /// Slot clock errors are mapped to `false`.
    fn altair_fork_activated(&self) -> bool {
        self.duties_service
            .spec
            .altair_fork_epoch
            .and_then(|fork_epoch| {
                let current_epoch = self.slot_clock.now()?.epoch(S::E::slots_per_epoch());
                Some(current_epoch >= fork_epoch)
            })
            .unwrap_or(false)
    }

    pub fn start_update_service(self, spec: &ChainSpec) -> Result<(), String> {
        if self.duties_service.disable_attesting {
            info!("Sync committee service disabled");
            return Ok(());
        }

        let slot_duration = spec.get_slot_duration();
        let duration_to_next_slot = self
            .slot_clock
            .duration_to_next_slot()
            .ok_or("Unable to determine duration to next slot")?;

        info!(
            next_update_millis = duration_to_next_slot.as_millis(),
            "Sync committee service started"
        );

        let executor = self.executor.clone();

        let sync_message_slot_component = spec.get_sync_message_due();

        let interval_fut = async move {
            loop {
                if let Some(duration_to_next_slot) = self.slot_clock.duration_to_next_slot() {
                    // Wait for contribution broadcast interval 1/3 of the way through the slot.
                    sleep(duration_to_next_slot + sync_message_slot_component).await;

                    // Do nothing if the Altair fork has not yet occurred.
                    if !self.altair_fork_activated() {
                        continue;
                    }

                    if let Err(e) = self.spawn_contribution_tasks().await {
                        crit!(
                            error = ?e,
                            "Failed to spawn sync contribution tasks"
                        )
                    } else {
                        trace!("Spawned sync contribution tasks")
                    }

                    // Do subscriptions for future slots/epochs.
                    self.spawn_subscription_tasks();
                } else {
                    error!("Failed to read slot clock");
                    // If we can't read the slot clock, just wait another slot.
                    sleep(slot_duration).await;
                }
            }
        };

        executor.spawn(interval_fut, "sync_committee_service");
        Ok(())
    }

    async fn spawn_contribution_tasks(&self) -> Result<(), String> {
        let spec = &self.duties_service.spec;
        let slot = self.slot_clock.now().ok_or("Failed to read slot clock")?;
        let duration_to_next_slot = self
            .slot_clock
            .duration_to_next_slot()
            .ok_or("Unable to determine duration to next slot")?;

        // If a validator needs to publish a sync aggregate, they must do so at 2/3
        // through the slot. This delay triggers at this time
        let aggregate_production_instant = Instant::now()
            + duration_to_next_slot
                .checked_add(spec.get_contribution_message_due())
                .and_then(|offset| offset.checked_sub(spec.get_slot_duration()))
                .unwrap_or_else(|| Duration::from_secs(0));

        let Some(slot_duties) = self
            .duties_service
            .sync_duties
            .get_duties_for_slot::<S::E>(slot, &self.duties_service.spec)
        else {
            debug!("No duties known for slot {}", slot);
            return Ok(());
        };

        if slot_duties.duties.is_empty() {
            debug!(%slot, "No local validators in current sync committee");
            return Ok(());
        }

        // Fetch `block_root` with non optimistic execution for `SyncCommitteeContribution`.
        let response = self
            .beacon_nodes
            .first_success(
                |beacon_node| async move {
                    match beacon_node.get_beacon_blocks_root(BlockId::Head).await {
                        Ok(Some(block)) if block.execution_optimistic == Some(false) => {
                            Ok(block)
                        }
                        Ok(Some(_)) => {
                            Err(format!("To sign sync committee messages for slot {slot} a non-optimistic head block is required"))
                        }
                        Ok(None) => Err(format!("No block root found for slot {}", slot)),
                        Err(e) => Err(e.to_string()),
                    }
                },
            )
            .await;

        let block_root = match response {
            Ok(block) => block.data.root,
            Err(errs) => {
                warn!(
                    errors = errs.to_string(),
                    %slot,
                    "Refusing to sign sync committee messages for an optimistic head block or \
                    a block head with unknown optimistic status"
                );
                return Ok(());
            }
        };

        // Spawn one task to publish all of the sync committee signatures.
        let validator_duties = slot_duties.duties;
        let service = self.clone();
        self.inner.executor.spawn(
            async move {
                service
                    .publish_sync_committee_signatures(slot, block_root, validator_duties)
                    .map(|_| ())
                    .await
            }
            .instrument(info_span!("lh_sync_committee_signature_publish", %slot)),
            "sync_committee_signature_publish",
        );

        let aggregators = slot_duties.aggregators;
        let service = self.clone();
        self.inner.executor.spawn(
            async move {
                service
                    .publish_sync_committee_aggregates(
                        slot,
                        block_root,
                        aggregators,
                        aggregate_production_instant,
                    )
                    .map(|_| ())
                    .await
            }
            .instrument(info_span!("lh_sync_committee_aggregate_publish", %slot)),
            "sync_committee_aggregate_publish",
        );

        Ok(())
    }

    /// Publish sync committee signatures.
    #[instrument(skip_all, fields(%slot, ?beacon_block_root))]
    async fn publish_sync_committee_signatures(
        &self,
        slot: Slot,
        beacon_block_root: Hash256,
        validator_duties: Vec<SyncDuty>,
    ) -> Result<(), ()> {
        let messages_to_sign: Vec<_> = validator_duties
            .iter()
            .map(|duty| SyncMessageToSign {
                slot,
                beacon_block_root,
                validator_index: duty.validator_index,
                pubkey: duty.pubkey,
            })
            .collect();

        let signature_stream = self
            .validator_store
            .sign_sync_committee_signatures(messages_to_sign);
        tokio::pin!(signature_stream);

        while let Some(result) = signature_stream.next().await {
            match result {
                Ok(committee_signatures) if !committee_signatures.is_empty() => {
                    let committee_signatures = &committee_signatures;
                    match self
                        .beacon_nodes
                        .request(ApiTopic::SyncCommittee, |beacon_node| async move {
                            beacon_node
                                .post_beacon_pool_sync_committee_signatures(committee_signatures)
                                .await
                        })
                        .instrument(info_span!(
                            "publish_sync_signatures",
                            count = committee_signatures.len()
                        ))
                        .await
                    {
                        Ok(()) => info!(
                            count = committee_signatures.len(),
                            head_block = ?beacon_block_root,
                            %slot,
                            "Successfully published sync committee messages"
                        ),
                        Err(e) => error!(
                            %slot,
                            error = %e,
                            "Unable to publish sync committee messages"
                        ),
                    }
                }
                Err(e) => {
                    crit!(%slot, error = ?e, "Failed to sign sync committee signatures");
                }
                _ => {}
            }
        }

        Ok(())
    }

    async fn publish_sync_committee_aggregates(
        &self,
        slot: Slot,
        beacon_block_root: Hash256,
        aggregators: HashMap<SyncSubnetId, Vec<(u64, PublicKeyBytes, SyncSelectionProof)>>,
        aggregate_instant: Instant,
    ) {
        for (subnet_id, subnet_aggregators) in aggregators {
            let service = self.clone();
            self.inner.executor.spawn(
                async move {
                    service
                        .publish_sync_committee_aggregate_for_subnet(
                            slot,
                            beacon_block_root,
                            subnet_id,
                            subnet_aggregators,
                            aggregate_instant,
                        )
                        .map(|_| ())
                        .await
                }
                .instrument(info_span!("lh_publish_sync_committee_aggregate_for_subnet", %slot, ?beacon_block_root, %subnet_id)),
                "sync_committee_aggregate_publish_subnet",
            );
        }
    }

    async fn publish_sync_committee_aggregate_for_subnet(
        &self,
        slot: Slot,
        beacon_block_root: Hash256,
        subnet_id: SyncSubnetId,
        subnet_aggregators: Vec<(u64, PublicKeyBytes, SyncSelectionProof)>,
        aggregate_instant: Instant,
    ) -> Result<(), ()> {
        sleep_until(aggregate_instant).await;

        let contribution = &self
            .beacon_nodes
            .first_success(|beacon_node| async move {
                let sync_contribution_data = SyncContributionData {
                    slot,
                    beacon_block_root,
                    subcommittee_index: subnet_id.into(),
                };

                beacon_node
                    .get_validator_sync_committee_contribution(&sync_contribution_data)
                    .await
            })
            .instrument(info_span!("fetch_sync_contribution"))
            .await
            .map_err(|e| {
                crit!(
                    %slot,
                    ?beacon_block_root,
                    error = %e,
                    "Failed to produce sync contribution"
                )
            })?
            .ok_or_else(|| {
                crit!(%slot, ?beacon_block_root, "No aggregate contribution found");
            })?
            .data;

        let contributions_to_sign: Vec<_> = subnet_aggregators
            .into_iter()
            .map(
                |(aggregator_index, aggregator_pk, selection_proof)| ContributionToSign {
                    aggregator_index,
                    aggregator_pubkey: aggregator_pk,
                    contribution: contribution.clone(),
                    selection_proof,
                },
            )
            .collect();

        let contribution_stream = self
            .validator_store
            .sign_sync_committee_contributions(contributions_to_sign);
        tokio::pin!(contribution_stream);

        while let Some(result) = contribution_stream.next().await {
            match result {
                Ok(signed_contributions) if !signed_contributions.is_empty() => {
                    let signed_contributions = &signed_contributions;
                    // Publish to the beacon node.
                    match self
                        .beacon_nodes
                        .first_success(|beacon_node| async move {
                            beacon_node
                                .post_validator_contribution_and_proofs(signed_contributions)
                                .await
                        })
                        .instrument(info_span!(
                            "publish_sync_contributions",
                            count = signed_contributions.len()
                        ))
                        .await
                    {
                        Ok(()) => info!(
                            subnet = %subnet_id,
                            beacon_block_root = %beacon_block_root,
                            num_signers = contribution.aggregation_bits.num_set_bits(),
                            %slot,
                            "Successfully published sync contributions"
                        ),
                        Err(e) => error!(
                            %slot,
                            error = %e,
                            "Unable to publish signed contributions and proofs"
                        ),
                    }
                }
                Err(e) => {
                    crit!(%slot, error = ?e, "Failed to sign sync committee contributions");
                }
                _ => {}
            }
        }

        Ok(())
    }

    fn spawn_subscription_tasks(&self) {
        let service = self.clone();

        self.inner.executor.spawn(
            async move {
                service.publish_subscriptions().await.unwrap_or_else(|e| {
                    error!(
                        error = ?e,
                        "Error publishing subscriptions"
                    )
                });
            },
            "sync_committee_subscription_publish",
        );
    }

    async fn publish_subscriptions(self) -> Result<(), String> {
        let spec = &self.duties_service.spec;
        let slot = self.slot_clock.now().ok_or("Failed to read slot clock")?;

        let mut duty_slots = vec![];
        let mut all_succeeded = true;

        // At the start of every epoch during the current period, re-post the subscriptions
        // to the beacon node. This covers the case where the BN has forgotten the subscriptions
        // due to a restart, or where the VC has switched to a fallback BN.
        let current_period = sync_period_of_slot::<S::E>(slot, spec)?;

        if !self.first_subscription_done.load(Ordering::Relaxed)
            || slot.as_u64() % S::E::slots_per_epoch() == 0
        {
            duty_slots.push((slot, current_period));
        }

        // Near the end of the current period, push subscriptions for the next period to the
        // beacon node. We aggressively push every slot in the lead-up, as this is the main way
        // that we want to ensure that the BN is subscribed (well in advance).
        let lookahead_slot = slot + SUBSCRIPTION_LOOKAHEAD_EPOCHS * S::E::slots_per_epoch();

        let lookahead_period = sync_period_of_slot::<S::E>(lookahead_slot, spec)?;

        if lookahead_period > current_period {
            duty_slots.push((lookahead_slot, lookahead_period));
        }

        if duty_slots.is_empty() {
            return Ok(());
        }

        // Collect subscriptions.
        let mut subscriptions = vec![];

        for (duty_slot, sync_committee_period) in duty_slots {
            debug!(%duty_slot, %slot, "Fetching subscription duties");
            match self
                .duties_service
                .sync_duties
                .get_duties_for_slot::<S::E>(duty_slot, spec)
            {
                Some(duties) => subscriptions.extend(subscriptions_from_sync_duties(
                    duties.duties,
                    sync_committee_period,
                    spec,
                )),
                None => {
                    debug!(
                        slot = %duty_slot,
                        "No duties for subscription"
                    );
                    all_succeeded = false;
                }
            }
        }

        if subscriptions.is_empty() {
            debug!(%slot, "No sync subscriptions to send");
            return Ok(());
        }

        // Post subscriptions to BN.
        debug!(
            count = subscriptions.len(),
            "Posting sync subscriptions to BN"
        );
        let subscriptions_slice = &subscriptions;

        for subscription in subscriptions_slice {
            debug!(
                validator_index = subscription.validator_index,
                validator_sync_committee_indices = ?subscription.sync_committee_indices,
                until_epoch = %subscription.until_epoch,
                "Subscription"
            );
        }

        if let Err(e) = self
            .beacon_nodes
            .request(ApiTopic::Subscriptions, |beacon_node| async move {
                beacon_node
                    .post_validator_sync_committee_subscriptions(subscriptions_slice)
                    .await
            })
            .await
        {
            error!(
                %slot,
                error = %e,
                "Unable to post sync committee subscriptions"
            );
            all_succeeded = false;
        }

        // Disable first-subscription latch once all duties have succeeded once.
        if all_succeeded {
            self.first_subscription_done.store(true, Ordering::Relaxed);
        }

        Ok(())
    }
}

fn sync_period_of_slot<E: EthSpec>(slot: Slot, spec: &ChainSpec) -> Result<u64, String> {
    slot.epoch(E::slots_per_epoch())
        .sync_committee_period(spec)
        .map_err(|e| format!("Error computing sync period: {:?}", e))
}

fn subscriptions_from_sync_duties(
    duties: Vec<SyncDuty>,
    sync_committee_period: u64,
    spec: &ChainSpec,
) -> impl Iterator<Item = SyncCommitteeSubscription> {
    let until_epoch = spec.epochs_per_sync_committee_period * (sync_committee_period + 1);
    duties
        .into_iter()
        .map(move |duty| SyncCommitteeSubscription {
            validator_index: duty.validator_index,
            sync_committee_indices: duty.validator_sync_committee_indices,
            until_epoch,
        })
}
