// Copyright 2024 The Erigon Authors
// This file is part of Erigon.
//
// Erigon is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Erigon is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with Erigon. If not, see <http://www.gnu.org/licenses/>.

package snapshot_format

import (
	"encoding/binary"
	"io"

	"github.com/erigontech/erigon/cl/clparams"
	"github.com/erigontech/erigon/cl/cltypes"
	"github.com/erigontech/erigon/cl/cltypes/solid"
	"github.com/erigontech/erigon/common"
	"github.com/erigontech/erigon/common/pool"
	"github.com/erigontech/erigon/execution/types"
)

type ExecutionBlockReaderByNumber interface {
	Transactions(number uint64, hash common.Hash) (*solid.TransactionsSSZ, error)
	Withdrawals(number uint64, hash common.Hash) (*solid.ListSSZ[*cltypes.Withdrawal], error)
	SetBeaconChainConfig(beaconCfg *clparams.BeaconChainConfig)
	// CacheBody stores a recently produced block body for immediate retrieval.
	// Implementations that don't support caching should no-op.
	CacheBody(blockNumber uint64, transactions [][]byte, withdrawals []*types.Withdrawal)
}

// WriteBlockForSnapshot writes a block to the given writer in the format expected by the snapshot.
// buf is just a reusable buffer. if it had to grow it will be returned back as grown.
//
// Format: [1 byte version][32 bytes bodyRoot][8 bytes length][N bytes SSZ payload]
//
// For pre-GLOAS blocks the SSZ payload is a SignedBlindedBeaconBlock (ExecutionPayload stripped
// to save space; transactions/withdrawals live in the EL snapshots).
// For GLOAS blocks (version >= GloasVersion) the SSZ payload is the full SignedBeaconBlock,
// because the beacon block contains no ExecutionPayload — there is nothing to blind.
func WriteBlockForSnapshot(w io.Writer, block *cltypes.SignedBeaconBlock, reusable []byte) ([]byte, error) {
	bodyRoot, err := block.Block.Body.HashSSZ()
	if err != nil {
		return reusable, err
	}
	reusable = reusable[:0]
	version := block.Version()

	var encoded []byte
	if version >= clparams.GloasVersion {
		// [New in Gloas:EIP7732] No ExecutionPayload in beacon block; store the full block directly.
		encoded, err = block.EncodeSSZ(reusable)
	} else {
		// Pre-GLOAS: strip the execution payload to avoid duplicating data already held by the EL.
		blinded, blindErr := block.Blinded()
		if blindErr != nil {
			return reusable, blindErr
		}
		encoded, err = blinded.EncodeSSZ(reusable)
	}
	if err != nil {
		return reusable, err
	}

	reusable = encoded
	if _, err := w.Write([]byte{byte(version)}); err != nil {
		return reusable, err
	}
	if _, err := w.Write(bodyRoot[:]); err != nil {
		return reusable, err
	}
	// Write the length of the buffer
	length := make([]byte, 8)
	binary.BigEndian.PutUint64(length, uint64(len(reusable)))
	if _, err := w.Write(length); err != nil {
		return reusable, err
	}
	// Write the buffer
	if _, err := w.Write(reusable); err != nil {
		return reusable, err
	}
	return reusable, nil
}

func readMetadataForBlock(r io.Reader, b []byte) (clparams.StateVersion, common.Hash, error) {
	if _, err := r.Read(b); err != nil {
		return 0, common.Hash{}, err
	}
	return clparams.StateVersion(b[0]), common.BytesToHash(b[1:]), nil
}

