use crate::test_utils::{DEFAULT_BUILDER_PAYLOAD_VALUE_WEI, DEFAULT_JWT_SECRET};
use crate::{Config, ExecutionLayer, PayloadAttributes};
use async_trait::async_trait;
use eth2::types::{BlobsBundle, BlockId, StateId, ValidatorId};
use eth2::{BeaconNodeHttpClient, Timeouts};
pub use ethereum_consensus::state_transition::Context;
use ethereum_consensus::{
    crypto::{SecretKey, Signature},
    primitives::{BlsPublicKey, BlsSignature, ExecutionAddress, Hash32, Root, U256},
    state_transition::Error,
};
use fork_choice::ForkchoiceUpdateParameters;
use mev_rs::{
    blinded_block_provider::Server as BlindedBlockProviderServer,
    signing::{sign_builder_message, verify_signed_builder_message},
    types::{
        bellatrix::{
            BuilderBid as BuilderBidBellatrix, SignedBuilderBid as SignedBuilderBidBellatrix,
        },
        capella::{BuilderBid as BuilderBidCapella, SignedBuilderBid as SignedBuilderBidCapella},
        deneb::{BuilderBid as BuilderBidDeneb, SignedBuilderBid as SignedBuilderBidDeneb},
        BidRequest, BuilderBid, ExecutionPayload as ServerPayload, SignedBlindedBeaconBlock,
        SignedBuilderBid, SignedValidatorRegistration,
    },
    Error as MevError,
};
use parking_lot::RwLock;
use sensitive_url::SensitiveUrl;
use ssz::{Decode, Encode};
use ssz_rs::{Merkleized, SimpleSerialize};
use std::collections::HashMap;
use std::fmt::Debug;
use std::net::Ipv4Addr;
use std::sync::Arc;
use std::time::Duration;
use task_executor::TaskExecutor;
use tempfile::NamedTempFile;
use tree_hash::TreeHash;
use types::builder_bid::BlindedBlobsBundle;
use types::{
    Address, BeaconState, ChainSpec, EthSpec, ExecPayload, ExecutionPayload,
    ExecutionPayloadHeader, ForkName, ForkVersionedResponse, Hash256, Slot, Uint256,
};

pub type MockBuilderServer = axum::Server<
    hyper::server::conn::AddrIncoming,
    axum::routing::IntoMakeService<axum::routing::Router>,
>;

#[derive(Clone)]
pub enum Operation {
    FeeRecipient(Address),
    GasLimit(usize),
    Value(Uint256),
    ParentHash(Hash256),
    PrevRandao(Hash256),
    BlockNumber(usize),
    Timestamp(usize),
    WithdrawalsRoot(Hash256),
}

impl Operation {
    fn apply<B: BidStuff>(self, bid: &mut B) -> Result<(), MevError> {
        match self {
            Operation::FeeRecipient(fee_recipient) => {
                *bid.fee_recipient_mut() = to_ssz_rs(&fee_recipient)?
            }
            Operation::GasLimit(gas_limit) => *bid.gas_limit_mut() = gas_limit as u64,
            Operation::Value(value) => *bid.value_mut() = to_ssz_rs(&value)?,
            Operation::ParentHash(parent_hash) => *bid.parent_hash_mut() = to_ssz_rs(&parent_hash)?,
            Operation::PrevRandao(prev_randao) => *bid.prev_randao_mut() = to_ssz_rs(&prev_randao)?,
            Operation::BlockNumber(block_number) => *bid.block_number_mut() = block_number as u64,
            Operation::Timestamp(timestamp) => *bid.timestamp_mut() = timestamp as u64,
            Operation::WithdrawalsRoot(root) => *bid.withdrawals_root_mut()? = to_ssz_rs(&root)?,
        }
        Ok(())
    }
}

