use std::{
    cmp::Ordering,
    collections::BTreeMap,
    fmt::{self, Debug, Display, Formatter},
    hash::{Hash, Hasher},
    ops::{Deref, Range, RangeInclusive},
    sync::Arc,
};

use super::block::{BlockRef, BlockTimestampMs, Round, TransactionIndex};
use super::{
    block::{BlockAPI, Slot, VerifiedBlock},
    leader_scoring::ReputationScores,
};
use crate::committee::AuthorityIndex;
use crate::crypto::{DIGEST_LENGTH, DefaultHash};
use crate::storage::consensus::Store;
use bytes::Bytes;
use enum_dispatch::enum_dispatch;
use fastcrypto::hash::{Digest, HashFunction as _};
use itertools::Itertools as _;
use serde::{Deserialize, Serialize};

/// Index of a commit among all consensus commits.
pub type CommitIndex = u32;

pub const GENESIS_COMMIT_INDEX: CommitIndex = 0;

/// Default wave length for all committers. A longer wave length increases the
/// chance of committing the leader under asynchrony at the cost of latency in
/// the common case.
// TODO: merge DEFAULT_WAVE_LENGTH and MINIMUM_WAVE_LENGTH into a single constant,
// because we are unlikely to change them via config in the forseeable future.
pub const DEFAULT_WAVE_LENGTH: Round = MINIMUM_WAVE_LENGTH;

/// We need at least one leader round, one voting round, and one decision round.
pub const MINIMUM_WAVE_LENGTH: Round = 3;

/// The consensus protocol operates in 'waves'. Each wave is composed of a leader
/// round, at least one voting round, and one decision round.
pub type WaveNumber = u32;

/// [`Commit`] summarizes [`CommittedSubDag`] for storage and network communications.
///
/// Validators should be able to reconstruct a sequence of CommittedSubDag from the
/// corresponding Commit and blocks referenced in the Commit.
/// A field must meet these requirements to be added to Commit:
/// - helps with recovery locally and for peers catching up.
/// - cannot be derived from a sequence of Commits and other persisted values.
///
/// For example, transactions in blocks should not be included in Commit, because they can be
/// retrieved from blocks specified in Commit. Last committed round per authority also should not
/// be included, because it can be derived from the latest value in storage and the additional
/// sequence of Commits.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[enum_dispatch(CommitAPI)]
pub enum Commit {
    V1(CommitV1),
}

impl Commit {
    /// Create a new commit.
    pub fn new(
        index: CommitIndex,
        previous_digest: CommitDigest,
        timestamp_ms: BlockTimestampMs,
        leader: BlockRef,
        blocks: Vec<BlockRef>,
    ) -> Self {
        Commit::V1(CommitV1 { index, previous_digest, timestamp_ms, leader, blocks })
    }

    pub fn serialize(&self) -> Result<Bytes, bcs::Error> {
        let bytes = bcs::to_bytes(self)?;
        Ok(bytes.into())
    }
}

/// Accessors to Commit info.
#[enum_dispatch]
pub trait CommitAPI {
    fn round(&self) -> Round;
    fn index(&self) -> CommitIndex;
    fn previous_digest(&self) -> CommitDigest;
    fn timestamp_ms(&self) -> BlockTimestampMs;
    fn leader(&self) -> BlockRef;
    fn blocks(&self) -> &[BlockRef];
}

/// Specifies one consensus commit.
/// It is stored on disk, so it does not contain blocks which are stored individually.
#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq)]
pub struct CommitV1 {
    /// Index of the commit.
    /// First commit after genesis has an index of 1, then every next commit has an index incremented by 1.
    index: CommitIndex,
    /// Digest of the previous commit.
    /// Set to CommitDigest::MIN for the first commit after genesis.
    previous_digest: CommitDigest,
    /// Timestamp of the commit, max of the timestamp of the leader block and previous Commit timestamp.
    timestamp_ms: BlockTimestampMs,
    /// A reference to the commit leader.
    leader: BlockRef,
    /// Refs to committed blocks, in the commit order.
    blocks: Vec<BlockRef>,
}

