#![allow(clippy::mutable_key_type)]

use crate::state::state_db::{StateContext, StateTree};
use crate::traits::chain_store::ChainStore;
use crate::traits::kv_store::{KVStore, KVStoreWrite};
use crate::{smt::smt_store::SMTStore, traits::kv_store::KVStoreRead};
use gw_common::h256_ext::H256Ext;
use gw_common::{merkle_utils::calculate_state_checkpoint, smt::SMT, H256};
use gw_db::schema::{
    Col, COLUMN_ACCOUNT_SMT_BRANCH, COLUMN_ACCOUNT_SMT_LEAF, COLUMN_ASSET_SCRIPT,
    COLUMN_BAD_BLOCK_CHALLENGE_TARGET, COLUMN_BLOCK, COLUMN_BLOCK_DEPOSIT_REQUESTS,
    COLUMN_BLOCK_GLOBAL_STATE, COLUMN_BLOCK_SMT_BRANCH, COLUMN_BLOCK_SMT_LEAF,
    COLUMN_BLOCK_STATE_RECORD, COLUMN_BLOCK_STATE_REVERSE_RECORD, COLUMN_INDEX,
    COLUMN_L2BLOCK_COMMITTED_INFO, COLUMN_MEM_POOL_TRANSACTION,
    COLUMN_MEM_POOL_TRANSACTION_RECEIPT, COLUMN_MEM_POOL_WITHDRAWAL, COLUMN_META,
    COLUMN_REVERTED_BLOCK_SMT_BRANCH, COLUMN_REVERTED_BLOCK_SMT_LEAF,
    COLUMN_REVERTED_BLOCK_SMT_ROOT, COLUMN_TRANSACTION, COLUMN_TRANSACTION_INFO,
    COLUMN_TRANSACTION_RECEIPT, COLUMN_WITHDRAWAL, COLUMN_WITHDRAWAL_INFO, META_BLOCK_SMT_ROOT_KEY,
    META_CHAIN_ID_KEY, META_LAST_VALID_TIP_BLOCK_HASH_KEY, META_REVERTED_BLOCK_SMT_ROOT_KEY,
    META_TIP_BLOCK_HASH_KEY,
};
use gw_db::{error::Error, iter::DBIter, DBIterator, IteratorMode, RocksDBTransaction};
use gw_db::{DBRawIterator, Direction};
use gw_types::packed::{Script, WithdrawalKey};
use gw_types::{
    packed::{
        self, AccountMerkleState, Byte32, ChallengeTarget, TransactionKey, WithdrawalReceipt,
    },
    prelude::*,
};
use std::collections::HashSet;

use crate::state::block_state_record::{BlockStateRecordKey, BlockStateRecordKeyReverse};

pub struct StoreTransaction {
    pub(crate) inner: RocksDBTransaction,
}

impl KVStoreRead for StoreTransaction {
    fn get(&self, col: Col, key: &[u8]) -> Option<Box<[u8]>> {
        self.inner
            .get(col, key)
            .expect("db operation should be ok")
            .map(|v| Box::<[u8]>::from(v.as_ref()))
    }
}

impl KVStoreWrite for StoreTransaction {
    fn insert_raw(&self, col: Col, key: &[u8], value: &[u8]) -> Result<(), Error> {
        self.inner.put(col, key, value)
    }

    fn delete(&self, col: Col, key: &[u8]) -> Result<(), Error> {
        self.inner.delete(col, key)
    }
}

impl KVStore for StoreTransaction {}

impl StoreTransaction {
    pub fn commit(&self) -> Result<(), Error> {
        self.inner.commit()
    }

    pub fn rollback(&self) -> Result<(), Error> {
        self.inner.rollback()
    }

    pub(crate) fn get_iter(&self, col: Col, mode: IteratorMode) -> DBIter {
        self.inner
            .iter(col, mode)
            .expect("db operation should be ok")
    }

    pub fn setup_chain_id(&self, chain_id: H256) -> Result<(), Error> {
        self.insert_raw(COLUMN_META, META_CHAIN_ID_KEY, chain_id.as_slice())?;
        Ok(())
    }