// contains functions we need for BuilderBids.. not sure what to call this
pub trait BidStuff {
    fn fee_recipient_mut(&mut self) -> &mut ExecutionAddress;
    fn gas_limit_mut(&mut self) -> &mut u64;
    fn value_mut(&mut self) -> &mut U256;
    fn parent_hash_mut(&mut self) -> &mut Hash32;
    fn prev_randao_mut(&mut self) -> &mut Hash32;
    fn block_number_mut(&mut self) -> &mut u64;
    fn timestamp_mut(&mut self) -> &mut u64;
    fn withdrawals_root_mut(&mut self) -> Result<&mut Root, MevError>;

    fn sign_builder_message(
        &mut self,
        signing_key: &SecretKey,
        context: &Context,
    ) -> Result<BlsSignature, Error>;

    fn to_signed_bid(self, signature: BlsSignature) -> SignedBuilderBid;
}

macro_rules! map_builder_bid {
    ($self_ident:ident, |$var:ident| $expr:expr) => {
        match $self_ident {
            BuilderBid::Bellatrix($var) => $expr,
            BuilderBid::Capella($var) => $expr,
            BuilderBid::Deneb($var) => $expr,
        }
    };
}

impl BidStuff for BuilderBid {
    fn fee_recipient_mut(&mut self) -> &mut ExecutionAddress {
        map_builder_bid!(self, |bid| &mut bid.header.fee_recipient)
    }

    fn gas_limit_mut(&mut self) -> &mut u64 {
        map_builder_bid!(self, |bid| &mut bid.header.gas_limit)
    }

    fn value_mut(&mut self) -> &mut U256 {
        map_builder_bid!(self, |bid| &mut bid.value)
    }

    fn parent_hash_mut(&mut self) -> &mut Hash32 {
        map_builder_bid!(self, |bid| &mut bid.header.parent_hash)
    }

    fn prev_randao_mut(&mut self) -> &mut Hash32 {
        map_builder_bid!(self, |bid| &mut bid.header.prev_randao)
    }

    fn block_number_mut(&mut self) -> &mut u64 {
        map_builder_bid!(self, |bid| &mut bid.header.block_number)
    }

    fn timestamp_mut(&mut self) -> &mut u64 {
        map_builder_bid!(self, |bid| &mut bid.header.timestamp)
    }

    fn withdrawals_root_mut(&mut self) -> Result<&mut Root, MevError> {
        match self {
            Self::Bellatrix(_) => Err(MevError::InvalidFork),
            Self::Capella(bid) => Ok(&mut bid.header.withdrawals_root),
            Self::Deneb(bid) => Ok(&mut bid.header.withdrawals_root),
        }
    }

    fn sign_builder_message(
        &mut self,
        signing_key: &SecretKey,
        context: &Context,
    ) -> Result<Signature, Error> {
        map_builder_bid!(self, |message| sign_builder_message(
            message,
            signing_key,
            context
        ))
    }

    fn to_signed_bid(self, signature: Signature) -> SignedBuilderBid {
        match self {
            Self::Bellatrix(message) => {
                SignedBuilderBid::Bellatrix(SignedBuilderBidBellatrix { message, signature })
            }
            Self::Capella(message) => {
                SignedBuilderBid::Capella(SignedBuilderBidCapella { message, signature })
            }
            Self::Deneb(message) => {
                SignedBuilderBid::Deneb(SignedBuilderBidDeneb { message, signature })
            }
        }
    }
}

#[derive(Clone)]
pub struct MockBuilder<E: EthSpec> {
    el: ExecutionLayer<E>,
    beacon_client: BeaconNodeHttpClient,
    spec: ChainSpec,
    context: Arc<Context>,
    val_registration_cache: Arc<RwLock<HashMap<BlsPublicKey, SignedValidatorRegistration>>>,
    builder_sk: SecretKey,
    operations: Arc<RwLock<Vec<Operation>>>,
    invalidate_signatures: Arc<RwLock<bool>>,
}