impl CommitAPI for CommitV1 {
    fn round(&self) -> Round {
        self.leader.round
    }

    fn index(&self) -> CommitIndex {
        self.index
    }

    fn previous_digest(&self) -> CommitDigest {
        self.previous_digest
    }

    fn timestamp_ms(&self) -> BlockTimestampMs {
        self.timestamp_ms
    }

    fn leader(&self) -> BlockRef {
        self.leader
    }

    fn blocks(&self) -> &[BlockRef] {
        &self.blocks
    }
}

/// A commit is trusted when it is produced locally or certified by a quorum of authorities.
/// Blocks referenced by TrustedCommit are assumed to be valid.
/// Only trusted Commit can be sent to execution.
///
/// Note: clone() is relatively cheap with the underlying data refcounted.
#[derive(Clone, Debug, PartialEq)]
pub struct TrustedCommit {
    inner: Arc<Commit>,

    // Cached digest and serialized value, to avoid re-computing these values.
    digest: CommitDigest,
    serialized: Bytes,
}

impl TrustedCommit {
    pub fn new_trusted(commit: Commit, serialized: Bytes) -> Self {
        let digest = Self::compute_digest(&serialized);
        Self { inner: Arc::new(commit), digest, serialized }
    }

    pub fn new_for_test(
        index: CommitIndex,
        previous_digest: CommitDigest,
        timestamp_ms: BlockTimestampMs,
        leader: BlockRef,
        blocks: Vec<BlockRef>,
    ) -> Self {
        let commit = Commit::new(index, previous_digest, timestamp_ms, leader, blocks);
        let serialized = commit.serialize().unwrap();
        Self::new_trusted(commit, serialized)
    }

    pub fn reference(&self) -> CommitRef {
        CommitRef { index: self.index(), digest: self.digest() }
    }

    pub fn digest(&self) -> CommitDigest {
        self.digest
    }

    pub fn serialized(&self) -> &Bytes {
        &self.serialized
    }

    pub fn compute_digest(serialized: &[u8]) -> CommitDigest {
        let mut hasher = DefaultHash::new();
        hasher.update(serialized);
        CommitDigest(hasher.finalize().into())
    }
}

/// Allow easy access on the underlying Commit.
impl Deref for TrustedCommit {
    type Target = Commit;

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

/// `CertifiedCommits` keeps the synchronized certified commits along with the corresponding votes received from the peer that provided these commits.
/// The `votes` contain the blocks as those provided by the peer, and certify the tip of the synced commits.
#[derive(Clone, Debug)]
pub struct CertifiedCommits {
    commits: Vec<CertifiedCommit>,
    votes: Vec<VerifiedBlock>,
}

impl CertifiedCommits {
    pub fn new(commits: Vec<CertifiedCommit>, votes: Vec<VerifiedBlock>) -> Self {
        Self { commits, votes }
    }

    pub fn commits(&self) -> &[CertifiedCommit] {
        &self.commits
    }

    pub fn votes(&self) -> &[VerifiedBlock] {
        &self.votes
    }
}

/// A commit that has been synced and certified by a quorum of authorities.
#[derive(Clone, Debug)]
pub struct CertifiedCommit {
    commit: Arc<TrustedCommit>,
    blocks: Vec<VerifiedBlock>,
}

impl CertifiedCommit {
    pub fn new_certified(commit: TrustedCommit, blocks: Vec<VerifiedBlock>) -> Self {
        Self { commit: Arc::new(commit), blocks }
    }

    pub fn blocks(&self) -> &[VerifiedBlock] {
        &self.blocks
    }
}

impl Deref for CertifiedCommit {
    type Target = TrustedCommit;

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

/// Digest of a consensus commit.
#[derive(Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
pub struct CommitDigest([u8; DIGEST_LENGTH]);

impl CommitDigest {
    /// Lexicographic min & max digest.
    pub const MIN: Self = Self([u8::MIN; DIGEST_LENGTH]);
    pub const MAX: Self = Self([u8::MAX; DIGEST_LENGTH]);

