use super::{Config, Mailbox, Message, Round};
use crate::{
    simplex::{
        actors::voter,
        config::ForwardingPolicy,
        interesting,
        metrics::{Inbound, Peer, TimeoutReason},
        scheme::Scheme,
        types::{Activity, Certificate, Proposal, Vote},
        Plan,
    },
    types::{Epoch, Participant, View, ViewDelta},
    Epochable, Relay, Reporter, Viewable,
};
use commonware_cryptography::Digest;
use commonware_macros::select_loop;
use commonware_p2p::{utils::codec::WrappedReceiver, Blocker, Receiver};
use commonware_parallel::Strategy;
use commonware_runtime::{
    spawn_cell,
    telemetry::metrics::{
        histogram::{self, Buckets},
        status::GaugeExt,
    },
    Clock, ContextCell, Handle, Metrics, Spawner,
};
use commonware_utils::{
    channel::{fallible::OneshotExt, mpsc},
    ordered::{Quorum, Set},
};
use prometheus_client::metrics::{
    counter::Counter, family::Family, gauge::Gauge, histogram::Histogram,
};
use rand_core::CryptoRngCore;
use std::{collections::BTreeMap, sync::Arc};
use tracing::{debug, trace};

/// Tracks the current view, its leader, and whether the voter has
/// already been told to timeout this view.
struct Current {
    view: View,
    leader: Option<Participant>,
    timed_out: bool,
}

pub struct Actor<E, S, B, D, Re, Rl, T>
where
    E: Spawner + Metrics + Clock + CryptoRngCore,
    S: Scheme<D>,
    B: Blocker<PublicKey = S::PublicKey>,
    D: Digest,
    Re: Reporter<Activity = Activity<S, D>>,
    Rl: Relay,
    T: Strategy,
{
    context: ContextCell<E>,

    participants: Set<S::PublicKey>,
    scheme: S,

    blocker: B,
    reporter: Re,
    relay: Rl,
    strategy: T,

    activity_timeout: ViewDelta,
    skip_timeout: ViewDelta,
    forwarding: ForwardingPolicy,
    epoch: Epoch,

    mailbox_receiver: mpsc::Receiver<Message<S, D>>,

    added: Counter,
    verified: Counter,
    inbound_messages: Family<Inbound, Counter>,
    latest_vote: Family<Peer, Gauge>,
    latest_seen: Vec<View>,
    batch_size: Histogram,
    verify_latency: histogram::Timed<E>,
    recover_latency: histogram::Timed<E>,
}