impl<E: EthSpec> MockBuilder<E> {
    pub fn new_for_testing(
        mock_el_url: SensitiveUrl,
        beacon_url: SensitiveUrl,
        spec: ChainSpec,
        executor: TaskExecutor,
    ) -> (Self, MockBuilderServer) {
        let file = NamedTempFile::new().unwrap();
        let path = file.path().into();
        std::fs::write(&path, hex::encode(DEFAULT_JWT_SECRET)).unwrap();

        // This EL should not talk to a builder
        let config = Config {
            execution_endpoints: vec![mock_el_url],
            secret_files: vec![path],
            suggested_fee_recipient: None,
            ..Default::default()
        };

        let el =
            ExecutionLayer::from_config(config, executor.clone(), executor.log().clone()).unwrap();

        // This should probably be done for all fields, we only update ones we are testing with so far.
        let mut context = Context::for_mainnet();
        context.terminal_total_difficulty = to_ssz_rs(&spec.terminal_total_difficulty).unwrap();
        context.terminal_block_hash = to_ssz_rs(&spec.terminal_block_hash).unwrap();
        context.terminal_block_hash_activation_epoch =
            to_ssz_rs(&spec.terminal_block_hash_activation_epoch).unwrap();

        let builder = MockBuilder::new(
            el,
            BeaconNodeHttpClient::new(beacon_url, Timeouts::set_all(Duration::from_secs(1))),
            spec,
            context,
        );
        let host: Ipv4Addr = Ipv4Addr::LOCALHOST;
        let port = 0;
        let provider = BlindedBlockProviderServer::new(host, port, builder.clone());
        let server = provider.serve();
        (builder, server)
    }

    pub fn new(
        el: ExecutionLayer<E>,
        beacon_client: BeaconNodeHttpClient,
        spec: ChainSpec,
        context: Context,
    ) -> Self {
        let sk = SecretKey::random(&mut rand::thread_rng()).unwrap();
        Self {
            el,
            beacon_client,
            // Should keep spec and context consistent somehow
            spec,
            context: Arc::new(context),
            val_registration_cache: Arc::new(RwLock::new(HashMap::new())),
            builder_sk: sk,
            operations: Arc::new(RwLock::new(vec![])),
            invalidate_signatures: Arc::new(RwLock::new(false)),
        }
    }

    pub fn add_operation(&self, op: Operation) {
        // Insert operations at the front of the vec to make sure `apply_operations` applies them
        // in the order they are added.
        self.operations.write().insert(0, op);
    }

    pub fn invalid_signatures(&self) {
        *self.invalidate_signatures.write() = true;
    }

    pub fn valid_signatures(&mut self) {
        *self.invalidate_signatures.write() = false;
    }

    fn apply_operations<B: BidStuff>(&self, bid: &mut B) -> Result<(), MevError> {
        let mut guard = self.operations.write();
        while let Some(op) = guard.pop() {
            op.apply(bid)?;
        }
        Ok(())
    }

    pub fn pubkey(&self) -> ethereum_consensus::crypto::PublicKey {
        self.builder_sk.public_key()
    }
}

#[async_trait]
impl<E: EthSpec> mev_rs::BlindedBlockProvider for MockBuilder<E> {
    async fn register_validators(
        &self,
        registrations: &mut [SignedValidatorRegistration],
    ) -> Result<(), MevError> {
        for registration in registrations {
            let pubkey = registration.message.public_key.clone();
            let message = &mut registration.message;
            verify_signed_builder_message(
                message,
                &registration.signature,
                &pubkey,
                &self.context,
            )?;
            self.val_registration_cache.write().insert(
                registration.message.public_key.clone(),
                registration.clone(),
            );
        }

        Ok(())
    }