    pub fn into_inner(self) -> [u8; DIGEST_LENGTH] {
        self.0
    }
}

impl Hash for CommitDigest {
    fn hash<H: Hasher>(&self, state: &mut H) {
        state.write(&self.0[..8]);
    }
}

impl From<CommitDigest> for Digest<{ DIGEST_LENGTH }> {
    fn from(hd: CommitDigest) -> Self {
        Digest::new(hd.0)
    }
}

impl fmt::Display for CommitDigest {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        write!(
            f,
            "{}",
            base64::Engine::encode(&base64::engine::general_purpose::STANDARD, self.0)
                .get(0..4)
                .ok_or(fmt::Error)?
        )
    }
}

impl fmt::Debug for CommitDigest {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        write!(f, "{}", base64::Engine::encode(&base64::engine::general_purpose::STANDARD, self.0))
    }
}

/// Uniquely identifies a commit with its index and digest.
#[derive(Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq, PartialOrd, Ord)]
pub struct CommitRef {
    pub index: CommitIndex,
    pub digest: CommitDigest,
}

impl CommitRef {
    pub fn new(index: CommitIndex, digest: CommitDigest) -> Self {
        Self { index, digest }
    }
}

impl fmt::Display for CommitRef {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        write!(f, "C{}({})", self.index, self.digest)
    }
}

impl fmt::Debug for CommitRef {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        write!(f, "C{}({:?})", self.index, self.digest)
    }
}

// Represents a vote on a Commit.
pub type CommitVote = CommitRef;

/// The output of consensus to execution is an ordered list of [`CommittedSubDag`].
/// Each CommittedSubDag contains the information needed to execution transactions in
/// the consensus commit.
///
/// The application processing CommittedSubDag can arbitrarily sort the blocks within
/// each sub-dag (but using a deterministic algorithm).
#[derive(Clone, PartialEq)]
pub struct CommittedSubDag {
    /// Set by Linearizer.
    ///
    /// A reference to the leader of the sub-dag
    pub leader: BlockRef,
    /// All the committed blocks that are part of this sub-dag
    pub blocks: Vec<VerifiedBlock>,
    /// The timestamp of the commit, obtained from the timestamp of the leader block.
    pub timestamp_ms: BlockTimestampMs,
    /// The reference of the commit.
    /// First commit after genesis has a index of 1, then every next commit has a
    /// index incremented by 1.
    pub commit_ref: CommitRef,

    /// Set by CommitObserver.
    ///
    /// Indicates whether the commit was decided locally based on the local DAG.
    ///
    /// If true, `CommitFinalizer` can then assume a quorum of certificates are available
    /// for each transaction in the commit if there is no reject vote, and proceed with
    /// optimistic finalization of transactions.
    ///
    /// If the commit was decided by `UniversalCommitter`, this must be true.
    /// If the commit was received from a peer via `CommitSyncer`, this must be false.
    /// There may not be enough blocks in local DAG to decide on the commit.
    ///
    /// For safety, a previously locally decided commit may be recovered after restarting as
    /// non-local, if its finalization state was not persisted.
    pub decided_with_local_blocks: bool,
    /// Whether rejected transactions in this commit have been recovered from storage.
    pub recovered_rejected_transactions: bool,
    /// Optional scores that are provided as part of the consensus output to Sui
    /// that can then be used by Sui for future submission to consensus.
    pub reputation_scores_desc: Vec<(AuthorityIndex, u64)>,

    /// Set by CommitFinalizer.
    ///
    /// Indices of rejected transactions in each block.
    pub rejected_transactions_by_block: BTreeMap<BlockRef, Vec<TransactionIndex>>,
}

impl CommittedSubDag {
    /// Creates a new committed sub dag.
    pub fn new(
        leader: BlockRef,
        blocks: Vec<VerifiedBlock>,
        timestamp_ms: BlockTimestampMs,
        commit_ref: CommitRef,
    ) -> Self {
        Self {
            leader,
            blocks,
            timestamp_ms,
            commit_ref,
            decided_with_local_blocks: true,
            recovered_rejected_transactions: false,
            reputation_scores_desc: vec![],
            rejected_transactions_by_block: BTreeMap::new(),
        }
    }
}

