use std::collections::BTreeMap;
use std::hash::Hash;
use std::sync::{Arc, OnceLock};
use std::time::Duration;

use ahash::HashMapExt;
use anyhow::{Result, anyhow, bail};
use parking_lot::lock_api::RwLockReadGuard;
use parking_lot::{MappedRwLockReadGuard, Mutex, RwLock};
use tl_proto::TlWrite;
use tycho_block_util::queue::QueuePartitionIdx;
use tycho_block_util::state::{RefMcStateHandle, ShardStateStuff};
use tycho_core::global_config::MempoolGlobalConfig;
use tycho_executor::{AccountMeta, PublicLibraryChange, TransactionMeta};
use tycho_types::boc;
use tycho_types::cell::{Cell, CellFamily, HashBytes, Lazy, UsageTree, UsageTreeMode};
use tycho_types::dict::{self, Dict};
use tycho_types::merkle::MerkleBuildResult;
use tycho_types::models::{
    AccountBlocks, AccountState, BlockId, BlockIdShort, BlockInfo, BlockLimits, BlockParamLimits,
    BlockRef, BlockchainConfig, CollationConfig, CurrencyCollection, ExtInMsgInfo, HashUpdate,
    ImportFees, InMsg, InMsgDescr, IntMsgInfo, LibDescr, MsgInfo, MsgsExecutionParams,
    OptionalAccount, OutMsg, OutMsgDescr, OwnedMessage, PrevBlockRef, ShardAccount, ShardAccounts,
    ShardDescription, ShardFeeCreated, ShardFees, ShardIdent, ShardIdentFull, ShardStateUnsplit,
    SpecialFlags, StateInit, StdAddr, Transaction, ValueFlow,
};
use tycho_types::num::Tokens;
use tycho_util::{FastHashMap, FastHashSet};

use super::do_collate::work_units::PrepareMsgGroupsWu;
use super::error::CollationAbortReason;
use super::messages_reader::MessagesReaderMetrics;
use crate::collator::do_collate::work_units::{DoCollateWu, ExecuteWu, FinalizeWu};
use crate::collator::messages_reader::MetricsTimer;
use crate::collator::messages_reader::state::ReaderState;
use crate::internal_queue::types::diff::QueueDiffWithMessages;
use crate::internal_queue::types::message::EnqueuedMessage;
use crate::tracing_targets;
use crate::types::processed_upto::{BlockSeqno, ProcessedUptoInfoStuff};
use crate::types::{
    BlockCandidate, BlockCollationResult, McData, ProcessedToByPartitions, TopShardBlockInfo,
};

pub(super) struct WorkingState {
    pub next_block_id_short: BlockIdShort,
    pub mc_data: Arc<McData>,
    pub collation_config: Arc<CollationConfig>,
    pub wu_used_from_last_anchor: u64,
    pub resume_collation_elapsed: Duration,
    pub prev_shard_data: Option<PrevData>,
    pub usage_tree: Option<UsageTree>,
    pub has_unprocessed_messages: Option<bool>,
    pub reader_state: ReaderState,
}

impl WorkingState {
    pub fn prev_shard_data_ref(&self) -> &PrevData {
        self.prev_shard_data.as_ref().unwrap()
    }
}

pub(super) struct PrevData {
    observable_states: Vec<ShardStateStuff>,
    observable_accounts: ShardAccounts,

    blocks_ids: Vec<BlockId>,

    pure_states: Vec<ShardStateStuff>,

    gen_chain_time: u64,
    gen_lt: u64,
    total_validator_fees: CurrencyCollection,
    wu_used_from_last_anchor: u64,

    processed_upto: ProcessedUptoInfoStuff,

    prev_queue_diff_hashes: Vec<HashBytes>,
}

impl PrevData {
    pub fn build(
        prev_states: Vec<ShardStateStuff>,
        prev_queue_diff_hashes: Vec<HashBytes>,
    ) -> Result<(Self, UsageTree)> {
        // TODO: consider split/merge logic
        //  Collator::prepare_data()
        //  Collator::unpack_last_state()

        let prev_blocks_ids: Vec<_> = prev_states.iter().map(|s| *s.block_id()).collect();
        let pure_prev_states = prev_states;

        let (usage_tree, observable_state) =
            pure_prev_states[0].track_usage(UsageTreeMode::OnLoad)?;
        let observable_states = vec![observable_state];

        let gen_chain_time = observable_states[0].get_gen_chain_time();
        let gen_lt = observable_states[0].state().gen_lt;
        let observable_accounts = observable_states[0].state().load_accounts()?;
        let total_validator_fees = observable_states[0].state().total_validator_fees.clone();
        let wu_used_from_last_anchor = observable_states[0].state().overload_history;

        let processed_upto_info = pure_prev_states[0].state().processed_upto.load()?;

        let prev_data = Self {
            observable_states,
            observable_accounts,

            blocks_ids: prev_blocks_ids,

            pure_states: pure_prev_states,

            gen_chain_time,
            gen_lt,
            total_validator_fees,
            wu_used_from_last_anchor,

            processed_upto: processed_upto_info.try_into()?,

            prev_queue_diff_hashes,
        };

        Ok((prev_data, usage_tree))
    }

    pub fn pure_states(&self) -> &Vec<ShardStateStuff> {
        &self.pure_states
    }

    pub fn observable_states(&self) -> &Vec<ShardStateStuff> {
        &self.observable_states
    }

    pub fn observable_accounts(&self) -> &ShardAccounts {
        &self.observable_accounts
    }

    pub fn blocks_ids(&self) -> &Vec<BlockId> {
        &self.blocks_ids
    }

    pub fn get_blocks_ref(&self) -> Result<PrevBlockRef> {
        if self.pure_states.is_empty() || self.pure_states.len() > 2 {
            bail!(
                "There should be 1 or 2 prev states. Actual count is {}",
                self.pure_states.len()
            )
        }

        let mut block_refs = vec![];
        for state in self.pure_states.iter() {
            block_refs.push(BlockRef {
                end_lt: state.state().gen_lt,
                seqno: state.block_id().seqno,
                root_hash: state.block_id().root_hash,
                file_hash: state.block_id().file_hash,
            });
        }

        let prev_ref = if block_refs.len() == 2 {
            PrevBlockRef::AfterMerge {
                left: block_refs.remove(0),
                right: block_refs.remove(0),
            }
        } else {
            PrevBlockRef::Single(block_refs.remove(0))
        };

        Ok(prev_ref)
    }