    pub fn set_block_smt_root(&self, root: H256) -> Result<(), Error> {
        self.insert_raw(COLUMN_META, META_BLOCK_SMT_ROOT_KEY, root.as_slice())?;
        Ok(())
    }

    pub fn set_tip_block_hash(&self, block_hash: H256) -> Result<(), Error> {
        let block_hash: [u8; 32] = block_hash.into();
        self.insert_raw(COLUMN_META, META_TIP_BLOCK_HASH_KEY, &block_hash)
    }

    pub fn set_bad_block_challenge_target(
        &self,
        block_hash: &H256,
        target: &ChallengeTarget,
    ) -> Result<(), Error> {
        self.insert_raw(
            COLUMN_BAD_BLOCK_CHALLENGE_TARGET,
            block_hash.as_slice(),
            target.as_slice(),
        )
    }

    pub fn set_reverted_block_hashes(
        &self,
        reverted_block_smt_root: &H256,
        block_hashes: Vec<H256>,
    ) -> Result<(), Error> {
        assert!(!block_hashes.is_empty(), "set empty reverted block hashes");

        self.insert_raw(
            COLUMN_REVERTED_BLOCK_SMT_ROOT,
            reverted_block_smt_root.as_slice(),
            block_hashes.pack().as_slice(),
        )
    }

    #[allow(clippy::too_many_arguments)]
    pub fn insert_block(
        &self,
        block: packed::L2Block,
        committed_info: packed::L2BlockCommittedInfo,
        global_state: packed::GlobalState,
        withdrawal_receipts: Vec<WithdrawalReceipt>,
        prev_txs_state: AccountMerkleState,
        tx_receipts: Vec<packed::TxReceipt>,
        deposit_requests: Vec<packed::DepositRequest>,
    ) -> Result<(), Error> {
        debug_assert_eq!(block.transactions().len(), tx_receipts.len());
        let block_hash = block.hash();
        self.insert_raw(COLUMN_BLOCK, &block_hash, block.as_slice())?;
        self.insert_raw(
            COLUMN_L2BLOCK_COMMITTED_INFO,
            &block_hash,
            committed_info.as_slice(),
        )?;
        self.insert_raw(
            COLUMN_BLOCK_GLOBAL_STATE,
            &block_hash,
            global_state.as_slice(),
        )?;
        let deposit_requests_vec: packed::DepositRequestVec = deposit_requests.pack();
        self.insert_raw(
            COLUMN_BLOCK_DEPOSIT_REQUESTS,
            &block_hash,
            deposit_requests_vec.as_slice(),
        )?;

        // Verify prev tx state and insert
        {
            let prev_txs_state_checkpoint = {
                let txs = block.as_reader().raw().submit_transactions();
                txs.prev_state_checkpoint().to_entity()
            };

            let root: [u8; 32] = prev_txs_state.merkle_root().unpack();
            let count: u32 = prev_txs_state.count().unpack();
            let checkpoint: Byte32 = {
                let checkpoint: [u8; 32] = calculate_state_checkpoint(&root.into(), count).into();
                checkpoint.pack()
            };
            if checkpoint != prev_txs_state_checkpoint {
                let root: Byte32 = root.pack();
                log::debug!("root: {} count: {}", root, count);
                return Err(Error::from(format!(
                    "unexpected prev tx state, checkpoint: {} prev_txs_state_checkpoint: {}",
                    checkpoint, prev_txs_state_checkpoint
                )));
            }

            let block_post_state = block.as_reader().raw().post_account();
            if tx_receipts.is_empty() && prev_txs_state.as_slice() != block_post_state.as_slice() {
                log::debug!(
                    "tx_receipts: {} prev_txs_state: {} post_state: {}",
                    tx_receipts.len(),
                    prev_txs_state,
                    block_post_state
                );
                return Err(Error::from("unexpected no tx post state".to_string()));
            }
        }

        for (index, (tx, tx_receipt)) in block
            .transactions()
            .into_iter()
            .zip(tx_receipts.iter())
            .enumerate()
        {
            let key = TransactionKey::build_transaction_key(block_hash.pack(), index as u32);
            self.insert_raw(COLUMN_TRANSACTION, key.as_slice(), tx.as_slice())?;
            self.insert_raw(
                COLUMN_TRANSACTION_RECEIPT,
                key.as_slice(),
                tx_receipt.as_slice(),
            )?;
        }
        for (index, withdrawal) in block.withdrawals().into_iter().enumerate() {
            let key = WithdrawalKey::build_withdrawal_key(block_hash.pack(), index as u32);
            self.insert_raw(COLUMN_WITHDRAWAL, key.as_slice(), withdrawal.as_slice())?;
        }

        let post_states: Vec<AccountMerkleState> = {
            let withdrawal_post_states = withdrawal_receipts.into_iter().map(|w| w.post_state());
            let tx_post_states = tx_receipts.iter().map(|t| t.post_state());
            withdrawal_post_states.chain(tx_post_states).collect()
        };

        let state_checkpoint_list = block.raw().state_checkpoint_list().into_iter();
        if post_states.len() != state_checkpoint_list.len() {
            return Err(Error::from("unexpected block post state length".to_owned()));
        }

        Ok(())
    }