// Sort the blocks of the sub-dag blocks by round number then authority index. Any
// deterministic & stable algorithm works.
pub fn sort_sub_dag_blocks(blocks: &mut [VerifiedBlock]) {
    blocks.sort_by(|a, b| a.round().cmp(&b.round()).then_with(|| a.author().cmp(&b.author())))
}

impl Display for CommittedSubDag {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}@{} [{}])",
            self.commit_ref,
            self.leader,
            self.blocks.iter().map(|b| b.reference().to_string()).join(", ")
        )
    }
}

impl fmt::Debug for CommittedSubDag {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}@{} [{}])",
            self.commit_ref,
            self.leader,
            self.blocks.iter().map(|b| b.reference().to_string()).join(", ")
        )?;
        write!(
            f,
            ";{}ms;rs{:?};{};{};[{}]",
            self.timestamp_ms,
            self.reputation_scores_desc,
            self.decided_with_local_blocks,
            self.recovered_rejected_transactions,
            self.rejected_transactions_by_block
                .iter()
                .map(|(block_ref, transactions)| {
                    format!("{}: {}, ", block_ref, transactions.len())
                })
                .collect::<Vec<_>>()
                .join(", "),
        )
    }
}

// Recovers the full CommittedSubDag from block store, based on Commit.
pub fn load_committed_subdag_from_store(
    store: &dyn Store,
    commit: TrustedCommit,
    reputation_scores_desc: Vec<(AuthorityIndex, u64)>,
) -> CommittedSubDag {
    let mut leader_block_idx = None;
    let commit_blocks = store
        .read_blocks(commit.blocks())
        .expect("We should have the block referenced in the commit data");
    let blocks = commit_blocks
        .into_iter()
        .enumerate()
        .map(|(idx, commit_block_opt)| {
            let commit_block =
                commit_block_opt.expect("We should have the block referenced in the commit data");
            if commit_block.reference() == commit.leader() {
                leader_block_idx = Some(idx);
            }
            commit_block
        })
        .collect::<Vec<_>>();
    let leader_block_idx = leader_block_idx.expect("Leader block must be in the sub-dag");
    let leader_block_ref = blocks[leader_block_idx].reference();

    let mut subdag =
        CommittedSubDag::new(leader_block_ref, blocks, commit.timestamp_ms(), commit.reference());

    subdag.reputation_scores_desc = reputation_scores_desc;

    let reject_votes = store.read_rejected_transactions(commit.reference()).unwrap();
    if let Some(reject_votes) = reject_votes {
        subdag.decided_with_local_blocks = true;
        subdag.recovered_rejected_transactions = true;
        subdag.rejected_transactions_by_block = reject_votes;
    } else {
        subdag.decided_with_local_blocks = false;
        subdag.recovered_rejected_transactions = false;
    }

    subdag
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum Decision {
    Direct,
    Indirect,
    Certified, // This is a commit certified leader so no commit decision was made locally.
}

/// The status of a leader slot from the direct and indirect commit rules.
#[derive(Debug, Clone, PartialEq)]
pub enum LeaderStatus {
    Commit(VerifiedBlock),
    Skip(Slot),
    Undecided(Slot),
}

impl LeaderStatus {
    pub fn round(&self) -> Round {
        match self {
            Self::Commit(block) => block.round(),
            Self::Skip(leader) => leader.round,
            Self::Undecided(leader) => leader.round,
        }
    }

    pub fn is_decided(&self) -> bool {
        match self {
            Self::Commit(_) => true,
            Self::Skip(_) => true,
            Self::Undecided(_) => false,
        }
    }

    pub fn into_decided_leader(self, direct: bool) -> Option<DecidedLeader> {
        match self {
            Self::Commit(block) => Some(DecidedLeader::Commit(block, direct)),
            Self::Skip(slot) => Some(DecidedLeader::Skip(slot)),
            Self::Undecided(..) => None,
        }
    }
}

impl Display for LeaderStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Commit(block) => write!(f, "Commit({})", block.reference()),
            Self::Skip(slot) => write!(f, "Skip({slot})"),
            Self::Undecided(slot) => write!(f, "Undecided({slot})"),
        }
    }
}