    pub fn ref_mc_state_handle(&self) -> &RefMcStateHandle {
        self.observable_states[0].ref_mc_state_handle()
    }

    pub fn pure_state_root(&self) -> &Cell {
        self.pure_states[0].root_cell()
    }

    pub fn gen_chain_time(&self) -> u64 {
        self.gen_chain_time
    }

    pub fn gen_lt(&self) -> u64 {
        self.gen_lt
    }

    pub fn wu_used_from_last_anchor(&self) -> u64 {
        self.wu_used_from_last_anchor
    }

    pub fn total_validator_fees(&self) -> &CurrencyCollection {
        &self.total_validator_fees
    }

    pub fn processed_upto(&self) -> &ProcessedUptoInfoStuff {
        &self.processed_upto
    }

    pub fn prev_queue_diff_hashes(&self) -> &Vec<HashBytes> {
        &self.prev_queue_diff_hashes
    }
}

pub(super) enum NextCollationFlowStep {
    Continue,
    Cancel(CancelledContext),
    Abort(AbortedContext),
}

#[derive(Debug)]
pub struct CancelledContext {
    pub prev_mc_block_id: BlockId,
    pub next_block_id_short: BlockIdShort,
}

#[derive(Debug)]
pub struct AbortedContext {
    pub prev_mc_block_id: BlockId,
    pub next_block_id_short: BlockIdShort,
    pub reason: CollationAbortReason,
}

pub enum CollatorResult {
    Skipped {
        prev_mc_block_id: BlockId,
        next_block_id_short: BlockIdShort,
        anchor_chain_time: u64,
        force_mc_block: ForceMasterCollation,
        collation_config: Arc<CollationConfig>,
    },
    Cancelled(CancelledContext),
    Aborted(AbortedContext),
    BlockCandidate {
        collation_result: BlockCollationResult,
    },
}

pub struct DebugCollatorResult<'a>(pub &'a CollatorResult);

impl std::fmt::Debug for DebugCollatorResult<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.0 {
            CollatorResult::Skipped {
                prev_mc_block_id,
                next_block_id_short,
                anchor_chain_time,
                force_mc_block,
                collation_config,
            } => f
                .debug_struct("Skipped")
                .field("prev_mc_block_id", prev_mc_block_id)
                .field("next_block_id_short", next_block_id_short)
                .field("anchor_chain_time", anchor_chain_time)
                .field("force_mc_block", force_mc_block)
                .field("collation_config", collation_config)
                .finish(),
            CollatorResult::Cancelled(cancel_ctx) => {
                f.debug_tuple("Cancelled").field(cancel_ctx).finish()
            }
            CollatorResult::Aborted(abort_ctx) => {
                f.debug_tuple("Aborted").field(abort_ctx).finish()
            }
            CollatorResult::BlockCandidate { .. } => {
                f.debug_struct("BlockCandidate").finish_non_exhaustive()
            }
        }
    }
}

impl std::fmt::Debug for CollatorResult {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        DebugCollatorResult(self).fmt(f)
    }
}

impl CollatorResult {
    pub fn skipped(
        prev_mc_block_id: BlockId,
        next_block_id_short: BlockIdShort,
        anchor_chain_time: u64,
        force_mc_block: ForceMasterCollation,
        collation_config: Arc<CollationConfig>,
    ) -> Self {
        Self::Skipped {
            prev_mc_block_id,
            next_block_id_short,
            anchor_chain_time,
            force_mc_block,
            collation_config,
        }
    }

    pub fn cancelled(prev_mc_block_id: BlockId, next_block_id_short: BlockIdShort) -> Self {
        Self::Cancelled(CancelledContext {
            prev_mc_block_id,
            next_block_id_short,
        })
    }

    pub fn aborted(
        prev_mc_block_id: BlockId,
        next_block_id_short: BlockIdShort,
        reason: CollationAbortReason,
    ) -> Self {
        Self::Aborted(AbortedContext {
            prev_mc_block_id,
            next_block_id_short,
            reason,
        })
    }
}

#[derive(Debug)]
pub(super) struct BlockCollationDataBuilder {
    pub block_id_short: BlockIdShort,
    pub gen_utime: u32,
    pub gen_utime_ms: u16,
    shards: Option<FastHashMap<ShardIdent, Box<ShardDescription>>>,
    pub shards_max_end_lt: u64,
    pub shard_fees: ShardFees,
    pub min_ref_mc_seqno: u32,
    pub rand_seed: HashBytes,
    pub created_by: HashBytes,
    pub top_shard_blocks: Vec<TopShardBlockInfo>,

    pub mc_shards_processed_to_by_partitions:
        FastHashMap<ShardIdent, (bool, ProcessedToByPartitions)>,

    /// Mempool config override for a new genesis
    pub mempool_config_override: Option<MempoolGlobalConfig>,
}

