use crate::{db::Database, RollupContract, CHAIN_ID, CHAIN_SPEC};
use alloy_consensus::{Blob, SidecarCoder, SimpleCoder};
use alloy_rlp::Decodable as _;
use eyre::OptionExt;
use reth::transaction_pool::TransactionPool;
use reth_execution_errors::BlockValidationError;
use reth_node_api::{ConfigureEvm, ConfigureEvmEnv};
use reth_node_ethereum::EthEvmConfig;
use reth_primitives::{
    constants,
    eip4844::kzg_to_versioned_hash,
    keccak256,
    revm::env::fill_tx_env,
    revm_primitives::{CfgEnvWithHandlerCfg, EVMError, ExecutionResult, ResultAndState},
    Address, Block, BlockWithSenders, Bytes, Hardfork, Header, Receipt, TransactionSigned, TxType,
    B256, U256,
};
use reth_revm::{
    db::{states::bundle_state::BundleRetention, BundleState},
    DBBox, DatabaseCommit, Evm, StateBuilder, StateDBBox,
};
use reth_tracing::tracing::debug;

/// Execute a rollup block and return (block with recovered senders)[BlockWithSenders], (bundle
/// state)[BundleState] and list of (receipts)[Receipt].
pub async fn execute_block<Pool: TransactionPool>(
    db: &mut Database,
    pool: &Pool,
    tx: &TransactionSigned,
    header: &RollupContract::BlockHeader,
    block_data: Bytes,
    block_data_hash: B256,
) -> eyre::Result<(BlockWithSenders, BundleState, Vec<Receipt>, Vec<ExecutionResult>)> {
    if header.rollupChainId != U256::from(CHAIN_ID) {
        eyre::bail!("Invalid rollup chain ID")
    }

    // Construct header
    let header = construct_header(db, header)?;

    // Decode transactions
    let transactions = decode_transactions(pool, tx, block_data, block_data_hash).await?;

    // Configure EVM
    let evm_config = EthEvmConfig::default();
    let mut evm = configure_evm(&evm_config, db, &header);

    // Execute transactions
    let (executed_txs, receipts, results) = execute_transactions(&mut evm, &header, transactions)?;

    // Construct block and recover senders
    let block = Block { header, body: executed_txs, ..Default::default() }
        .with_recovered_senders()
        .ok_or_eyre("failed to recover senders")?;

    let bundle = evm.db_mut().take_bundle();

    Ok((block, bundle, receipts, results))
}

/// Construct header from the given rollup header.
fn construct_header(db: &Database, header: &RollupContract::BlockHeader) -> eyre::Result<Header> {
    let parent_block = if !header.sequence.is_zero() {
        db.get_block(header.sequence - U256::from(1))?
    } else {
        None
    };

    let block_number = u64::try_from(header.sequence)?;

    // Calculate base fee per gas for EIP-1559 transactions
    let base_fee_per_gas = if CHAIN_SPEC.fork(Hardfork::London).transitions_at_block(block_number) {
        constants::EIP1559_INITIAL_BASE_FEE
    } else {
        parent_block
            .as_ref()
            .ok_or(eyre::eyre!("parent block not found"))?
            .header
            .next_block_base_fee(CHAIN_SPEC.base_fee_params_at_block(block_number))
            .ok_or(eyre::eyre!("failed to calculate base fee"))?
    };

    // Construct header
    Ok(Header {
        parent_hash: parent_block.map(|block| block.header.hash()).unwrap_or_default(),
        number: block_number,
        gas_limit: u64::try_from(header.gasLimit)?,
        timestamp: u64::try_from(header.confirmBy)?,
        base_fee_per_gas: Some(base_fee_per_gas),
        ..Default::default()
    })
}