    pub fn insert_asset_scripts(&self, scripts: HashSet<Script>) -> Result<(), Error> {
        for script in scripts.into_iter() {
            self.insert_raw(COLUMN_ASSET_SCRIPT, &script.hash(), script.as_slice())?;
        }

        Ok(())
    }

    pub fn block_smt(&self) -> Result<SMT<SMTStore<Self>>, Error> {
        let root = self.get_block_smt_root()?;
        let smt_store = SMTStore::new(COLUMN_BLOCK_SMT_LEAF, COLUMN_BLOCK_SMT_BRANCH, self);
        Ok(SMT::new(root, smt_store))
    }

    pub fn reverted_block_smt(&self) -> Result<SMT<SMTStore<Self>>, Error> {
        let root = self.get_reverted_block_smt_root()?;
        let smt_store = SMTStore::new(
            COLUMN_REVERTED_BLOCK_SMT_LEAF,
            COLUMN_REVERTED_BLOCK_SMT_BRANCH,
            self,
        );
        Ok(SMT::new(root, smt_store))
    }

    // TODO: prune db state
    pub fn get_reverted_block_hashes(&self) -> Result<HashSet<H256>, Error> {
        let iter = self.get_iter(COLUMN_REVERTED_BLOCK_SMT_LEAF, IteratorMode::End);
        let to_byte32 = iter.map(|(key, _value)| {
            packed::Byte32Reader::from_slice_should_be_ok(key.as_ref()).to_entity()
        });
        let to_h256 = to_byte32.map(|byte32| byte32.unpack());

        Ok(to_h256.collect())
    }

    pub fn rewind_reverted_block_smt(&self, block_hashes: Vec<H256>) -> Result<(), Error> {
        let mut reverted_block_smt = self.reverted_block_smt()?;

        for block_hash in block_hashes.into_iter() {
            reverted_block_smt
                .update(block_hash, H256::zero())
                .map_err(|err| Error::from(format!("reset reverted block smt error {}", err)))?;
        }

        self.set_reverted_block_smt_root(*reverted_block_smt.root())
    }

    pub fn rewind_block_smt(&self, block: &packed::L2Block) -> Result<(), Error> {
        let mut block_smt = self.block_smt()?;
        block_smt
            .update(block.smt_key().into(), H256::zero())
            .map_err(|err| Error::from(format!("reset block smt error {}", err)))?;

        self.set_block_smt_root(*block_smt.root())
    }

    fn set_last_valid_tip_block_hash(&self, block_hash: &H256) -> Result<(), Error> {
        self.insert_raw(
            COLUMN_META,
            META_LAST_VALID_TIP_BLOCK_HASH_KEY,
            block_hash.as_slice(),
        )
    }

    pub fn set_reverted_block_smt_root(&self, root: H256) -> Result<(), Error> {
        self.insert_raw(
            COLUMN_META,
            META_REVERTED_BLOCK_SMT_ROOT_KEY,
            root.as_slice(),
        )?;
        Ok(())
    }