func ReadBlockFromSnapshot(r io.Reader, executionReader ExecutionBlockReaderByNumber, cfg *clparams.BeaconChainConfig) (*cltypes.SignedBeaconBlock, error) {
	buffer := pool.GetBuffer()
	defer pool.PutBuffer(buffer)

	// Read the metadata
	metadataSlab := make([]byte, 33)
	v, _, err := readMetadataForBlock(r, metadataSlab)
	if err != nil {
		return nil, err
	}
	// Read the length
	length := make([]byte, 8)
	if _, err := io.ReadFull(r, length); err != nil {
		return nil, err
	}
	// Read the block bytes
	if _, err := io.CopyN(buffer, r, int64(binary.BigEndian.Uint64(length))); err != nil {
		return nil, err
	}

	// [New in Gloas:EIP7732] GLOAS blocks are stored as full SignedBeaconBlock (no blinding).
	if v >= clparams.GloasVersion {
		block := cltypes.NewSignedBeaconBlock(cfg, v)
		if err := block.DecodeSSZ(buffer.Bytes(), int(v)); err != nil {
			return nil, err
		}
		return block, nil
	}

	// Pre-GLOAS: stored as blinded block; reconstruct full block with execution data from EL.
	blindedBlock := cltypes.NewSignedBlindedBeaconBlock(cfg, v)
	if err := blindedBlock.DecodeSSZ(buffer.Bytes(), int(v)); err != nil {
		return nil, err
	}
	// No execution data for pre-altair blocks
	if v <= clparams.AltairVersion {
		return blindedBlock.Full(nil, nil), nil
	}
	// No execution reader: return with nil txs/withdrawals (sufficient for block parent lookup).
	if executionReader == nil {
		return blindedBlock.Full(nil, nil), nil
	}
	blockNumber := blindedBlock.Block.Body.ExecutionPayload.BlockNumber
	blockHash := blindedBlock.Block.Body.ExecutionPayload.BlockHash
	// Pre-Merge Bellatrix blocks carry a default (all-zero) execution payload header.
	// Their BlockHash is the zero hash, meaning no EL block exists to fetch transactions
	// from. This matches the consensus spec's is_execution_enabled check:
	// is_merge_transition_complete(state) || block.body.execution_payload.block_hash != Bytes32()
	// We don't have beacon state here, but a zero BlockHash is a reliable indicator that
	// execution is not yet enabled for this block.
	if blockHash == (common.Hash{}) {
		return blindedBlock.Full(nil, nil), nil
	}
	txs, err := executionReader.Transactions(blockNumber, blockHash)
	if err != nil {
		return nil, err
	}
	ws, err := executionReader.Withdrawals(blockNumber, blockHash)
	if err != nil {
		return nil, err
	}
	return blindedBlock.Full(txs, ws), nil
}