impl BlockCollationDataBuilder {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        block_id_short: BlockIdShort,
        rand_seed: HashBytes,
        min_ref_mc_seqno: u32,
        next_chain_time: u64,
        created_by: HashBytes,
        mempool_config_override: Option<MempoolGlobalConfig>,
    ) -> Self {
        let gen_utime = (next_chain_time / 1000) as u32;
        let gen_utime_ms = (next_chain_time % 1000) as u16;
        Self {
            block_id_short,
            gen_utime,
            gen_utime_ms,
            shards_max_end_lt: 0,
            shard_fees: Default::default(),
            min_ref_mc_seqno,
            rand_seed,
            created_by,
            shards: None,
            top_shard_blocks: vec![],
            mc_shards_processed_to_by_partitions: Default::default(),
            mempool_config_override,
        }
    }

    pub fn set_shards(&mut self, shards: FastHashMap<ShardIdent, Box<ShardDescription>>) {
        self.shards = Some(shards);
    }

    pub fn shards_mut(&mut self) -> Result<&mut FastHashMap<ShardIdent, Box<ShardDescription>>> {
        self.shards
            .as_mut()
            .ok_or_else(|| anyhow!("`shards` is not initialized yet"))
    }

    pub fn shards(&self) -> Result<&FastHashMap<ShardIdent, Box<ShardDescription>>> {
        self.shards
            .as_ref()
            .ok_or_else(|| anyhow!("`shards` is not initialized yet"))
    }

    pub fn update_shards_max_end_lt(&mut self, val: u64) {
        if val > self.shards_max_end_lt {
            self.shards_max_end_lt = val;
        }
    }

    pub fn store_shard_fees(
        &mut self,
        shard_id: ShardIdent,
        proof_funds: ShardFeeCreated,
    ) -> Result<()> {
        self.shard_fees.set(
            ShardIdentFull::from(shard_id),
            proof_funds.clone(),
            proof_funds,
        )?;
        Ok(())
    }

    pub fn build(self, start_lt: u64, block_limits: BlockLimits) -> BlockCollationData {
        let block_limit = BlockLimitStats::new(block_limits, start_lt);
        BlockCollationData {
            block_id_short: self.block_id_short,
            gen_utime: self.gen_utime,
            gen_utime_ms: self.gen_utime_ms,
            min_ref_mc_seqno: self.min_ref_mc_seqno,
            rand_seed: self.rand_seed,
            created_by: self.created_by,
            shards: self.shards,
            top_shard_blocks: self.top_shard_blocks,
            shard_fees: self.shard_fees,
            value_flow: Default::default(),
            block_limit,
            start_lt,
            next_lt: start_lt + 1,
            tx_count: 0,
            touched_accounts: Default::default(),
            shard_accounts_count: 0,
            updated_accounts_count: 0,
            added_accounts_count: 0,
            removed_accounts_count: 0,
            total_execute_msgs_time_mc: 0,
            execute_count_all: 0,
            execute_count_ext: 0,
            ext_msgs_error_count: 0,
            ext_msgs_skipped_count: 0,
            execute_count_int: 0,
            execute_count_new_int: 0,
            int_enqueue_count: 0,
            int_dequeue_count: 0,
            // read_ext_msgs_count: 0,
            // read_int_msgs_from_iterator_count: 0,
            new_msgs_created_count: 0,
            inserted_new_msgs_count: 0,
            // read_new_msgs_count: 0,
            in_msgs: Default::default(),
            out_msgs: Default::default(),
            mint_msg: None,
            recover_create_msg: None,
            mempool_config_override: self.mempool_config_override,
            consensus_config_changed: None,
            diff_tail_len: 0,
            mc_shards_processed_to_by_partitions: self.mc_shards_processed_to_by_partitions,
        }
    }
}

#[derive(Debug)]
pub(super) struct BlockCollationData {
    pub block_id_short: BlockIdShort,
    pub gen_utime: u32,
    pub gen_utime_ms: u16,

    pub tx_count: u64,
    pub touched_accounts: FastHashSet<HashBytes>,

    pub shard_accounts_count: u64,
    pub updated_accounts_count: u64,
    pub added_accounts_count: u64,
    pub removed_accounts_count: u64,

    pub block_limit: BlockLimitStats,

    pub total_execute_msgs_time_mc: u128,

    pub execute_count_all: u64,
    pub execute_count_ext: u64,
    pub execute_count_int: u64,
    pub execute_count_new_int: u64,

    pub ext_msgs_error_count: u64,
    pub ext_msgs_skipped_count: u64,

    pub int_enqueue_count: u64,
    pub int_dequeue_count: u64,

    // pub read_ext_msgs_count: u64,
    // pub read_int_msgs_from_iterator_count: u64,
    pub new_msgs_created_count: u64,
    pub inserted_new_msgs_count: u64,
    // pub read_new_msgs_count: u64,
    pub start_lt: u64,
    // Should be updated on each tx finalization from MessagesPreparer.max_lt
    // which is updating during tx execution
    pub next_lt: u64,

    pub in_msgs: BTreeMap<HashBytes, PreparedInMsg>,
    pub out_msgs: BTreeMap<HashBytes, PreparedOutMsg>,

    /// Ids of top blocks from shards that were included in the master block
    pub top_shard_blocks: Vec<TopShardBlockInfo>,

    /// Actual processed to info by each shard from referenced master block:
    /// * will contain copy of data from `McData` when collating shard blocks
    /// * will contain updated info by top shard blocks info when collating master blocks
    ///
    /// Stores a tuple for each shard:
    /// (referenced shard is updated in master block, processed to info by partitions)
    pub mc_shards_processed_to_by_partitions:
        FastHashMap<ShardIdent, (bool, ProcessedToByPartitions)>,

    shards: Option<FastHashMap<ShardIdent, Box<ShardDescription>>>,

    pub shard_fees: ShardFees,

    pub mint_msg: Option<InMsg>,
    pub recover_create_msg: Option<InMsg>,

    pub value_flow: PartialValueFlow,
    pub min_ref_mc_seqno: u32,

    pub rand_seed: HashBytes,

    pub created_by: HashBytes,

    /// Mempool config override for a new genesis
    pub mempool_config_override: Option<MempoolGlobalConfig>,

    /// If current block is a key master block and `ConsensusConfig` was changed.
    /// `None` - if it is a shard block or not a key master block.
    pub consensus_config_changed: Option<bool>,

    pub diff_tail_len: u32,
}

impl BlockCollationData {
    pub fn get_gen_chain_time(&self) -> u64 {
        self.gen_utime as u64 * 1000 + self.gen_utime_ms as u64
    }