    pub fn insert_bad_block(
        &self,
        block: &packed::L2Block,
        committed_info: &packed::L2BlockCommittedInfo,
        global_state: &packed::GlobalState,
    ) -> Result<(), Error> {
        let block_hash = block.hash();
        let block_number = block.raw().number();

        let committed_info = committed_info.as_slice();
        let global_state = global_state.as_slice();

        self.insert_raw(COLUMN_L2BLOCK_COMMITTED_INFO, &block_hash, committed_info)?;
        self.insert_raw(COLUMN_BLOCK_GLOBAL_STATE, &block_hash, global_state)?;

        self.insert_raw(COLUMN_BLOCK, &block_hash, block.as_slice())?;

        self.insert_raw(COLUMN_INDEX, block_number.as_slice(), &block_hash)?;
        self.insert_raw(COLUMN_INDEX, &block_hash, block_number.as_slice())?;

        // Add to block smt
        let mut block_smt = self.block_smt()?;
        block_smt
            .update(block.smt_key().into(), block_hash.into())
            .map_err(|err| Error::from(format!("block smt error {}", err)))?;
        self.set_block_smt_root(*block_smt.root())?;

        // Update tip block
        self.insert_raw(COLUMN_META, META_TIP_BLOCK_HASH_KEY, &block_hash)?;

        Ok(())
    }

    pub fn revert_bad_blocks(&self, bad_blocks: &[packed::L2Block]) -> Result<(), Error> {
        if bad_blocks.is_empty() {
            return Ok(());
        }

        let mut block_smt = self.block_smt()?;
        let mut reverted_block_smt = self.reverted_block_smt()?;

        for block in bad_blocks {
            let block_hash = block.hash();
            let block_number = block.raw().number();

            self.delete(COLUMN_INDEX, &block_hash)?;
            self.delete(COLUMN_INDEX, block_number.as_slice())?;

            // Remove block from smt
            block_smt
                .update(block.smt_key().into(), H256::zero())
                .map_err(|err| Error::from(format!("block smt error {}", err)))?;

            // Add block to reverted smt
            reverted_block_smt
                .update(block_hash.into(), H256::one())
                .map_err(|err| Error::from(format!("reverted block smt error {}", err)))?;
        }

        self.set_block_smt_root(*block_smt.root())?;
        self.set_reverted_block_smt_root(*reverted_block_smt.root())?;

        // Revert tip block to parent block
        let parent_block_hash: [u8; 32] = {
            let first_bad_block = bad_blocks.first().expect("exists");
            first_bad_block.raw().parent_block_hash().unpack()
        };
        self.insert_raw(COLUMN_META, META_TIP_BLOCK_HASH_KEY, &parent_block_hash)
    }

    /// Attach block to the rollup main chain
    pub fn attach_block(&self, block: packed::L2Block) -> Result<(), Error> {
        let raw = block.raw();
        let raw_number = raw.number();
        let block_hash = raw.hash();

        // build tx info
        for (index, tx) in block.transactions().into_iter().enumerate() {
            let key = TransactionKey::build_transaction_key(block_hash.pack(), index as u32);
            let info = packed::TransactionInfo::new_builder()
                .key(key)
                .block_number(raw_number.clone())
                .build();
            let tx_hash = tx.hash();
            self.insert_raw(COLUMN_TRANSACTION_INFO, &tx_hash, info.as_slice())?;
        }

        // build withdrawal info
        for (index, withdrawal) in block.withdrawals().into_iter().enumerate() {
            let key = WithdrawalKey::build_withdrawal_key(block_hash.pack(), index as u32);
            let info = packed::WithdrawalInfo::new_builder()
                .key(key)
                .block_number(raw_number.clone())
                .build();
            let withdrawal_hash = withdrawal.hash();
            self.insert_raw(COLUMN_WITHDRAWAL_INFO, &withdrawal_hash, info.as_slice())?;
        }

        // build main chain index
        self.insert_raw(COLUMN_INDEX, raw_number.as_slice(), &block_hash)?;
        self.insert_raw(COLUMN_INDEX, &block_hash, raw_number.as_slice())?;

        // update block tree
        let mut block_smt = self.block_smt()?;
        block_smt
            .update(raw.smt_key().into(), raw.hash().into())
            .map_err(|err| Error::from(format!("SMT error {}", err)))?;
        let root = block_smt.root();
        self.set_block_smt_root(*root)?;

        // update tip
        self.insert_raw(COLUMN_META, META_TIP_BLOCK_HASH_KEY, &block_hash)?;
        self.set_last_valid_tip_block_hash(&block_hash.into())?;

        Ok(())
    }