// ReadBlockHeaderFromSnapshotWithExecutionData reads the beacon block header and the EL block number and block hash.
func ReadBlockHeaderFromSnapshotWithExecutionData(r io.Reader, cfg *clparams.BeaconChainConfig) (*cltypes.SignedBeaconBlockHeader, uint64, common.Hash, error) {
	buffer := pool.GetBuffer()
	defer pool.PutBuffer(buffer)

	// Read the metadata
	metadataSlab := make([]byte, 33)
	v, bodyRoot, err := readMetadataForBlock(r, metadataSlab)
	if err != nil {
		return nil, 0, common.Hash{}, err
	}
	// Read the length
	length := make([]byte, 8)
	if _, err := io.ReadFull(r, length); err != nil {
		return nil, 0, common.Hash{}, err
	}
	// Read the block bytes
	if _, err := io.CopyN(buffer, r, int64(binary.BigEndian.Uint64(length))); err != nil {
		return nil, 0, common.Hash{}, err
	}

	// [New in Gloas:EIP7732] GLOAS blocks are stored as full SignedBeaconBlock (no blinding).
	// Execution block number and hash are not present in the beacon block; they come from the
	// SignedExecutionPayloadEnvelope and are tracked separately via WriteExecutionPayloadEnvelopeIndicies.
	if v >= clparams.GloasVersion {
		block := cltypes.NewSignedBeaconBlock(cfg, v)
		if err := block.DecodeSSZ(buffer.Bytes(), int(v)); err != nil {
			return nil, 0, common.Hash{}, err
		}
		blockHeader := &cltypes.SignedBeaconBlockHeader{
			Signature: block.Signature,
			Header: &cltypes.BeaconBlockHeader{
				Slot:          block.Block.Slot,
				ProposerIndex: block.Block.ProposerIndex,
				ParentRoot:    block.Block.ParentRoot,
				Root:          block.Block.StateRoot,
				BodyRoot:      bodyRoot,
			},
		}
		return blockHeader, 0, common.Hash{}, nil
	}

	// Pre-GLOAS: stored as blinded block.
	blindedBlock := cltypes.NewSignedBlindedBeaconBlock(cfg, v)
	if err := blindedBlock.DecodeSSZ(buffer.Bytes(), int(v)); err != nil {
		return nil, 0, common.Hash{}, err
	}
	blockHeader := &cltypes.SignedBeaconBlockHeader{
		Signature: blindedBlock.Signature,
		Header: &cltypes.BeaconBlockHeader{
			Slot:          blindedBlock.Block.Slot,
			ProposerIndex: blindedBlock.Block.ProposerIndex,
			ParentRoot:    blindedBlock.Block.ParentRoot,
			Root:          blindedBlock.Block.StateRoot,
			BodyRoot:      bodyRoot,
		},
	}
	// No execution data for pre-altair blocks
	if v <= clparams.AltairVersion {
		return blockHeader, 0, common.Hash{}, nil
	}
	blockNumber := blindedBlock.Block.Body.ExecutionPayload.BlockNumber
	blockHash := blindedBlock.Block.Body.ExecutionPayload.BlockHash
	return blockHeader, blockNumber, blockHash, nil
}

// ReadBeaconBlockBodyFromSnapshot reads a beacon block from the snapshot without reconstructing
// execution payload transactions or withdrawals. Works for both pre- and post-GLOAS:
//   - Pre-GLOAS: stored as SignedBlindedBeaconBlock; decoded and promoted to SignedBeaconBlock
//     with nil transactions/withdrawals (suitable for attestation replay, blob counting, etc.).
//     Callers that need full EL data should use ReadBlockFromSnapshot instead.
//   - GLOAS: stored as full SignedBeaconBlock (no payload to blind); decoded directly.
func ReadBeaconBlockBodyFromSnapshot(r io.Reader, cfg *clparams.BeaconChainConfig) (*cltypes.SignedBeaconBlock, error) {
	buffer := pool.GetBuffer()
	defer pool.PutBuffer(buffer)

	// Read the metadata
	metadataSlab := make([]byte, 33)
	v, _, err := readMetadataForBlock(r, metadataSlab)
	if err != nil {
		return nil, err
	}
	// Read the length
	length := make([]byte, 8)
	if _, err := io.ReadFull(r, length); err != nil {
		return nil, err
	}
	// Read the block bytes
	if _, err := io.CopyN(buffer, r, int64(binary.BigEndian.Uint64(length))); err != nil {
		return nil, err
	}

	// [New in Gloas:EIP7732] GLOAS blocks are stored as full SignedBeaconBlock.
	if v >= clparams.GloasVersion {
		block := cltypes.NewSignedBeaconBlock(cfg, v)
		if err := block.DecodeSSZ(buffer.Bytes(), int(v)); err != nil {
			return nil, err
		}
		return block, nil
	}

	// Pre-GLOAS: stored as blinded block; promote to full block with empty execution data.
	blindedBlock := cltypes.NewSignedBlindedBeaconBlock(cfg, v)
	if err := blindedBlock.DecodeSSZ(buffer.Bytes(), int(v)); err != nil {
		return nil, err
	}
	txs := &solid.TransactionsSSZ{}
	var ws *solid.ListSSZ[*cltypes.Withdrawal]
	if v >= clparams.CapellaVersion {
		ws = solid.NewStaticListSSZ[*cltypes.Withdrawal](int(cfg.MaxWithdrawalsPerPayload), 44)
	}
	return blindedBlock.Full(txs, ws), nil
}