    /// Updates [`Self::value_flow`] and returns its full version.
    pub fn finalize_value_flow(
        &mut self,
        account_blocks: &AccountBlocks,
        shard_accounts: &ShardAccounts,
        in_msgs: &InMsgDescr,
        out_msgs: &OutMsgDescr,
        config: &BlockchainConfig,
    ) -> Result<ValueFlow> {
        let is_masterchain = self.block_id_short.is_masterchain();
        let burning = config.get_burning_config().unwrap_or_default();

        // Add message and transaction fees collected in this block.
        let mut new_fees = account_blocks.root_extra().clone();
        new_fees.try_add_assign_tokens(in_msgs.root_extra().fees_collected)?;

        // Burn a fraction of fees in masterchain.
        if is_masterchain {
            let burned = burning.compute_burned_fees(new_fees.tokens)?;

            self.value_flow.burned.try_add_assign_tokens(burned)?;
            new_fees.try_sub_assign_tokens(burned)?;
        }

        self.value_flow.fees_collected.try_add_assign(&new_fees)?;

        // Finalize value flow.
        Ok(ValueFlow {
            from_prev_block: self.value_flow.from_prev_block.clone(),
            to_next_block: shard_accounts.root_extra().balance.clone(),
            imported: in_msgs.root_extra().value_imported.clone(),
            exported: out_msgs.root_extra().clone(),
            fees_collected: self.value_flow.fees_collected.clone(),
            burned: self.value_flow.burned.clone(),
            fees_imported: self.shard_fees.root_extra().fees.clone(),
            recovered: self.value_flow.recovered.clone(),
            created: self.value_flow.created.clone(),
            minted: self.value_flow.minted.clone(),
        })
    }
}

#[derive(Default, Debug)]
pub struct PartialValueFlow {
    /// Total amount transferred from the previous block.
    pub from_prev_block: CurrencyCollection,
    /// Fee recovery.
    pub recovered: CurrencyCollection,
    /// Block creation fees.
    pub created: CurrencyCollection,
    /// Minted native tokens and extra currencies.
    pub minted: CurrencyCollection,
    /// Burned native tokens and extra currencies.
    pub burned: CurrencyCollection,
    /// Fees collected on init.
    pub fees_collected: CurrencyCollection,
}

#[derive(Debug)]
pub struct BlockLimitStats {
    pub gas_used: u64,
    pub lt_current: u64,
    pub total_accounts: u32,
    pub total_items: u32,
    pub block_limits: BlockLimits,
}

impl BlockLimitStats {
    pub fn new(block_limits: BlockLimits, lt_start: u64) -> Self {
        Self {
            gas_used: 0,
            lt_current: lt_start,
            total_accounts: 0,
            total_items: 0,
            block_limits,
        }
    }

    pub fn reached(&self, level: BlockLimitsLevel) -> Option<BlockLimitReachType> {
        let BlockLimits {
            bytes,
            gas,
            lt_delta,
        } = &self.block_limits;

        let BlockParamLimits {
            soft_limit,
            hard_limit,
            ..
        } = bytes;

        if self.total_accounts >= *hard_limit {
            return Some(BlockLimitReachType::TotalAccountsHard);
        }
        if self.total_accounts >= *soft_limit && level == BlockLimitsLevel::Soft {
            return Some(BlockLimitReachType::TotalAccountsSoft);
        }

        let BlockParamLimits {
            soft_limit,
            hard_limit,
            ..
        } = gas;

        if self.gas_used >= *hard_limit as u64 {
            return Some(BlockLimitReachType::GasUsedHard);
        }
        if self.gas_used >= *soft_limit as u64 && level == BlockLimitsLevel::Soft {
            return Some(BlockLimitReachType::GasUsedSoft);
        }

        let BlockParamLimits {
            soft_limit,
            hard_limit,
            ..
        } = lt_delta;

        // created transactions + created out messages
        if self.total_items >= *hard_limit {
            return Some(BlockLimitReachType::TotalItemsHard);
        }
        if self.total_items >= *soft_limit && level == BlockLimitsLevel::Soft {
            return Some(BlockLimitReachType::TotalItemsSoft);
        }
        None
    }
}

#[derive(Debug)]
pub enum BlockLimitReachType {
    TotalAccountsSoft,
    TotalAccountsHard,
    TotalItemsSoft,
    TotalItemsHard,
    GasUsedSoft,
    GasUsedHard,
}

impl BlockLimitReachType {
    pub fn metric_label(&self) -> &'static str {
        match self {
            BlockLimitReachType::TotalAccountsSoft => "total_accounts_soft",
            BlockLimitReachType::TotalAccountsHard => "total_accounts_hard",
            BlockLimitReachType::TotalItemsSoft => "total_items_soft",
            BlockLimitReachType::TotalItemsHard => "total_items_hard",
            BlockLimitReachType::GasUsedSoft => "gas_used_soft",
            BlockLimitReachType::GasUsedHard => "gas_used_hard",
        }
    }
}

#[derive(Debug, Clone, Copy, Eq, Ord, PartialEq, PartialOrd)]
pub enum BlockLimitsLevel {
    Soft,
    Hard,
}

#[derive(Debug)]
pub struct PreparedInMsg {
    pub in_msg: Lazy<InMsg>,
    pub import_fees: ImportFees,
}

#[derive(Debug)]
pub struct PreparedOutMsg {
    pub out_msg: Lazy<OutMsg>,
    pub exported_value: CurrencyCollection,
    pub new_tx: Option<Lazy<Transaction>>,
}

impl BlockCollationData {
    pub fn shards(&self) -> Option<&FastHashMap<ShardIdent, Box<ShardDescription>>> {
        self.shards.as_ref()
    }

    pub fn get_shards(&self) -> Result<&FastHashMap<ShardIdent, Box<ShardDescription>>> {
        self.shards
            .as_ref()
            .ok_or_else(|| anyhow!("`shards` is not initialized yet"))
    }

    pub fn get_shards_mut(
        &mut self,
    ) -> Result<&mut FastHashMap<ShardIdent, Box<ShardDescription>>> {
        self.shards
            .as_mut()
            .ok_or_else(|| anyhow!("`shards` is not initialized yet"))
    }

    pub fn update_ref_min_mc_seqno(&mut self, mc_seqno: u32) -> u32 {
        self.min_ref_mc_seqno = std::cmp::min(self.min_ref_mc_seqno, mc_seqno);
        self.min_ref_mc_seqno
    }
}

#[derive(Debug, Default)]
pub(super) struct CollatorStats {
    pub total_execute_msgs_time_mc: u128,
    pub avg_exec_msgs_per_1000_ms: u128,

    pub total_execute_count_all: u64,
    pub total_execute_count_ext: u64,
    pub total_execute_count_int: u64,
    pub total_execute_count_new_int: u64,
    pub int_queue_length: u64,
    pub tps_block: u32,
    pub tps_timer: Option<std::time::Instant>,
    pub tps_execute_count: u64,
    pub tps: u128,
}