    /// Delete block from DB
    pub fn detach_block(&self, block: &packed::L2Block) -> Result<(), Error> {
        // check
        {
            let tip = self.get_last_valid_tip_block_hash()?;
            assert_eq!(tip, H256::from(block.raw().hash()), "Must detach from tip");
        }
        {
            let number: u64 = block.raw().number().unpack();
            let hash: Byte32 = block.hash().pack();
            log::warn!("detach block #{} {}", number, hash);
        }
        // remove transaction info
        for tx in block.transactions().into_iter() {
            let tx_hash = tx.hash();
            self.delete(COLUMN_TRANSACTION_INFO, &tx_hash)?;
        }

        let block_hash: H256 = block.hash().into();

        // remove index
        let block_number = block.raw().number();
        self.delete(COLUMN_INDEX, block_number.as_slice())?;
        self.delete(COLUMN_INDEX, block_hash.as_slice())?;

        // update block tree
        let mut block_smt = self.block_smt()?;
        block_smt
            .update(block.smt_key().into(), H256::zero())
            .map_err(|err| Error::from(format!("SMT error {}", err)))?;
        let root = block_smt.root();
        self.set_block_smt_root(*root)?;

        // update tip
        let block_number: u64 = block_number.unpack();
        let parent_number = block_number.saturating_sub(1);
        let parent_block_hash = self
            .get_block_hash_by_number(parent_number)?
            .expect("parent block hash");
        self.insert_raw(
            COLUMN_META,
            META_TIP_BLOCK_HASH_KEY,
            parent_block_hash.as_slice(),
        )?;
        self.set_last_valid_tip_block_hash(&parent_block_hash)?;

        Ok(())
    }

    pub fn remove_block_state_record(&self, block_number: u64) -> Result<(), Error> {
        let iter = self.iter_block_state_record(block_number);
        for record_key in iter {
            // delete record key
            self.delete(COLUMN_BLOCK_STATE_RECORD, record_key.as_slice())?;
            // delete reverse record key
            let reverse_key =
                BlockStateRecordKeyReverse::new(record_key.block_number(), &record_key.state_key());
            self.delete(COLUMN_BLOCK_STATE_REVERSE_RECORD, reverse_key.as_slice())?;
        }
        Ok(())
    }

    pub fn get_history_state(&self, block_number: u64, state_key: &H256) -> Option<H256> {
        let key = BlockStateRecordKeyReverse::new(block_number, state_key);
        let mut raw_iter: DBRawIterator = self
            .get_iter(COLUMN_BLOCK_STATE_REVERSE_RECORD, IteratorMode::Start)
            .into();
        raw_iter.seek_for_prev(key.as_slice());

        if !raw_iter.valid() {
            return None;
        }
        match raw_iter.key() {
            Some(prev_key) => {
                // not a some key
                if &prev_key[..32] != key.state_key().as_slice() {
                    return None;
                }

                // get old value
                let prev_reverse_key = BlockStateRecordKeyReverse::from_slice(prev_key);
                let prev_key = BlockStateRecordKey::new(
                    prev_reverse_key.block_number(),
                    &prev_reverse_key.state_key(),
                );

                self.get(COLUMN_BLOCK_STATE_RECORD, prev_key.as_slice())
                    .map(|raw| {
                        let mut buf = [0u8; 32];
                        buf.copy_from_slice(&raw);
                        buf.into()
                    })
            }
            _ => None,
        }
    }