    async fn fetch_best_bid(&self, bid_request: &BidRequest) -> Result<SignedBuilderBid, MevError> {
        let slot = Slot::new(bid_request.slot);
        let fork = self.spec.fork_name_at_slot::<E>(slot);
        let signed_cached_data = self
            .val_registration_cache
            .read()
            .get(&bid_request.public_key)
            .ok_or_else(|| convert_err("missing registration"))?
            .clone();
        let cached_data = signed_cached_data.message;

        let head = self
            .beacon_client
            .get_beacon_blocks::<E>(BlockId::Head)
            .await
            .map_err(convert_err)?
            .ok_or_else(|| convert_err("missing head block"))?;

        let block = head.data.message();
        let head_block_root = block.tree_hash_root();
        let head_execution_hash = block
            .body()
            .execution_payload()
            .map_err(convert_err)?
            .block_hash();
        if head_execution_hash != from_ssz_rs(&bid_request.parent_hash)? {
            return Err(custom_err(format!(
                "head mismatch: {} {}",
                head_execution_hash, bid_request.parent_hash
            )));
        }

        let finalized_execution_hash = self
            .beacon_client
            .get_beacon_blocks::<E>(BlockId::Finalized)
            .await
            .map_err(convert_err)?
            .ok_or_else(|| convert_err("missing finalized block"))?
            .data
            .message()
            .body()
            .execution_payload()
            .map_err(convert_err)?
            .block_hash();

        let justified_execution_hash = self
            .beacon_client
            .get_beacon_blocks::<E>(BlockId::Justified)
            .await
            .map_err(convert_err)?
            .ok_or_else(|| convert_err("missing finalized block"))?
            .data
            .message()
            .body()
            .execution_payload()
            .map_err(convert_err)?
            .block_hash();

        let val_index = self
            .beacon_client
            .get_beacon_states_validator_id(
                StateId::Head,
                &ValidatorId::PublicKey(from_ssz_rs(&cached_data.public_key)?),
            )
            .await
            .map_err(convert_err)?
            .ok_or_else(|| convert_err("missing validator from state"))?
            .data
            .index;
        let fee_recipient = from_ssz_rs(&cached_data.fee_recipient)?;
        let slots_since_genesis = slot.as_u64() - self.spec.genesis_slot.as_u64();

        let genesis_time = self
            .beacon_client
            .get_beacon_genesis()
            .await
            .map_err(convert_err)?
            .data
            .genesis_time;
        let timestamp = (slots_since_genesis * self.spec.seconds_per_slot) + genesis_time;

        let head_state: BeaconState<E> = self
            .beacon_client
            .get_debug_beacon_states(StateId::Head)
            .await
            .map_err(convert_err)?
            .ok_or_else(|| custom_err("missing head state".to_string()))?
            .data;
        let prev_randao = head_state
            .get_randao_mix(head_state.current_epoch())
            .map_err(convert_err)?;
        let expected_withdrawals = match fork {
            ForkName::Base | ForkName::Altair | ForkName::Merge => None,
            ForkName::Capella | ForkName::Deneb => Some(
                self.beacon_client
                    .get_expected_withdrawals(&StateId::Head)
                    .await
                    .unwrap()
                    .data,
            ),
        };

        let payload_attributes = match fork {
            // the withdrawals root is filled in by operations, but we supply the valid withdrawals
            // first to avoid polluting the execution block generator with invalid payload attributes
            // NOTE: this was part of an effort to add payload attribute uniqueness checks,
            // which was abandoned because it broke too many tests in subtle ways.
            ForkName::Merge | ForkName::Capella => PayloadAttributes::new(
                timestamp,
                *prev_randao,
                fee_recipient,
                expected_withdrawals,
                None,
            ),
            ForkName::Deneb => PayloadAttributes::new(
                timestamp,
                *prev_randao,
                fee_recipient,
                expected_withdrawals,
                Some(head_block_root),
            ),
            ForkName::Base | ForkName::Altair => {
                return Err(MevError::InvalidFork);
            }
        };

        self.el
            .insert_proposer(slot, head_block_root, val_index, payload_attributes.clone())
            .await;

        let forkchoice_update_params = ForkchoiceUpdateParameters {
            head_root: Hash256::zero(),
            head_hash: None,
            justified_hash: Some(justified_execution_hash),
            finalized_hash: Some(finalized_execution_hash),
        };

        let (payload, _block_value, maybe_blobs_bundle): (
            ExecutionPayload<E>,
            Uint256,
            Option<BlobsBundle<E>>,
        ) = self
            .el
            .get_full_payload_caching(
                head_execution_hash,
                &payload_attributes,
                forkchoice_update_params,
                fork,
            )
            .await
            .map_err(convert_err)?
            .into();

        let header = match payload {
            ExecutionPayload::Merge(payload) => ExecutionPayloadHeader::Merge((&payload).into()),
            ExecutionPayload::Capella(payload) => {
                ExecutionPayloadHeader::Capella((&payload).into())
            }
            ExecutionPayload::Deneb(payload) => ExecutionPayloadHeader::Deneb((&payload).into()),
        };

        let mut message = match fork {
            ForkName::Deneb => {
                let blinded_blobs: BlindedBlobsBundle<E> =
                    maybe_blobs_bundle.map(Into::into).unwrap_or_default();
                BuilderBid::Deneb(BuilderBidDeneb {
                    header: to_ssz_rs(&header)?,
                    blinded_blobs_bundle: to_ssz_rs(&blinded_blobs)?,
                    value: to_ssz_rs(&Uint256::from(DEFAULT_BUILDER_PAYLOAD_VALUE_WEI))?,
                    public_key: self.builder_sk.public_key(),
                })
            }
            ForkName::Capella => BuilderBid::Capella(BuilderBidCapella {
                header: to_ssz_rs(&header)?,
                value: to_ssz_rs(&Uint256::from(DEFAULT_BUILDER_PAYLOAD_VALUE_WEI))?,
                public_key: self.builder_sk.public_key(),
            }),
            ForkName::Merge => BuilderBid::Bellatrix(BuilderBidBellatrix {
                header: to_ssz_rs(&header)?,
                value: to_ssz_rs(&Uint256::from(DEFAULT_BUILDER_PAYLOAD_VALUE_WEI))?,
                public_key: self.builder_sk.public_key(),
            }),
            ForkName::Base | ForkName::Altair => return Err(MevError::InvalidFork),
        };
        *message.gas_limit_mut() = cached_data.gas_limit;

        self.apply_operations(&mut message)?;
        let mut signature =
            message.sign_builder_message(&self.builder_sk, self.context.as_ref())?;

        if *self.invalidate_signatures.read() {
            signature = Signature::default();
        }

        Ok(message.to_signed_bid(signature))
    }