#[allow(unused_imports)]
pub(super) use super::anchors_cache::{AnchorInfo, AnchorsCache, CachedAnchor};

pub(super) type AccountId = HashBytes;

#[derive(Clone)]
pub(super) struct ShardAccountStuff {
    pub workchain_id: i8,
    pub account_addr: AccountId,
    pub shard_account: ShardAccount,
    pub special: SpecialFlags,
    pub initial_state_hash: HashBytes,
    pub balance: CurrencyCollection,
    pub public_libs_diff: AccountPublicLibsDiff,
    pub exists: AccountExistence,
    pub transactions: BTreeMap<u64, (CurrencyCollection, Lazy<Transaction>)>,
    pub min_next_lt: u64,
}

#[derive(Clone, Copy, PartialEq, Eq)]
pub(super) enum AccountExistence {
    Exists,
    NotExists,
    Created,
    Removed,
}

pub type AccountPublicLibsDiff = FastHashMap<HashBytes, Option<Cell>>;

impl ShardAccountStuff {
    pub fn new(
        workchain_id: i8,
        account_addr: &AccountId,
        shard_account: ShardAccount,
    ) -> Result<Self> {
        let initial_state_hash = *shard_account.account.inner().repr_hash();

        let mut special = SpecialFlags::default();
        let balance;
        let exists;

        if let Some(account) = shard_account.load_account()? {
            if let AccountState::Active(StateInit {
                special: acc_flags, ..
            }) = account.state
            {
                special = acc_flags.unwrap_or_default();
            }
            balance = account.balance;
            exists = AccountExistence::Exists;
        } else {
            balance = CurrencyCollection::ZERO;
            exists = AccountExistence::NotExists;
        }

        Ok(Self {
            workchain_id,
            account_addr: *account_addr,
            shard_account,
            special,
            initial_state_hash,
            balance,
            public_libs_diff: FastHashMap::new(),
            exists,
            transactions: Default::default(),
            // NOTE: Since there were no transactions on this account yet we don't
            //       need to additionaly align the execution LT.
            //       (considering that the block LT is also aligned)
            min_next_lt: 0,
        })
    }

    pub fn new_empty(workchain_id: i8, account_addr: &AccountId) -> Self {
        static EMPTY_SHARD_ACCOUNT: OnceLock<ShardAccount> = OnceLock::new();

        let shard_account = EMPTY_SHARD_ACCOUNT
            .get_or_init(|| ShardAccount {
                account: Lazy::new(&OptionalAccount::EMPTY).unwrap(),
                last_trans_hash: Default::default(),
                last_trans_lt: 0,
            })
            .clone();

        let initial_state_hash = *shard_account.account.inner().repr_hash();

        Self {
            workchain_id,
            account_addr: *account_addr,
            shard_account,
            special: Default::default(),
            initial_state_hash,
            balance: CurrencyCollection::ZERO,
            public_libs_diff: FastHashMap::new(),
            exists: AccountExistence::NotExists,
            transactions: Default::default(),
            // NOTE: No need to align logical time for non-existing accounts.
            min_next_lt: 0,
        }
    }

    pub fn is_empty(&self) -> Result<bool> {
        Ok(self.shard_account.load_account()?.is_none())
    }

    /// Aligns the provided `lt` with the logical time of the last transaction on account.
    /// Prevents reusing the same lt for uninit accounts in a single execution group.
    pub fn align_min_lt(&self, lt: u64) -> u64 {
        std::cmp::max(lt, self.min_next_lt)
    }

    pub fn make_std_addr(&self) -> StdAddr {
        StdAddr::new(self.workchain_id, self.account_addr)
    }

    pub fn build_hash_update(&self) -> Lazy<HashUpdate> {
        Lazy::new(&HashUpdate {
            old: self.initial_state_hash,
            new: *self.shard_account.account.inner().repr_hash(),
        })
        .unwrap()
    }

    pub fn apply_transaction(
        &mut self,
        new_state: ShardAccount,
        account_meta: AccountMeta,
        tx: Lazy<Transaction>,
        tx_meta: &TransactionMeta,
        mut public_libs_diff: Vec<PublicLibraryChange>,
    ) {
        use std::collections::hash_map;

        let is_masterchain = self.workchain_id as i32 == ShardIdent::MASTERCHAIN.workchain();
        debug_assert!(
            is_masterchain || public_libs_diff.is_empty(),
            "non-empty public libs diff for a non-masterchain account"
        );

        debug_assert!(new_state.last_trans_lt >= self.min_next_lt);

        let lt = new_state.last_trans_lt;

        self.shard_account = new_state;
        self.transactions
            .insert(lt, (tx_meta.total_fees.into(), tx));
        self.balance = account_meta.balance;
        self.min_next_lt = tx_meta.next_lt;

        // update account exists flag
        self.exists = match (self.exists, account_meta.exists) {
            (AccountExistence::NotExists, true) => AccountExistence::Created,
            (AccountExistence::Created, false) => AccountExistence::NotExists,
            (AccountExistence::Exists, false) => AccountExistence::Removed,
            (AccountExistence::Removed, true) => AccountExistence::Exists,
            (state, _) => state,
        };

        if is_masterchain {
            // Sort diff in reverse order (sort must be stable here).
            public_libs_diff.sort_by(|a, b| a.lib_hash().cmp(b.lib_hash()).reverse());

            // Merge diff with the previous.
            let mut prev_hash = None::<HashBytes>;
            for op in public_libs_diff {
                let hash = op.lib_hash();
                if matches!(&prev_hash, Some(prev_hash) if prev_hash == hash) {
                    // Skip duplicates (only the last change will be used because of reverse order).
                    continue;
                }
                prev_hash = Some(*hash);

                match self.public_libs_diff.entry(*hash) {
                    hash_map::Entry::Vacant(entry) => {
                        entry.insert(match op {
                            PublicLibraryChange::Add(cell) => Some(cell),
                            PublicLibraryChange::Remove(_) => None,
                        });
                    }
                    hash_map::Entry::Occupied(entry) => match (entry.get(), op) {
                        // Removed before, added later -> no diff.
                        // Added before, removed later -> no diff.
                        (None, PublicLibraryChange::Add(_))
                        | (Some(_), PublicLibraryChange::Remove(_)) => {
                            entry.remove();
                        }
                        // Removed before, removed now -> diff persists.
                        // Added before, added now -> diff persists.
                        (None, PublicLibraryChange::Remove(_))
                        | (Some(_), PublicLibraryChange::Add(_)) => {}
                    },
                }
            }
        }
    }
}

