// Copyright (c) 2024 XDC Network
// Snapshot management for XDPoS 2.0

package engine_v2

import (
	"encoding/json"
	"errors"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/ethdb"
	"github.com/ethereum/go-ethereum/log"
)

var ErrNotFoundBlockByNum = errors.New("not found block by number")

// SnapshotV2 represents the state of masternodes at a given block.
// Aligned with v2.6.8: JSON tag must be "masterNodes" for DB compatibility.
//
// Field semantics (checkpoint-sync aware):
//   - NextEpochCandidates: the ACTIVE masternodes for the epoch this snapshot
//     governs (top-N after penalty filtering). This is what
//     getEpochSwitchInfo's synthetic path (engine.go:1818) treats as the
//     active list for proposer/leader rotation. v2.6.8 production snapshots
//     also populate this with the post-filter active set when
//     UpdateMasternodesFromHeader runs against contract state.
//   - Penalties: the penalty addresses associated with this snapshot, stored
//     SEPARATELY so calcMasternodes can hand HookPenalty a "candidates pool"
//     that's a superset of `pens` (the comeback algorithm intersects
//     prev-epoch penalties against this pool). Only populated by
//     checkpoint-sync seed paths; v2.6.8 snapshots leave it empty
//     (omitempty preserves on-disk compat).
type SnapshotV2 struct {
	Version             uint64           `json:"version,omitempty"` // 1 = legacy, 2 = V2 with proper masternode handling (omitempty for v2.6.8 compat)
	Number              uint64           `json:"number"`
	Hash                common.Hash      `json:"hash"`
	NextEpochCandidates []common.Address `json:"masterNodes"`
	Penalties           []common.Address `json:"penalties,omitempty"`
}

// newSnapshot creates a new snapshot with the current version
func newSnapshot(number uint64, hash common.Hash, masternodes []common.Address) *SnapshotV2 {
	snap := &SnapshotV2{
		Version:             4, // Bumped to invalidate corrupted v57/v60 snapshots (v62 fix)
		Number:              number,
		Hash:                hash,
		NextEpochCandidates: make([]common.Address, len(masternodes)),
	}
	copy(snap.NextEpochCandidates, masternodes)
	return snap
}

// newSnapshotWithPenalties creates a checkpoint-sync snapshot that records
// both the active masternode list and the penalty pool from the same epoch.
// The two are kept in distinct fields so the synthetic-epoch-info path
// (which wants active masternodes for proposer rotation) and HookPenalty's
// comeback path (which wants the broader candidate pool to intersect against
// prev-epoch penalties) don't fight over a single field.
func newSnapshotWithPenalties(number uint64, hash common.Hash, masternodes, penalties []common.Address) *SnapshotV2 {
	snap := newSnapshot(number, hash, masternodes)
	if len(penalties) > 0 {
		snap.Penalties = make([]common.Address, len(penalties))
		copy(snap.Penalties, penalties)
	}
	return snap
}

// CandidatePool returns the full candidate pool for HookPenalty:
// active masternodes ∪ recorded penalties (deduped). For v2.6.8-produced
// snapshots (Penalties empty) this is just NextEpochCandidates, preserving
// the canonical algorithm.
func (s *SnapshotV2) CandidatePool() []common.Address {
	if len(s.Penalties) == 0 {
		return s.NextEpochCandidates
	}
	pool := make([]common.Address, len(s.NextEpochCandidates), len(s.NextEpochCandidates)+len(s.Penalties))
	copy(pool, s.NextEpochCandidates)
	seen := make(map[common.Address]struct{}, len(pool))
	for _, a := range pool {
		seen[a] = struct{}{}
	}
	for _, p := range s.Penalties {
		if _, ok := seen[p]; !ok {
			pool = append(pool, p)
			seen[p] = struct{}{}
		}
	}
	return pool
}

// loadSnapshot loads a snapshot from the database
// Aligned with v2.6.8: DB key prefix is "XDPoS-V2-" + hash[:]
func loadSnapshot(db ethdb.Database, hash common.Hash) (*SnapshotV2, error) {
	key := append([]byte("XDPoS-V2-"), hash[:]...)
	blob, err := db.Get(key)
	if err != nil {
		return nil, err
	}
	snap := new(SnapshotV2)
	if err := json.Unmarshal(blob, snap); err != nil {
		return nil, err
	}
	return snap, nil
}

// storeSnapshot stores a snapshot to the database
// Aligned with v2.6.8: DB key prefix is "XDPoS-V2-" + hash[:]
func storeSnapshot(snap *SnapshotV2, db ethdb.Database) error {
	blob, err := json.Marshal(snap)
	if err != nil {
		return err
	}
	key := append([]byte("XDPoS-V2-"), snap.Hash[:]...)
	if err := db.Put(key, blob); err != nil {
		log.Error("Failed to store snapshot", "hash", snap.Hash, "error", err)
		return err
	}
	return nil
}

// Copy creates a copy of the snapshot
func (s *SnapshotV2) Copy() *SnapshotV2 {
	return newSnapshot(s.Number, s.Hash, s.NextEpochCandidates)
}

// GetSigners returns the list of masternodes
func (s *SnapshotV2) GetSigners() []common.Address {
	return s.NextEpochCandidates
}

// GetMappedCandidates returns the candidate list as a map for O(1) lookups.
// Matches v2.6.8 engines/engine_v2/snapshot.go.
func (s *SnapshotV2) GetMappedCandidates() map[common.Address]struct{} {
	ms := make(map[common.Address]struct{})
	for _, n := range s.NextEpochCandidates {
		ms[n] = struct{}{}
	}
	return ms
}

// IsCandidates checks if an address is in the candidate list.
// Matches v2.6.8 engines/engine_v2/snapshot.go.
func (s *SnapshotV2) IsCandidates(address common.Address) bool {
	for _, n := range s.NextEpochCandidates {
		if n == address {
			return true
		}
	}
	return false
}

// IsValidForV2Switch returns whether this snapshot is valid for use at the V2 switch block.
// Legacy snapshots (version < 2) may have incorrect masternode data at the V1->V2 boundary.
func (s *SnapshotV2) IsValidForV2Switch() bool {
	return s.Version >= 2
}