/// Decision of each leader slot.
#[derive(Debug, Clone, PartialEq)]
pub enum DecidedLeader {
    /// The committed leader block and whether it is a direct commit.
    /// It is incorrect to trigger the direct commit optimization when the commit is not.
    /// So when it is unknown if the commit is direct, the boolean flag should be false.
    Commit(VerifiedBlock, bool),
    /// The skipped leader slot where no block is committed.
    Skip(Slot),
}

impl DecidedLeader {
    // Slot where the leader is decided.
    pub fn slot(&self) -> Slot {
        match self {
            Self::Commit(block, _direct) => block.reference().into(),
            Self::Skip(slot) => *slot,
        }
    }

    // Converts to committed block if the decision is to commit. Returns None otherwise.
    pub fn into_committed_block(self) -> Option<VerifiedBlock> {
        match self {
            Self::Commit(block, _direct) => Some(block),
            Self::Skip(_) => None,
        }
    }

    #[cfg(test)]
    pub fn round(&self) -> Round {
        match self {
            Self::Commit(block, _direct) => block.round(),
            Self::Skip(leader) => leader.round,
        }
    }

    #[cfg(test)]
    pub fn authority(&self) -> AuthorityIndex {
        match self {
            Self::Commit(block, _direct) => block.author(),
            Self::Skip(leader) => leader.authority,
        }
    }
}

impl Display for DecidedLeader {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Commit(block, _direct) => write!(f, "Commit({})", block.reference()),
            Self::Skip(slot) => write!(f, "Skip({slot})"),
        }
    }
}

/// Per-commit properties that can be regenerated from past values, and do not need to be part of
/// the Commit struct.
/// Only the latest version is needed for recovery, but more versions are stored for debugging,
/// and potentially restoring from an earlier state.
// TODO: version this struct.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct CommitInfo {
    pub committed_rounds: Vec<Round>,
    pub reputation_scores: ReputationScores,
}

/// `CommitRange` stores a range of `CommitIndex`. The range contains the start (inclusive)
/// and end (inclusive) commit indices and can be ordered for use as the key of a table.
///
/// NOTE: using `Range<CommitIndex>` for internal representation for backward compatibility.
/// The external semantics of `CommitRange` is closer to `RangeInclusive<CommitIndex>`.
#[derive(Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct CommitRange(Range<CommitIndex>);

impl CommitRange {
    pub fn new(range: RangeInclusive<CommitIndex>) -> Self {
        // When end is CommitIndex::MAX, the range can be considered as unbounded
        // so it is ok to saturate at the end.
        Self(*range.start()..(*range.end()).saturating_add(1))
    }

    // Inclusive
    pub fn start(&self) -> CommitIndex {
        self.0.start
    }

    // Inclusive
    pub fn end(&self) -> CommitIndex {
        self.0.end.saturating_sub(1)
    }

    pub fn extend_to(&mut self, other: CommitIndex) {
        let new_end = other.saturating_add(1);
        assert!(self.0.end <= new_end);
        self.0 = self.0.start..new_end;
    }

    pub fn size(&self) -> usize {
        self.0.end.checked_sub(self.0.start).expect("Range should never have end < start") as usize
    }

    /// Check whether the two ranges have the same size.
    pub fn is_equal_size(&self, other: &Self) -> bool {
        self.size() == other.size()
    }

    /// Check if the provided range is sequentially after this range.
    pub fn is_next_range(&self, other: &Self) -> bool {
        self.0.end == other.0.start
    }
}

impl Ord for CommitRange {
    fn cmp(&self, other: &Self) -> Ordering {
        self.start().cmp(&other.start()).then_with(|| self.end().cmp(&other.end()))
    }
}

impl PartialOrd for CommitRange {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl From<RangeInclusive<CommitIndex>> for CommitRange {
    fn from(range: RangeInclusive<CommitIndex>) -> Self {
        Self::new(range)
    }
}

/// Display CommitRange as an inclusive range.
impl Debug for CommitRange {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "CommitRange({}..={})", self.start(), self.end())
    }
}