pub struct PublicLibsDiff {
    pub original: Dict<HashBytes, LibDescr>,
    pub changes: BTreeMap<HashBytes, PublicLibrariesDiffItem>,
}

impl PublicLibsDiff {
    pub fn new(original: Dict<HashBytes, LibDescr>) -> Self {
        Self {
            original,
            changes: BTreeMap::new(),
        }
    }

    pub fn finalize(mut self) -> Result<Dict<HashBytes, LibDescr>> {
        const SMALL_CHANGES: usize = 3;

        match self.changes.len() {
            0 => {}
            // Simple set is faster on a small number of changes.
            1..=SMALL_CHANGES => {
                for (lib_hash, change) in self.changes {
                    match change.finalize()? {
                        None => {
                            self.original.remove(lib_hash)?;
                        }
                        Some(new_descr) => {
                            self.original.set(lib_hash, new_descr)?;
                        }
                    }
                }
            }
            // Otherwise apply a full modify.
            _ => {
                self.original.modify_with_sorted_iter_ext(
                    self.changes,
                    |(key, _)| *key,
                    |(_, value)| value.finalize(),
                    Cell::empty_context(),
                )?;
            }
        }

        Ok(self.original)
    }

    pub fn merge(&mut self, account: &HashBytes, diff: AccountPublicLibsDiff) -> Result<()> {
        use std::collections::btree_map;

        for (lib_hash, new_lib) in diff {
            match self.changes.entry(lib_hash) {
                // A first change of this library in the block.
                btree_map::Entry::Vacant(entry) => {
                    // Find an existing published library first.
                    let existing = self.original.get(lib_hash)?;

                    let (origin, add) = match (existing, new_lib) {
                        // Trying to remove a non-existing library does nothing.
                        // TODO: Check if this even possible, maybe panic/error instead.
                        (None, None) => {
                            tracing::warn!(
                                target: tracing_targets::COLLATOR,
                                %account, %lib_hash,
                                "removing a nonexisting library",
                            );
                            continue;
                        }
                        // Publishing a new library.
                        (None, Some(new_lib)) => (PublicLibraryOrigin::New(new_lib), true),
                        // Library exists, only update publishers.
                        (Some(existing), new_lib) => {
                            (PublicLibraryOrigin::Existing(existing), new_lib.is_some())
                        }
                    };

                    entry.insert(PublicLibrariesDiffItem {
                        origin,
                        publishers: vec![(*account, add)],
                    });
                }
                // Changes to this library by multiple accounts.
                btree_map::Entry::Occupied(mut entry) => {
                    let entry = entry.get_mut();
                    if let PublicLibraryOrigin::New(_) = &entry.origin {
                        // Trying to remove a non-existing library does nothing.
                        // TODO: Check if this even possible, maybe panic/error instead.
                        if new_lib.is_none() {
                            tracing::warn!(
                                target: tracing_targets::COLLATOR,
                                %account, %lib_hash,
                                "removing a nonexisting library",
                            );
                            continue;
                        }
                    }

                    entry.publishers.push((*account, new_lib.is_some()));
                }
            }
        }

        Ok(())
    }
}

pub struct PublicLibrariesDiffItem {
    pub origin: PublicLibraryOrigin,
    pub publishers: Vec<(HashBytes, bool)>,
}

impl PublicLibrariesDiffItem {
    pub fn finalize(mut self) -> Result<Option<LibDescr>, tycho_types::error::Error> {
        match self.origin {
            // New `LibDescr` can be simply built from a list of publishers.
            PublicLibraryOrigin::New(lib) => {
                if self.publishers.is_empty() {
                    // Might be unreachable but handle just in case.
                    return Ok(None);
                }

                self.publishers.sort_unstable_by_key(|(a, _)| *a);
                let root = dict::build_dict_from_sorted_iter(
                    self.publishers.into_iter().map(|(key, add)| {
                        debug_assert!(add, "must not remove publishers from nonexisting libraries");
                        (key, ())
                    }),
                    Cell::empty_context(),
                )?;
                debug_assert!(root.is_some());

                Ok(Some(LibDescr {
                    lib,
                    publishers: Dict::from_raw(root),
                }))
            }
            // Existing `LibDescr` must be modified.
            PublicLibraryOrigin::Existing(mut descr) => {
                if self.publishers.is_empty() {
                    return Ok(Some(descr));
                }

                self.publishers.sort_unstable_by_key(|(a, _)| *a);
                descr.publishers.modify_with_sorted_iter_ext(
                    self.publishers,
                    |(key, _)| *key,
                    |(_, add)| Ok(add.then_some(())),
                    Cell::empty_context(),
                )?;
                if descr.publishers.is_empty() {
                    // All publishers were removed so the library should also be removed.
                    return Ok(None);
                }

                Ok(Some(descr))
            }
        }
    }
}

pub enum PublicLibraryOrigin {
    Existing(LibDescr),
    New(Cell),
}

pub trait ShardDescriptionExt {
    fn from_block_info(
        block_id: BlockId,
        block_info: &BlockInfo,
        ext_processed_to_anchor_id: u32,
        value_flow: &ValueFlow,
    ) -> ShardDescription;
}

