package blocks

import (
	"bytes"

	"github.com/pkg/errors"
	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
	"github.com/prysmaticlabs/prysm/beacon-chain/core/time"
	"github.com/prysmaticlabs/prysm/beacon-chain/state"
	"github.com/prysmaticlabs/prysm/consensus-types/forks/bellatrix"
	"github.com/prysmaticlabs/prysm/consensus-types/interfaces"
	"github.com/prysmaticlabs/prysm/consensus-types/wrapper"
	"github.com/prysmaticlabs/prysm/encoding/bytesutil"
	enginev1 "github.com/prysmaticlabs/prysm/proto/engine/v1"
	"github.com/prysmaticlabs/prysm/runtime/version"
	"github.com/prysmaticlabs/prysm/time/slots"
)

var (
	ErrInvalidPayloadBlockHash  = errors.New("invalid payload block hash")
	ErrInvalidPayloadTimeStamp  = errors.New("invalid payload timestamp")
	ErrInvalidPayloadPrevRandao = errors.New("invalid payload previous randao")
)

// IsMergeTransitionComplete returns true if the transition to Bellatrix has completed.
// Meaning the payload header in beacon state is not `ExecutionPayloadHeader()` (i.e. not empty).
//
// Spec code:
// def is_merge_transition_complete(state: BeaconState) -> bool:
//    return state.latest_execution_payload_header != ExecutionPayloadHeader()
func IsMergeTransitionComplete(st state.BeaconState) (bool, error) {
	if st == nil {
		return false, errors.New("nil state")
	}
	if IsPreBellatrixVersion(st.Version()) {
		return false, nil
	}
	h, err := st.LatestExecutionPayloadHeader()
	if err != nil {
		return false, err
	}
	return !bellatrix.IsEmptyHeader(h), nil
}

// IsMergeTransitionBlockUsingPreStatePayloadHeader returns true if the input block is the terminal merge block.
// Terminal merge block must be associated with an empty payload header.
// This assumes the header `h` is referenced as the parent state for block body `body.
func IsMergeTransitionBlockUsingPreStatePayloadHeader(h *enginev1.ExecutionPayloadHeader, body interfaces.BeaconBlockBody) (bool, error) {
	if h == nil || body == nil {
		return false, errors.New("nil header or block body")
	}
	if !bellatrix.IsEmptyHeader(h) {
		return false, nil
	}
	return IsExecutionBlock(body)
}

// IsExecutionBlock returns whether the block has a non-empty ExecutionPayload.
//
// Spec code:
// def is_execution_block(block: BeaconBlock) -> bool:
//     return block.body.execution_payload != ExecutionPayload()
func IsExecutionBlock(body interfaces.BeaconBlockBody) (bool, error) {
	if body == nil {
		return false, errors.New("nil block body")
	}
	payload, err := body.ExecutionPayload()
	switch {
	case errors.Is(err, wrapper.ErrUnsupportedField):
		return false, nil
	case err != nil:
		return false, err
	default:
	}
	return !bellatrix.IsEmptyPayload(payload), nil
}

// IsExecutionEnabled returns true if the beacon chain can begin executing.
// Meaning the payload header is beacon state is non-empty or the payload in block body is non-empty.
//
// Spec code:
// def is_execution_enabled(state: BeaconState, body: BeaconBlockBody) -> bool:
//    return is_merge_block(state, body) or is_merge_complete(state)
func IsExecutionEnabled(st state.BeaconState, body interfaces.BeaconBlockBody) (bool, error) {
	if st == nil || body == nil {
		return false, errors.New("nil state or block body")
	}
	if IsPreBellatrixVersion(st.Version()) {
		return false, nil
	}
	header, err := st.LatestExecutionPayloadHeader()
	if err != nil {
		return false, err
	}
	return IsExecutionEnabledUsingHeader(header, body)
}

// IsExecutionEnabledUsingHeader returns true if the execution is enabled using post processed payload header and block body.
// This is an optimized version of IsExecutionEnabled where beacon state is not required as an argument.
func IsExecutionEnabledUsingHeader(header *enginev1.ExecutionPayloadHeader, body interfaces.BeaconBlockBody) (bool, error) {
	if !bellatrix.IsEmptyHeader(header) {
		return true, nil
	}
	return IsExecutionBlock(body)
}

// IsPreBellatrixVersion returns true if input version is before bellatrix fork.
func IsPreBellatrixVersion(v int) bool {
	return v < version.Bellatrix
}

// ValidatePayloadWhenMergeCompletes validates if payload is valid versus input beacon state.
// These validation steps ONLY apply to post merge.
//
// Spec code:
//    # Verify consistency of the parent hash with respect to the previous execution payload header
//    if is_merge_complete(state):
//        assert payload.parent_hash == state.latest_execution_payload_header.block_hash
func ValidatePayloadWhenMergeCompletes(st state.BeaconState, payload *enginev1.ExecutionPayload) error {
	complete, err := IsMergeTransitionComplete(st)
	if err != nil {
		return err
	}
	if !complete {
		return nil
	}

	header, err := st.LatestExecutionPayloadHeader()
	if err != nil {
		return err
	}
	if !bytes.Equal(payload.ParentHash, header.BlockHash) {
		return errors.New("incorrect block hash")
	}
	return nil
}