    pub fn iter_block_state_record(
        &self,
        block_number: u64,
    ) -> impl Iterator<Item = BlockStateRecordKey> + '_ {
        let start_key = BlockStateRecordKey::new(block_number, &H256::zero());
        self.get_iter(
            COLUMN_BLOCK_STATE_RECORD,
            IteratorMode::From(start_key.as_slice(), Direction::Forward),
        )
        .map(|(key, _value)| BlockStateRecordKey::from_slice(&key))
        .take_while(move |key| key.block_number() == block_number)
    }

    // FIXME: This method may running into inconsistent state if current state is dirty.
    // We should separate the StateDB into ReadOnly & WriteOnly,
    // The ReadOnly is for fetching history state, and the write only is for writing new state.
    // This function should only be added on the ReadOnly state.
    pub fn account_smt(&self) -> Result<SMT<SMTStore<'_, Self>>, Error> {
        let block = self.get_last_valid_tip_block()?;
        let merkle_state = block.raw().post_account();
        self.account_smt_with_merkle_state(merkle_state)
    }

    pub fn account_smt_store(&self) -> Result<SMTStore<Self>, Error> {
        let smt_store = SMTStore::new(COLUMN_ACCOUNT_SMT_LEAF, COLUMN_ACCOUNT_SMT_BRANCH, self);
        Ok(smt_store)
    }

    pub fn account_smt_with_merkle_state(
        &self,
        merkle_state: AccountMerkleState,
    ) -> Result<SMT<SMTStore<Self>>, Error> {
        let smt_store = self.account_smt_store()?;
        Ok(SMT::new(merkle_state.merkle_root().unpack(), smt_store))
    }

    // get state tree
    pub fn state_tree(&self, context: StateContext) -> Result<StateTree, Error> {
        let block = match context {
            StateContext::ReadOnlyHistory(block_number) => {
                let block_hash = self
                    .get_block_hash_by_number(block_number)?
                    .ok_or_else(|| Error::from("can't find block".to_string()))?;
                self.get_block(&block_hash)?
                    .ok_or_else(|| "can't find block".to_string())?
            }
            _ => self.get_last_valid_tip_block()?,
        };
        let merkle_state = block.raw().post_account();
        let account_count = merkle_state.count().unpack();
        let tree = self.account_smt_with_merkle_state(merkle_state)?;
        Ok(StateTree::new(tree, account_count, context))
    }

    pub fn insert_mem_pool_transaction(
        &self,
        tx_hash: &H256,
        tx: packed::L2Transaction,
    ) -> Result<(), Error> {
        self.insert_raw(
            COLUMN_MEM_POOL_TRANSACTION,
            tx_hash.as_slice(),
            tx.as_slice(),
        )
    }

    pub fn remove_mem_pool_transaction(&self, tx_hash: &H256) -> Result<(), Error> {
        self.delete(COLUMN_MEM_POOL_TRANSACTION, tx_hash.as_slice())?;
        self.delete(COLUMN_MEM_POOL_TRANSACTION_RECEIPT, tx_hash.as_slice())?;
        Ok(())
    }

    pub fn insert_mem_pool_transaction_receipt(
        &self,
        tx_hash: &H256,
        tx_receipt: packed::TxReceipt,
    ) -> Result<(), Error> {
        self.insert_raw(
            COLUMN_MEM_POOL_TRANSACTION_RECEIPT,
            tx_hash.as_slice(),
            tx_receipt.as_slice(),
        )
    }

    pub fn insert_mem_pool_withdrawal(
        &self,
        withdrawal_hash: &H256,
        withdrawal: packed::WithdrawalRequestExtra,
    ) -> Result<(), Error> {
        self.insert_raw(
            COLUMN_MEM_POOL_WITHDRAWAL,
            withdrawal_hash.as_slice(),
            withdrawal.as_slice(),
        )
    }

    pub fn remove_mem_pool_withdrawal(&self, withdrawal_hash: &H256) -> Result<(), Error> {
        self.delete(COLUMN_MEM_POOL_WITHDRAWAL, withdrawal_hash.as_slice())?;
        Ok(())
    }

    pub fn record_block_state(
        &self,
        block_number: u64,
        state_key: H256,
        value: H256,
    ) -> Result<(), Error> {
        // insert block state key value
        let key = BlockStateRecordKey::new(block_number, &state_key);
        self.insert_raw(COLUMN_BLOCK_STATE_RECORD, key.as_slice(), value.as_slice())?;
        // insert block state reverse key
        let reverse_key = BlockStateRecordKeyReverse::new(block_number, &state_key);
        self.insert_raw(
            COLUMN_BLOCK_STATE_REVERSE_RECORD,
            reverse_key.as_slice(),
            &[],
        )
    }
}

impl ChainStore for StoreTransaction {}