impl<E, S, B, D, Re, Rl, T> Actor<E, S, B, D, Re, Rl, T>
where
    E: Spawner + Metrics + Clock + CryptoRngCore,
    S: Scheme<D>,
    B: Blocker<PublicKey = S::PublicKey>,
    D: Digest,
    Re: Reporter<Activity = Activity<S, D>>,
    Rl: Relay<Digest = D, PublicKey = S::PublicKey, Plan = Plan<S::PublicKey>>,
    T: Strategy,
{
    pub fn new(context: E, cfg: Config<S, B, Re, Rl, T>) -> (Self, Mailbox<S, D>) {
        let participants = cfg.scheme.participants().clone();
        let participant_count = participants.len();
        let added = Counter::default();
        let verified = Counter::default();
        let inbound_messages = Family::<Inbound, Counter>::default();
        let batch_size =
            Histogram::new([1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0]);
        context.register(
            "added",
            "number of messages added to the verifier",
            added.clone(),
        );
        context.register("verified", "number of messages verified", verified.clone());
        context.register(
            "inbound_messages",
            "number of inbound messages",
            inbound_messages.clone(),
        );
        let latest_vote = Family::<Peer, Gauge>::default();
        context.register(
            "latest_vote",
            "view of latest vote received per peer",
            latest_vote.clone(),
        );
        for participant in participants.iter() {
            latest_vote.get_or_create(&Peer::new(participant)).set(0);
        }
        context.register(
            "batch_size",
            "number of messages in a signature verification batch",
            batch_size.clone(),
        );
        let verify_latency = Histogram::new(Buckets::CRYPTOGRAPHY);
        context.register(
            "verify_latency",
            "latency of signature verification",
            verify_latency.clone(),
        );
        let recover_latency = Histogram::new(Buckets::CRYPTOGRAPHY);
        context.register(
            "recover_latency",
            "certificate recover latency",
            recover_latency.clone(),
        );
        // TODO(#1833): Metrics should use the post-start context
        let clock = Arc::new(context.clone());
        let (sender, receiver) = mpsc::channel(cfg.mailbox_size);
        (
            Self {
                context: ContextCell::new(context),

                participants,
                scheme: cfg.scheme,

                blocker: cfg.blocker,
                reporter: cfg.reporter,
                relay: cfg.relay,
                strategy: cfg.strategy,

                activity_timeout: cfg.activity_timeout,
                skip_timeout: cfg.skip_timeout,
                forwarding: cfg.forwarding,
                epoch: cfg.epoch,

                mailbox_receiver: receiver,

                added,
                verified,
                inbound_messages,
                latest_vote,
                latest_seen: vec![View::zero(); participant_count],
                batch_size,
                verify_latency: histogram::Timed::new(verify_latency, clock.clone()),
                recover_latency: histogram::Timed::new(recover_latency, clock),
            },
            Mailbox::new(sender),
        )
    }

    fn new_round(&self) -> Round<S, B, D, Re> {
        Round::new(
            self.participants.clone(),
            self.scheme.clone(),
            self.blocker.clone(),
            self.reporter.clone(),
        )
    }

    /// Records the latest view message received from a participant.
    ///
    /// This mechanism is not resistant to malicious validators (nor is
    /// it meant to be). If a peer sends us a certificate very far in the future,
    /// we will record that as their latest activity (and not attempt to skip them).
    fn record_activity(&mut self, sender: &S::PublicKey, view: View) {
        let Some(participant) = self.participants.index(sender) else {
            return;
        };
        let seen_view = &mut self.latest_seen[usize::from(participant)];
        if *seen_view < view {
            *seen_view = view;
        }
    }

    /// Returns true if the participant has sent a recent message.
    fn is_active(
        &self,
        work: &BTreeMap<View, Round<S, B, D, Re>>,
        view: View,
        participant: Participant,
    ) -> bool {
        // Until we have tracked `skip_timeout` views, everyone stays active so startup does not
        // immediately skip leaders whose `latest_seen` entry is still at the default view.
        if work.len() < self.skip_timeout.get() as usize {
            return true;
        }
        let seen_view = self.latest_seen[usize::from(participant)];
        view.get().saturating_sub(seen_view.get()) < self.skip_timeout.get()
    }

    /// Maps `missing` participants to targeted forward recipients, excluding self.
    fn forward_recipients(&self, missing: &[Participant]) -> Vec<S::PublicKey> {
        let me = self.scheme.me();
        missing
            .iter()
            .filter(|&&p| Some(p) != me)
            .filter_map(|&p| self.participants.key(p).cloned())
            .collect()
    }

    /// Selects forwarding targets for a certified proposal under the active policy.
    fn forward_targets(
        &self,
        round: &Round<S, B, D, Re>,
        proposal: &Proposal<D>,
        next_leader: Participant,
    ) -> Vec<Participant> {
        match self.forwarding {
            ForwardingPolicy::Disabled => Vec::new(),
            ForwardingPolicy::SilentVoters => round.missing_voters(proposal),
            ForwardingPolicy::SilentLeader => round
                .is_missing_voter(proposal, next_leader)
                .then_some(next_leader)
                .into_iter()
                .collect(),
        }
    }

    /// Forwards a proposal to the requested peers.
    async fn forward_proposal(&mut self, proposal: Proposal<D>, missing: Vec<Participant>) {
        let peers = self.forward_recipients(&missing);
        if peers.is_empty() {
            return;
        }
        self.relay
            .broadcast(
                proposal.payload,
                Plan::Forward {
                    round: proposal.round,
                    peers,
                },
            )
            .await;
    }

    /// Returns true if the leader has nullified the current view
    /// and we have not yet notified the voter.
    fn leader_nullified(current: &Current, work: &BTreeMap<View, Round<S, B, D, Re>>) -> bool {
        if current.timed_out {
            return false;
        }
        let Some(leader) = current.leader else {
            return false;
        };
        work.get(&current.view)
            .is_some_and(|round| round.has_nullify(leader))
    }

    pub fn start(
        mut self,
        voter: voter::Mailbox<S, D>,
        vote_receiver: impl Receiver<PublicKey = S::PublicKey>,
        certificate_receiver: impl Receiver<PublicKey = S::PublicKey>,
    ) -> Handle<()> {
        spawn_cell!(
            self.context,
            self.run(voter, vote_receiver, certificate_receiver).await
        )
    }

    pub async fn run(
        mut self,
        mut voter: voter::Mailbox<S, D>,
        vote_receiver: impl Receiver<PublicKey = S::PublicKey>,
        certificate_receiver: impl Receiver<PublicKey = S::PublicKey>,
    ) {
        // Wrap channels
        let mut vote_receiver: WrappedReceiver<_, Vote<S, D>> =
            WrappedReceiver::new((), vote_receiver);
        let mut certificate_receiver: WrappedReceiver<_, Certificate<S, D>> =
            WrappedReceiver::new(self.scheme.certificate_codec_config(), certificate_receiver);

        // Initialize view data structures
        let mut current = Current {
            view: View::zero(),
            leader: None,
            timed_out: false,
        };
        let mut finalized = View::zero();
        let mut work: BTreeMap<View, Round<S, B, D, Re>> = BTreeMap::new();
        select_loop! {
            self.context,
            on_start => {
                // Track which view was modified (if any) for certificate construction
                let updated_view;
            },
            on_stopped => {
                debug!("context shutdown, stopping batcher");
            },
            Some(message) = self.mailbox_receiver.recv() else break => match message {
                Message::Update {
                    current: new_current,
                    leader,
                    finalized: new_finalized,
                    forwardable_proposal,
                    response,
                } => {
                    let am_leader = self.scheme.me().is_some_and(|me| me == leader);
                    current = Current {
                        view: new_current,
                        leader: Some(leader),
                        timed_out: false,
                    };
                    finalized = new_finalized;
                    work.entry(current.view)
                        .or_insert_with(|| self.new_round())
                        .set_leader(leader);

                    // If the leader nullified this view or has not been active
                    // recently, tell the voter to reduce the leader timeout to now
                    #[allow(clippy::collapsible_else_if)]
                    let timeout_reason = if Self::leader_nullified(&current, &work) {
                        // Leader already buffered a nullify for this now-current view
                        // (allowed because we accept votes up to `current+1`)
                        Some(TimeoutReason::LeaderNullify)
                    } else {
                        if am_leader {
                            // If we are the leader, we should not timeout
                            None
                        } else {
                            // If we are not the leader and the leader isn't active, we should timeout
                            (!self.is_active(&work, current.view, leader))
                                .then_some(TimeoutReason::Inactivity)
                        }
                    };
                    if timeout_reason.is_some() {
                        current.timed_out = true;
                    }
                    response.send_lossy(timeout_reason);

                    // Forward the proposal, if enabled and we have something to forward
                    if let Some((proposal, round)) = forwardable_proposal
                        .filter(|_| self.forwarding.is_enabled())
                        .and_then(|proposal| {
                            work.get(&proposal.view()).map(|round| (proposal, round))
                        })
                    {
                        let participants = self.forward_targets(round, &proposal, leader);
                        self.forward_proposal(proposal, participants).await;
                    }

                    // Setting leader may enable batch verification
                    updated_view = current.view;
                }
                Message::Constructed(message) => {
                    // If the view isn't interesting, we can skip
                    let view = message.view();
                    if !interesting(self.activity_timeout, finalized, current.view, view, false) {
                        continue;
                    }

                    // Add the message to the verifier
                    work.entry(view)
                        .or_insert_with(|| self.new_round())
                        .add_constructed(message)
                        .await;
                    self.added.inc();
                    updated_view = view;
                }
            },
            // Handle certificates from the network
            Ok((sender, message)) = certificate_receiver.recv() else break => {
                // If there is a decoding error, block
                let Ok(message) = message else {
                    commonware_p2p::block!(self.blocker, sender, "decoding error");
                    continue;
                };

                // Update metrics
                let label = match &message {
                    Certificate::Notarization(_) => Inbound::notarization(&sender),
                    Certificate::Nullification(_) => Inbound::nullification(&sender),
                    Certificate::Finalization(_) => Inbound::finalization(&sender),
                };
                self.inbound_messages.get_or_create(&label).inc();

                // If the epoch is not the current epoch, block
                if message.epoch() != self.epoch {
                    commonware_p2p::block!(self.blocker, sender, "epoch mismatch");
                    continue;
                }

                // Allow future certificates (they advance our view)
                let view = message.view();
                if !interesting(
                    self.activity_timeout,
                    finalized,
                    current.view,
                    view,
                    true, // allow future
                ) {
                    continue;
                }
                self.record_activity(&sender, view);

                match message {
                    Certificate::Notarization(notarization) => {
                        // Skip if we already have a notarization for this view
                        if work.get(&view).is_some_and(|r| r.has_notarization()) {
                            trace!(%view, "skipping duplicate notarization");
                            continue;
                        }

                        // Verify the certificate
                        if !notarization.verify(&mut self.context, &self.scheme, &self.strategy) {
                            commonware_p2p::block!(self.blocker, sender, %view, "invalid notarization");
                            continue;
                        }

                        // Store and forward to voter
                        work.entry(view)
                            .or_insert_with(|| self.new_round())
                            .set_notarization(notarization.clone());
                        voter
                            .recovered(Certificate::Notarization(notarization))
                            .await;
                    }
                    Certificate::Nullification(nullification) => {
                        // Skip if we already have a nullification for this view
                        if work.get(&view).is_some_and(|r| r.has_nullification()) {
                            trace!(%view, "skipping duplicate nullification");
                            continue;
                        }

                        // Verify the certificate
                        if !nullification.verify::<_, D>(
                            &mut self.context,
                            &self.scheme,
                            &self.strategy,
                        ) {
                            commonware_p2p::block!(self.blocker, sender, %view, "invalid nullification");
                            continue;
                        }

                        // Store and forward to voter
                        work.entry(view)
                            .or_insert_with(|| self.new_round())
                            .set_nullification(nullification.clone());
                        voter
                            .recovered(Certificate::Nullification(nullification))
                            .await;
                    }
                    Certificate::Finalization(finalization) => {
                        // Skip if we already have a finalization for this view
                        if work.get(&view).is_some_and(|r| r.has_finalization()) {
                            trace!(%view, "skipping duplicate finalization");
                            continue;
                        }

                        // Verify the certificate
                        if !finalization.verify(&mut self.context, &self.scheme, &self.strategy) {
                            commonware_p2p::block!(self.blocker, sender, %view, "invalid finalization");
                            continue;
                        }

                        // Store and forward to voter
                        work.entry(view)
                            .or_insert_with(|| self.new_round())
                            .set_finalization(finalization.clone());
                        voter
                            .recovered(Certificate::Finalization(finalization))
                            .await;
                    }
                }

                // Certificates are already forwarded to voter, no need for construction
                continue;
            },
            // Handle votes from the network
            Ok((sender, message)) = vote_receiver.recv() else break => {
                // If there is a decoding error, block
                let Ok(message) = message else {
                    commonware_p2p::block!(self.blocker, sender, "decoding error");
                    continue;
                };

                // Update metrics
                let label = match &message {
                    Vote::Notarize(_) => Inbound::notarize(&sender),
                    Vote::Nullify(_) => Inbound::nullify(&sender),
                    Vote::Finalize(_) => Inbound::finalize(&sender),
                };
                self.inbound_messages.get_or_create(&label).inc();

                // If the epoch is not the current epoch, block
                if message.epoch() != self.epoch {
                    commonware_p2p::block!(self.blocker, sender, "epoch mismatch");
                    continue;
                }

                // If the view isn't interesting, we can skip
                let view = message.view();
                if !interesting(self.activity_timeout, finalized, current.view, view, false) {
                    continue;
                }
                self.record_activity(&sender, view);

                // Add the vote to the verifier
                let peer = Peer::new(&sender);
                if work
                    .entry(view)
                    .or_insert_with(|| self.new_round())
                    .add_network(sender, message)
                    .await
                {
                    self.added.inc();

                    // Update per-peer latest vote metric (only if higher than current)
                    let _ = self
                        .latest_vote
                        .get_or_create(&peer)
                        .try_set_max(view.get());

                    // If the current leader explicitly nullifies the current view, signal
                    // the voter so it can fast-path timeout without waiting for its local
                    // timer. We check after adding because duplicate votes are rejected.
                    if Self::leader_nullified(&current, &work) {
                        current.timed_out = true;
                        voter
                            .timeout(current.view, TimeoutReason::LeaderNullify)
                            .await;
                    }
                }
                updated_view = view;
            },
            on_end => {
                assert!(
                    updated_view != View::zero(),
                    "updated view must be greater than zero"
                );

                // Forward leader's proposal to voter (if we're not the leader and haven't already)
                if let Some(round) = work.get_mut(&current.view) {
                    if let Some(me) = self.scheme.me() {
                        if let Some(proposal) = round.forward_proposal(me) {
                            voter.proposal(proposal).await;
                        }
                    }
                }

                // Skip verification and construction for views at or below finalized.
                //
                // We still use interesting() for filtering votes because we want to
                // notify the reporter of all votes within the activity timeout (even
                // if we don't need them in the voter).
                if updated_view <= finalized {
                    continue;
                }

                // Process the updated view (if any)
                let Some(round) = work.get_mut(&updated_view) else {
                    continue;
                };

                // Batch verify votes if ready
                let mut timer = self.verify_latency.timer();
                let verified = if round.ready_notarizes() {
                    Some(round.verify_notarizes(&mut self.context, &self.strategy))
                } else if round.ready_nullifies() {
                    Some(round.verify_nullifies(&mut self.context, &self.strategy))
                } else if round.ready_finalizes() {
                    Some(round.verify_finalizes(&mut self.context, &self.strategy))
                } else {
                    None
                };

                // Process batch verification results
                if let Some((voters, failed)) = verified {
                    timer.observe();

                    // Process verified votes
                    let batch = voters.len() + failed.len();
                    trace!(view = %updated_view, batch, "batch verified votes");
                    self.verified.inc_by(batch as u64);
                    self.batch_size.observe(batch as f64);

                    // Block invalid signers
                    for invalid in failed {
                        if let Some(signer) = self.participants.key(invalid) {
                            commonware_p2p::block!(
                                self.blocker,
                                signer.clone(),
                                "invalid signature"
                            );
                        }
                    }

                    // Store verified votes for certificate construction
                    for valid in voters {
                        round.add_verified(valid);
                    }
                } else {
                    timer.cancel();
                    trace!(
                        current = %current.view,
                        %finalized,
                        "no verifier ready"
                    );
                }

                // Try to construct and forward certificates
                if let Some(notarization) = self
                    .recover_latency
                    .time_some(|| round.try_construct_notarization(&self.scheme, &self.strategy))
                {
                    debug!(view = %updated_view, "constructed notarization, forwarding to voter");

                    // Forward notarization to voter
                    voter
                        .recovered(Certificate::Notarization(notarization))
                        .await;
                }
                if let Some(nullification) = self
                    .recover_latency
                    .time_some(|| round.try_construct_nullification(&self.scheme, &self.strategy))
                {
                    debug!(view = %updated_view, "constructed nullification, forwarding to voter");
                    voter
                        .recovered(Certificate::Nullification(nullification))
                        .await;
                }
                if let Some(finalization) = self
                    .recover_latency
                    .time_some(|| round.try_construct_finalization(&self.scheme, &self.strategy))
                {
                    debug!(view = %updated_view, "constructed finalization, forwarding to voter");
                    voter
                        .recovered(Certificate::Finalization(finalization))
                        .await;
                }

                // Drop any rounds that are no longer interesting
                while work.first_key_value().is_some_and(|(&view, _)| {
                    !interesting(self.activity_timeout, finalized, current.view, view, false)
                }) {
                    work.pop_first();
                }
            },
        }
    }
}