impl ShardDescriptionExt for ShardDescription {
    fn from_block_info(
        block_id: BlockId,
        block_info: &BlockInfo,
        ext_processed_to_anchor_id: u32,
        value_flow: &ValueFlow,
    ) -> ShardDescription {
        ShardDescription {
            seqno: block_id.seqno,
            reg_mc_seqno: 0,
            start_lt: block_info.start_lt,
            end_lt: block_info.end_lt,
            root_hash: block_id.root_hash,
            file_hash: block_id.file_hash,
            before_split: block_info.before_split, // not used in the current implementation
            before_merge: false,                   // not used in the current implementation
            want_split: block_info.want_split,     // not used in the current implementation
            want_merge: block_info.want_merge,     // not used in the current implementation
            nx_cc_updated: false,                  // not used in the current implementation
            next_catchain_seqno: block_info.gen_catchain_seqno,
            ext_processed_to_anchor_id,
            top_sc_block_updated: false, // initially set false, will be updated further
            min_ref_mc_seqno: block_info.min_ref_mc_seqno,
            gen_utime: block_info.gen_utime,
            split_merge_at: None, // not used in the current implementation
            fees_collected: value_flow.fees_collected.clone(),
            funds_created: value_flow.created.clone(),
        }
    }
}

#[derive(Clone)]
pub struct BlockSerializerCache {
    inner: Arc<Mutex<BlockSerializerCacheInner>>,
}

impl BlockSerializerCache {
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            inner: Arc::new(Mutex::new(BlockSerializerCacheInner {
                boc_header: BocHeaderCache::with_capacity(capacity),
            })),
        }
    }

    pub fn take_boc_header_cache(&self) -> BocHeaderCache {
        std::mem::take(&mut self.inner.lock().boc_header)
    }

    pub fn set_boc_header_cache(&self, boc_header: BocHeaderCache) {
        self.inner.lock().boc_header = boc_header;
    }
}

struct BlockSerializerCacheInner {
    boc_header: BocHeaderCache,
}

type BocHeaderCache = boc::ser::BocHeaderCache<ahash::RandomState>;

#[derive(Clone, Debug)]
pub struct ExecutedTransaction {
    pub transaction: Lazy<Transaction>,
    pub out_msgs: Vec<Lazy<OwnedMessage>>,
    pub gas_used: u64,
    pub next_lt: u64,
    pub burned: Tokens,
}

#[derive(Clone, Debug)]
pub struct SkippedTransaction {
    pub gas_used: u64,
}

#[derive(Clone, Debug)]
pub struct OutMessageData {
    pub int_msg_info: IntMsgInfo,
    pub cell: Cell,
    pub dst_in_current_shard: bool,
}

#[derive(Debug)]
pub struct ParsedMessageInner {
    pub info: MsgInfo,
    pub dst_in_current_shard: bool,
    pub cell: Cell,
    pub special_origin: Option<SpecialOrigin>,
    pub block_seqno: Option<BlockSeqno>,
    pub from_same_shard: Option<bool>,
    pub ext_msg_chain_time: Option<u64>,
}

#[derive(Clone, Debug)]
pub struct ParsedMessage(Arc<ParsedMessageInner>);

impl ParsedMessage {
    /// Creates a new parsed internal message.
    pub fn from_int(
        info: IntMsgInfo,
        cell: Cell,
        dst_in_current_shard: bool,
        block_seqno: Option<BlockSeqno>,
        from_same_shard: Option<bool>,
    ) -> Self {
        Self(Arc::new(ParsedMessageInner {
            info: MsgInfo::Int(info),
            dst_in_current_shard,
            cell,
            special_origin: None,
            block_seqno,
            from_same_shard,
            ext_msg_chain_time: None,
        }))
    }

    /// Creates a new parsed external message.
    pub fn from_ext(
        info: ExtInMsgInfo,
        cell: Cell,
        dst_in_current_shard: bool,
        chain_time: u64,
    ) -> Self {
        Self(Arc::new(ParsedMessageInner {
            info: MsgInfo::ExtIn(info),
            dst_in_current_shard,
            cell,
            special_origin: None,
            block_seqno: None,
            from_same_shard: None,
            ext_msg_chain_time: Some(chain_time),
        }))
    }

    /// Creates a new message with special origin.
    pub fn from_special(
        info: MsgInfo,
        cell: Cell,
        special_origin: SpecialOrigin,
        block_seqno: BlockSeqno,
    ) -> Self {
        Self(Arc::new(ParsedMessageInner {
            info,
            dst_in_current_shard: true,
            cell,
            special_origin: Some(special_origin),
            block_seqno: Some(block_seqno),
            from_same_shard: None,
            ext_msg_chain_time: None,
        }))
    }

    pub fn block_seqno(&self) -> Option<BlockSeqno> {
        self.0.block_seqno
    }
    pub fn is_from_same_shard(&self) -> Option<bool> {
        self.0.from_same_shard
    }
    pub fn info(&self) -> &MsgInfo {
        &self.0.info
    }
    pub fn cell(&self) -> &Cell {
        &self.0.cell
    }
    pub fn ext_msg_chain_time(&self) -> Option<u64> {
        self.0.ext_msg_chain_time
    }
    pub fn dst_in_current_shard(&self) -> bool {
        self.0.dst_in_current_shard
    }
    pub fn special_origin(&self) -> Option<SpecialOrigin> {
        self.0.special_origin
    }

    pub fn kind(&self) -> ParsedMessageKind {
        match (self.info(), self.special_origin()) {
            (_, Some(SpecialOrigin::Recover)) => ParsedMessageKind::Recover,
            (_, Some(SpecialOrigin::Mint)) => ParsedMessageKind::Mint,
            (MsgInfo::ExtIn(_), _) => ParsedMessageKind::ExtIn,
            (MsgInfo::Int(_), _) => ParsedMessageKind::Int,
            (MsgInfo::ExtOut(_), _) => ParsedMessageKind::ExtOut,
        }
    }

    pub fn is_external(&self) -> bool {
        matches!(self.info(), MsgInfo::ExtIn(_) | MsgInfo::ExtOut(_))
    }