    async fn open_bid(
        &self,
        signed_block: &mut SignedBlindedBeaconBlock,
    ) -> Result<ServerPayload, MevError> {
        let node = match signed_block {
            SignedBlindedBeaconBlock::Bellatrix(block) => {
                block.message.body.execution_payload_header.hash_tree_root()
            }
            SignedBlindedBeaconBlock::Capella(block) => {
                block.message.body.execution_payload_header.hash_tree_root()
            }
            SignedBlindedBeaconBlock::Deneb(block_and_blobs) => block_and_blobs
                .signed_blinded_block
                .message
                .body
                .execution_payload_header
                .hash_tree_root(),
        }
        .map_err(convert_err)?;

        let payload = self
            .el
            .get_payload_by_root(&from_ssz_rs(&node)?)
            .ok_or_else(|| convert_err("missing payload for tx root"))?;

        let fork = payload.payload_ref().fork_name();
        let resp = ForkVersionedResponse {
            version: Some(fork),
            data: payload,
        };

        let json_payload = serde_json::to_string(&resp).map_err(convert_err)?;
        serde_json::from_str(json_payload.as_str()).map_err(convert_err)
    }
}

pub fn from_ssz_rs<T: SimpleSerialize, U: Decode>(ssz_rs_data: &T) -> Result<U, MevError> {
    U::from_ssz_bytes(
        ssz_rs::serialize(ssz_rs_data)
            .map_err(convert_err)?
            .as_ref(),
    )
    .map_err(convert_err)
}

pub fn to_ssz_rs<T: Encode, U: SimpleSerialize>(ssz_data: &T) -> Result<U, MevError> {
    ssz_rs::deserialize::<U>(&ssz_data.as_ssz_bytes()).map_err(convert_err)
}

pub fn convert_err<E: Debug>(e: E) -> MevError {
    custom_err(format!("{e:?}"))
}

// This is a bit of a hack since the `Custom` variant was removed from `mev_rs::Error`.
pub fn custom_err(s: String) -> MevError {
    MevError::Consensus(ethereum_consensus::state_transition::Error::Io(
        std::io::Error::new(std::io::ErrorKind::Other, s),
    ))
}