/// Configure EVM with the given database and header.
fn configure_evm<'a>(
    config: &'a EthEvmConfig,
    db: &'a mut Database,
    header: &Header,
) -> Evm<'a, (), StateDBBox<'a, eyre::Report>> {
    let mut evm = config.evm(
        StateBuilder::new_with_database(Box::new(db) as DBBox<'_, eyre::Report>)
            .with_bundle_update()
            .build(),
    );
    evm.db_mut().set_state_clear_flag(
        CHAIN_SPEC.fork(Hardfork::SpuriousDragon).active_at_block(header.number),
    );

    let mut cfg = CfgEnvWithHandlerCfg::new_with_spec_id(evm.cfg().clone(), evm.spec_id());
    EthEvmConfig::fill_cfg_and_block_env(
        &mut cfg,
        evm.block_mut(),
        &CHAIN_SPEC,
        header,
        U256::ZERO,
    );
    *evm.cfg_mut() = cfg.cfg_env;

    evm
}

/// Decode transactions from the block data and recover senders.
/// - If the transaction is a blob-carrying one, decode the blobs either using the local transaction
///   pool, or querying Blobscan.
/// - If the transaction is a regular one, decode the block data directly.
async fn decode_transactions<Pool: TransactionPool>(
    pool: &Pool,
    tx: &TransactionSigned,
    block_data: Bytes,
    block_data_hash: B256,
) -> eyre::Result<Vec<(TransactionSigned, Address)>> {
    // Get raw transactions either from the blobs, or directly from the block data
    let raw_transactions = if matches!(tx.tx_type(), TxType::Eip4844) {
        let blobs: Vec<_> = if let Some(sidecar) = pool.get_blob(tx.hash)? {
            // Try to get blobs from the transaction pool
            sidecar.blobs.into_iter().zip(sidecar.commitments).collect()
        } else {
            // If transaction is not found in the pool, try to get blobs from Blobscan
            let blobscan_client = foundry_blob_explorers::Client::holesky();
            let sidecar = blobscan_client.transaction(tx.hash).await?.blob_sidecar();
            sidecar
                .blobs
                .into_iter()
                .map(|blob| (*blob).into())
                .zip(sidecar.commitments.into_iter().map(|commitment| (*commitment).into()))
                .collect()
        };

        // Decode blob hashes from block data
        let blob_hashes = Vec::<B256>::decode(&mut block_data.as_ref())?;

        // Filter blobs that are present in the block data
        let blobs = blobs
            .into_iter()
            // Convert blob KZG commitments to versioned hashes
            .map(|(blob, commitment)| (blob, kzg_to_versioned_hash(commitment.as_slice())))
            // Filter only blobs that are present in the block data
            .filter(|(_, hash)| blob_hashes.contains(hash))
            .map(|(blob, _)| Blob::from(*blob))
            .collect::<Vec<_>>();
        if blobs.len() != blob_hashes.len() {
            eyre::bail!("some blobs not found")
        }

        // Decode blobs and concatenate them to get the raw transactions
        let data = SimpleCoder::default()
            .decode_all(&blobs)
            .ok_or(eyre::eyre!("failed to decode blobs"))?
            .concat();

        data.into()
    } else {
        block_data
    };

    let raw_transaction_hash = keccak256(&raw_transactions);
    if raw_transaction_hash != block_data_hash {
        eyre::bail!("block data hash mismatch")
    }

    // Decode block data, filter only transactions with the correct chain ID and recover senders
    let transactions = Vec::<TransactionSigned>::decode(&mut raw_transactions.as_ref())?
        .into_iter()
        .filter(|tx| tx.chain_id() == Some(CHAIN_ID))
        .map(|tx| {
            let sender = tx.recover_signer().ok_or(eyre::eyre!("failed to recover signer"))?;
            Ok((tx, sender))
        })
        .collect::<eyre::Result<_>>()?;

    Ok(transactions)
}

/// Execute transactions and return the list of executed transactions, receipts and
/// execution results.
fn execute_transactions(
    evm: &mut Evm<'_, (), StateDBBox<'_, eyre::Report>>,
    header: &Header,
    transactions: Vec<(TransactionSigned, Address)>,
) -> eyre::Result<(Vec<TransactionSigned>, Vec<Receipt>, Vec<ExecutionResult>)> {
    let mut receipts = Vec::with_capacity(transactions.len());
    let mut executed_txs = Vec::with_capacity(transactions.len());
    let mut results = Vec::with_capacity(transactions.len());
    if !transactions.is_empty() {
        let mut cumulative_gas_used = 0;
        for (transaction, sender) in transactions {
            // The sum of the transaction’s gas limit, Tg, and the gas utilized in this block prior,
            // must be no greater than the block’s gasLimit.
            let block_available_gas = header.gas_limit - cumulative_gas_used;
            if transaction.gas_limit() > block_available_gas {
                return Err(BlockValidationError::TransactionGasLimitMoreThanAvailableBlockGas {
                    transaction_gas_limit: transaction.gas_limit(),
                    block_available_gas,
                }
                .into())
            }
            // Execute transaction.
            // Fill revm structure.
            fill_tx_env(evm.tx_mut(), &transaction, sender);

            let ResultAndState { result, state } = match evm.transact() {
                Ok(result) => result,
                Err(err) => {
                    match err {
                        EVMError::Transaction(err) => {
                            // if the transaction is invalid, we can skip it
                            debug!(%err, ?transaction, "Skipping invalid transaction");
                            continue
                        }
                        err => {
                            // this is an error that we should treat as fatal for this attempt
                            eyre::bail!(err)
                        }
                    }
                }
            };

            debug!(?transaction, ?result, ?state, "Executed transaction");

            evm.db_mut().commit(state);

            // append gas used
            cumulative_gas_used += result.gas_used();

            // Push transaction changeset and calculate header bloom filter for receipt.
            #[allow(clippy::needless_update)] // side-effect of optimism fields
            receipts.push(Receipt {
                tx_type: transaction.tx_type(),
                success: result.is_success(),
                cumulative_gas_used,
                logs: result.logs().iter().cloned().map(Into::into).collect(),
                ..Default::default()
            });

            // append transaction to the list of executed transactions
            executed_txs.push(transaction);
            results.push(result);
        }

        evm.db_mut().merge_transitions(BundleRetention::Reverts);
    }

    Ok((executed_txs, receipts, results))
}

#[cfg(test)]
mod tests {
    use std::time::{SystemTime, UNIX_EPOCH};

    use alloy_consensus::{SidecarBuilder, SimpleCoder};
    use alloy_sol_types::{sol, SolCall};
    use reth::transaction_pool::{
        test_utils::{testing_pool, MockTransaction},
        TransactionOrigin, TransactionPool,
    };
    use reth_primitives::{
        bytes,
        constants::ETH_TO_WEI,
        keccak256, public_key_to_address,
        revm_primitives::{AccountInfo, ExecutionResult, Output, TransactTo, TxEnv},
        BlockNumber, Receipt, SealedBlockWithSenders, Transaction, TxEip2930, TxKind, U256,
    };
    use reth_revm::Evm;
    use reth_testing_utils::generators::{self, sign_tx_with_key_pair};
    use rusqlite::Connection;
    use secp256k1::{Keypair, Secp256k1};

    use crate::{
        db::Database, execute_block, RollupContract::BlockHeader, CHAIN_ID,
        ROLLUP_SUBMITTER_ADDRESS,
    };

    sol!(
        WETH,
        r#"
[
   {
      "constant":true,
      "inputs":[
         {
            "name":"",
            "type":"address"
         }
      ],
      "name":"balanceOf",
      "outputs":[
         {
            "name":"",
            "type":"uint256"
         }
      ],
      "payable":false,
      "stateMutability":"view",
      "type":"function"
   }
]
        "#
    );

    #[tokio::test]
    async fn test_execute_block() -> eyre::Result<()> {
        reth_tracing::init_test_tracing();

        let mut database = Database::new(Connection::open_in_memory()?)?;

        // Create key pair
        let secp = Secp256k1::new();
        let key_pair = Keypair::new(&secp, &mut generators::rng());
        let sender_address = public_key_to_address(key_pair.public_key());

        // Deposit some ETH to the sender and insert it into database
        database.upsert_account(sender_address, |_| {
            Ok(AccountInfo { balance: U256::from(ETH_TO_WEI), nonce: 1, ..Default::default() })
        })?;

        // WETH deployment transaction sent using calldata
        let (_, _, results) = execute_transaction(
            &mut database,
            key_pair,
            0,
            Transaction::Eip2930(TxEip2930 {
                chain_id: CHAIN_ID,
                nonce: 1,
                gas_limit: 1_500_000,
                gas_price: 1_500_000_000,
                to: TxKind::Create,
                // WETH9 bytecode
                input: bytes!(" "),
                ..Default::default()
            }),
            BlockDataSource::Calldata
        ).await?;

        let weth_address = match results.first() {
            Some(ExecutionResult::Success { output: Output::Create(_, Some(address)), .. }) => {
                *address
            }
            _ => eyre::bail!("WETH contract address not found"),
        };

        // WETH deposit transaction sent using blobs
        execute_transaction(
            &mut database,
            key_pair,
            1,
            Transaction::Eip2930(TxEip2930 {
                chain_id: CHAIN_ID,
                nonce: 2,
                gas_limit: 50000,
                gas_price: 1_500_000_000,
                to: TxKind::Call(weth_address),
                value: U256::from(0.5 * ETH_TO_WEI as f64),
                input: bytes!("d0e30db0"),
                ..Default::default()
            }),
            BlockDataSource::Blobs,
        )
        .await?;

        // Verify WETH balance
        let mut evm = Evm::builder()
            .with_db(&mut database)
            .with_tx_env(TxEnv {
                caller: sender_address,
                gas_limit: 50_000_000,
                transact_to: TransactTo::Call(weth_address),
                data: WETH::balanceOfCall::new((sender_address,)).abi_encode().into(),
                ..Default::default()
            })
            .build();
        let result = evm.transact().map_err(|err| eyre::eyre!(err))?.result;
        assert_eq!(
            result.output(),
            Some(&U256::from(0.5 * ETH_TO_WEI as f64).to_be_bytes_vec().into())
        );
        drop(evm);

        // Verify nonce
        let account = database.get_account(sender_address)?.unwrap();
        assert_eq!(account.nonce, 3);

        // Revert block with WETH deposit transaction
        database.revert_tip_block(U256::from(1))?;

        // Verify WETH balance after revert
        let mut evm = Evm::builder()
            .with_db(&mut database)
            .with_tx_env(TxEnv {
                caller: sender_address,
                gas_limit: 50_000_000,
                transact_to: TransactTo::Call(weth_address),
                data: WETH::balanceOfCall::new((sender_address,)).abi_encode().into(),
                ..Default::default()
            })
            .build();
        let result = evm.transact().map_err(|err| eyre::eyre!(err))?.result;
        assert_eq!(result.output(), Some(&U256::ZERO.to_be_bytes_vec().into()));
        drop(evm);

        // Verify nonce after revert
        let account = database.get_account(sender_address)?.unwrap();
        assert_eq!(account.nonce, 2);

        Ok(())
    }

    enum BlockDataSource {
        Calldata,
        Blobs,
    }

    async fn execute_transaction(
        database: &mut Database,
        key_pair: Keypair,
        sequence: BlockNumber,
        tx: Transaction,
        block_data_source: BlockDataSource,
    ) -> eyre::Result<(SealedBlockWithSenders, Vec<Receipt>, Vec<ExecutionResult>)> {
        // Construct block header
        let block_header = BlockHeader {
            rollupChainId: U256::from(CHAIN_ID),
            sequence: U256::from(sequence),
            confirmBy: U256::from(SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs()),
            gasLimit: U256::from(30_000_000),
            rewardAddress: ROLLUP_SUBMITTER_ADDRESS,
        };
        let encoded_transactions =
            alloy_rlp::encode(vec![sign_tx_with_key_pair(key_pair, tx).envelope_encoded()]);
        let block_data_hash = keccak256(&encoded_transactions);

        let pool = testing_pool();

        let (block_data, l1_transaction) = match block_data_source {
            BlockDataSource::Calldata => (
                encoded_transactions,
                sign_tx_with_key_pair(key_pair, Transaction::Eip2930(TxEip2930::default())),
            ),
            BlockDataSource::Blobs => {
                let sidecar =
                    SidecarBuilder::<SimpleCoder>::from_slice(&encoded_transactions).build()?;
                let blob_hashes = alloy_rlp::encode(sidecar.versioned_hashes().collect::<Vec<_>>());

                let mut mock_transaction = MockTransaction::eip4844_with_sidecar(sidecar);
                let transaction =
                    sign_tx_with_key_pair(key_pair, Transaction::from(mock_transaction.clone()));
                mock_transaction.set_hash(transaction.hash);
                pool.add_transaction(TransactionOrigin::Local, mock_transaction).await?;
                (blob_hashes, transaction)
            }
        };

        // Execute block and insert into database
        let (block, bundle, receipts, results) = execute_block(
            database,
            &pool,
            &l1_transaction,
            &block_header,
            block_data.into(),
            block_data_hash,
        )
        .await?;
        let block = block.seal_slow();
        database.insert_block_with_bundle(&block, bundle)?;

        Ok((block, receipts, results))
    }
}