    pub fn arc(&self) -> Arc<ParsedMessageInner> {
        Arc::clone(&self.0)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ParsedMessageKind {
    Recover,
    Mint,
    ExtIn,
    Int,
    ExtOut,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SpecialOrigin {
    Recover,
    Mint,
}

pub struct FinalizeMessagesReaderResult {
    pub has_unprocessed_messages: bool,
    pub queue_diff_with_msgs: QueueDiffWithMessages<EnqueuedMessage>,
    pub current_msgs_exec_params: MsgsExecutionParams,
}

pub struct FinalizeCollationResult {
    pub collation_result: BlockCollationResult,
}

pub struct ExecuteResult {
    pub prepare_msg_groups_wu: PrepareMsgGroupsWu,
    /// Accumulated messages reader metrics across all partitions
    pub msgs_reader_metrics: MessagesReaderMetrics,

    pub execute_wu: ExecuteWu,
    pub execute_metrics: ExecuteMetrics,
}

pub struct FinalizeBlockResult {
    pub collation_data: Box<BlockCollationData>,
    pub block_candidate: Box<BlockCandidate>,
    pub mc_data: Option<Arc<McData>>,
    pub old_mc_data: Arc<McData>,
    pub merkle_build: MerkleBuildResult,
    pub new_state_root: Cell,
    pub new_observable_state: Box<ShardStateUnsplit>,
    pub finalize_wu: FinalizeWu,
    pub finalize_metrics: FinalizeMetrics,
    pub do_collate_wu: DoCollateWu,
    pub collation_config: Arc<CollationConfig>,
}

pub struct CollationResult {
    pub final_result: FinalResult,
    pub finalized: FinalizeBlockResult,
    pub execute_result: ExecuteResult,
}

pub struct FinalResult {
    pub prepare_elapsed: Duration,
    pub has_unprocessed_messages: bool,
}

#[derive(Debug)]
pub enum ForceMasterCollation {
    No,
    ByUncommittedChain,
    ByAnchorImportSkipped,
    ByUnprocessedMessages,
    NoPendingMessagesAfterShardBlocks,
}
impl ForceMasterCollation {
    pub fn is_forced(&self) -> bool {
        !matches!(self, Self::No | Self::NoPendingMessagesAfterShardBlocks)
    }
}

/// Rand seed for block source data.
#[derive(Debug, Clone, Hash, PartialEq, Eq, TlWrite)]
#[tl(boxed, id = "collator.randSeed", scheme = "proto.tl")]
pub struct RandSeed {
    #[tl(with = "tycho_block_util::tl::shard_ident")]
    pub shard: ShardIdent,
    pub seqno: u32,
    pub next_chain_time: u64,
}

pub trait MsgsExecutionParamsExtension {
    fn group_slots_fractions(&self) -> Result<BTreeMap<QueuePartitionIdx, u8>>;
    fn open_ranges_limit(&self) -> usize;
}

impl MsgsExecutionParamsExtension for MsgsExecutionParams {
    fn group_slots_fractions(&self) -> Result<BTreeMap<QueuePartitionIdx, u8>> {
        let mut res = BTreeMap::new();
        for item in self.group_slots_fractions.iter() {
            let (par_id, fraction) = item?;
            res.insert(QueuePartitionIdx(par_id), fraction);
        }
        Ok(res)
    }

    fn open_ranges_limit(&self) -> usize {
        self.open_ranges_limit.max(2) as usize
    }
}

#[derive(Default, Clone)]
struct MsgsExecutionParamsStuffInner {
    current: MsgsExecutionParams,
    new: Option<MsgsExecutionParams>,
}

#[derive(Default, Clone)]
pub(super) struct MsgsExecutionParamsStuff {
    inner: Arc<RwLock<MsgsExecutionParamsStuffInner>>,
}

impl MsgsExecutionParamsStuff {
    pub fn create(
        prev_params: Option<MsgsExecutionParams>,
        mc_data_params: MsgsExecutionParams,
    ) -> Self {
        match prev_params {
            Some(prev_params) => {
                let new = (prev_params != mc_data_params).then_some(mc_data_params);
                Self {
                    inner: Arc::new(RwLock::new(MsgsExecutionParamsStuffInner {
                        current: prev_params,
                        new,
                    })),
                }
            }
            _ => Self {
                inner: Arc::new(RwLock::new(MsgsExecutionParamsStuffInner {
                    current: mc_data_params,
                    new: None,
                })),
            },
        }
    }

    #[cfg(test)]
    pub fn with_current_and_new(
        current: MsgsExecutionParams,
        new: Option<MsgsExecutionParams>,
    ) -> Self {
        Self {
            inner: Arc::new(RwLock::new(MsgsExecutionParamsStuffInner { current, new })),
        }
    }

    pub fn new_is_some(&self) -> bool {
        self.inner.read().new.is_some()
    }

    pub fn update(&self) {
        let mut guard = self.inner.write();

        let new = guard.new.take();

        if let Some(new) = new {
            guard.current = new;
        }
    }

    #[allow(clippy::wrong_self_convention, clippy::new_ret_no_self)]
    pub fn new(&self) -> MappedRwLockReadGuard<'_, Option<MsgsExecutionParams>> {
        RwLockReadGuard::map(self.inner.read(), |x| &x.new)
    }

    pub fn current(&self) -> MappedRwLockReadGuard<'_, MsgsExecutionParams> {
        RwLockReadGuard::map(self.inner.read(), |x| &x.current)
    }
}

#[derive(Debug, Default)]
pub(super) struct ExecuteMetrics {
    pub execute_groups_vm_only_timer: MetricsTimer,
    pub process_txs_timer: MetricsTimer,

    pub execute_tick_elapsed: Duration,
    pub execute_special_elapsed: Duration,
    pub execute_incoming_msgs_elapsed: Duration,
    pub execute_tock_elapsed: Duration,
}

impl ExecuteMetrics {
    pub fn execute_groups_total_elapsed(&self) -> Duration {
        self.execute_groups_vm_only_timer
            .total_elapsed
            .saturating_add(self.process_txs_timer.total_elapsed)
    }
}

#[derive(Debug, Default)]
pub(super) struct FinalizeMetrics {
    /// Measures queue diff creation together with queue diff serialization
    pub create_queue_diff_timer: MetricsTimer,

    pub apply_queue_diff_elapsed: Duration,

    pub update_shard_accounts_elapsed: Duration,
    pub build_accounts_blocks_elapsed: Duration,

    pub build_accounts_elapsed: Duration,

    pub build_in_msgs_elapsed: Duration,
    pub build_out_msgs_elapsed: Duration,

    pub build_accounts_and_messages_in_parallel_elased: Duration,

    pub build_mc_state_extra_elapsed: Duration,

    pub build_state_update_elapsed: Duration,
    pub build_block_elapsed: Duration,

    pub finalize_block_elapsed: Duration,

    pub total_timer: MetricsTimer,
}