// ValidatePayload validates if payload is valid versus input beacon state.
// These validation steps apply to both pre merge and post merge.
//
// Spec code:
//    # Verify random
//    assert payload.random == get_randao_mix(state, get_current_epoch(state))
//    # Verify timestamp
//    assert payload.timestamp == compute_timestamp_at_slot(state, state.slot)
func ValidatePayload(st state.BeaconState, payload *enginev1.ExecutionPayload) error {
	random, err := helpers.RandaoMix(st, time.CurrentEpoch(st))
	if err != nil {
		return err
	}

	if !bytes.Equal(payload.PrevRandao, random) {
		return ErrInvalidPayloadPrevRandao
	}
	t, err := slots.ToTime(st.GenesisTime(), st.Slot())
	if err != nil {
		return err
	}
	if payload.Timestamp != uint64(t.Unix()) {
		return ErrInvalidPayloadTimeStamp
	}
	return nil
}

// ProcessPayload processes input execution payload using beacon state.
// ValidatePayloadWhenMergeCompletes validates if payload is valid versus input beacon state.
// These validation steps ONLY apply to post merge.
//
// Spec code:
// def process_execution_payload(state: BeaconState, payload: ExecutionPayload, execution_engine: ExecutionEngine) -> None:
//    # Verify consistency of the parent hash with respect to the previous execution payload header
//    if is_merge_complete(state):
//        assert payload.parent_hash == state.latest_execution_payload_header.block_hash
//    # Verify random
//    assert payload.random == get_randao_mix(state, get_current_epoch(state))
//    # Verify timestamp
//    assert payload.timestamp == compute_timestamp_at_slot(state, state.slot)
//    # Verify the execution payload is valid
//    assert execution_engine.execute_payload(payload)
//    # Cache execution payload header
//    state.latest_execution_payload_header = ExecutionPayloadHeader(
//        parent_hash=payload.parent_hash,
//        FeeRecipient=payload.FeeRecipient,
//        state_root=payload.state_root,
//        receipt_root=payload.receipt_root,
//        logs_bloom=payload.logs_bloom,
//        random=payload.random,
//        block_number=payload.block_number,
//        gas_limit=payload.gas_limit,
//        gas_used=payload.gas_used,
//        timestamp=payload.timestamp,
//        extra_data=payload.extra_data,
//        base_fee_per_gas=payload.base_fee_per_gas,
//        block_hash=payload.block_hash,
//        transactions_root=hash_tree_root(payload.transactions),
//    )
func ProcessPayload(st state.BeaconState, payload *enginev1.ExecutionPayload) (state.BeaconState, error) {
	if err := ValidatePayloadWhenMergeCompletes(st, payload); err != nil {
		return nil, err
	}

	if err := ValidatePayload(st, payload); err != nil {
		return nil, err
	}

	header, err := bellatrix.PayloadToHeader(payload)
	if err != nil {
		return nil, err
	}
	if err := st.SetLatestExecutionPayloadHeader(header); err != nil {
		return nil, err
	}
	return st, nil
}

// ValidatePayloadHeaderWhenMergeCompletes validates the payload header when the merge completes.
func ValidatePayloadHeaderWhenMergeCompletes(st state.BeaconState, header *enginev1.ExecutionPayloadHeader) error {
	// Skip validation if the state is not merge compatible.
	complete, err := IsMergeTransitionComplete(st)
	if err != nil {
		return err
	}
	if !complete {
		return nil
	}
	// Validate current header's parent hash matches state header's block hash.
	h, err := st.LatestExecutionPayloadHeader()
	if err != nil {
		return err
	}
	if !bytes.Equal(header.ParentHash, h.BlockHash) {
		return ErrInvalidPayloadBlockHash
	}
	return nil
}

// ValidatePayloadHeader validates the payload header.
func ValidatePayloadHeader(st state.BeaconState, header *enginev1.ExecutionPayloadHeader) error {
	// Validate header's random mix matches with state in current epoch
	random, err := helpers.RandaoMix(st, time.CurrentEpoch(st))
	if err != nil {
		return err
	}
	if !bytes.Equal(header.PrevRandao, random) {
		return ErrInvalidPayloadPrevRandao
	}

	// Validate header's timestamp matches with state in current slot.
	t, err := slots.ToTime(st.GenesisTime(), st.Slot())
	if err != nil {
		return err
	}
	if header.Timestamp != uint64(t.Unix()) {
		return ErrInvalidPayloadTimeStamp
	}
	return nil
}

// ProcessPayloadHeader processes the payload header.
func ProcessPayloadHeader(st state.BeaconState, header *enginev1.ExecutionPayloadHeader) (state.BeaconState, error) {
	if err := ValidatePayloadHeaderWhenMergeCompletes(st, header); err != nil {
		return nil, err
	}

	if err := ValidatePayloadHeader(st, header); err != nil {
		return nil, err
	}

	if err := st.SetLatestExecutionPayloadHeader(header); err != nil {
		return nil, err
	}
	return st, nil
}

// GetBlockPayloadHash returns the hash of the execution payload of the block
func GetBlockPayloadHash(blk interfaces.BeaconBlock) ([32]byte, error) {
	payloadHash := [32]byte{}
	if IsPreBellatrixVersion(blk.Version()) {
		return payloadHash, nil
	}
	payload, err := blk.Body().ExecutionPayload()
	if err != nil {
		return payloadHash, err
	}
	return bytesutil.ToBytes32(payload.BlockHash), nil
}