// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package state

import (
	"time"

	"github.com/luxfi/codec"
	"github.com/luxfi/constants"
	"github.com/luxfi/database"
	"github.com/luxfi/ids"
	"github.com/luxfi/math/set"
	"github.com/luxfi/codec/wrappers"
)

// preDelegateeRewardSize is the size of codec marshalling
// [preDelegateeRewardMetadata].
//
// CodecVersionLen + UpDurationLen + LastUpdatedLen + PotentialRewardLen
const preDelegateeRewardSize = codec.VersionSize + 3*wrappers.LongLen

// preStakerStartTimeSize is the size of codec marshalling
// [preStakerStartTimeMetadata].
//
// CodecVersionLen + UpDurationLen + LastUpdatedLen + PotentialRewardLen + PotentialDelegateeRewardLen
const preStakerStartTimeSize = codec.VersionSize + 4*wrappers.LongLen

var _ validatorState = (*metadata)(nil)

type preDelegateeRewardMetadata struct {
	UpDuration      time.Duration `serialize:"true"`
	LastUpdated     uint64        `serialize:"true"` // Unix time in seconds
	PotentialReward uint64        `serialize:"true"`
}

// preStakerStartTimeMetadata is used for backward compatibility with data
// that was written before StakerStartTime was added.
type preStakerStartTimeMetadata struct {
	UpDuration               time.Duration `serialize:"true"`
	LastUpdated              uint64        `serialize:"true"` // Unix time in seconds
	PotentialReward          uint64        `serialize:"true"`
	PotentialDelegateeReward uint64        `serialize:"true"`
}

type validatorMetadata struct {
	UpDuration               time.Duration `serialize:"true"`
	LastUpdated              uint64        `serialize:"true"` // Unix time in seconds
	PotentialReward          uint64        `serialize:"true"`
	PotentialDelegateeReward uint64        `serialize:"true"`
	StakerStartTime          uint64        `serialize:"true"`

	txID        ids.ID
	lastUpdated time.Time
}

// Permissioned validators originally wrote their values as nil.
// We now write the uptime, reward, and delegatee reward together.
func parseValidatorMetadata(bytes []byte, metadata *validatorMetadata) error {
	switch len(bytes) {
	case 0:
	// nothing was stored

	case database.Uint64Size:
		// only potential reward was stored
		var err error
		metadata.PotentialReward, err = database.ParseUInt64(bytes)
		if err != nil {
			return err
		}

	case preDelegateeRewardSize:
		// potential reward and uptime was stored but potential delegatee reward
		// was not
		tmp := preDelegateeRewardMetadata{}
		if _, err := MetadataCodec.Unmarshal(bytes, &tmp); err != nil {
			return err
		}

		metadata.UpDuration = tmp.UpDuration
		metadata.LastUpdated = tmp.LastUpdated
		metadata.PotentialReward = tmp.PotentialReward

	case preStakerStartTimeSize:
		// All fields except StakerStartTime were stored (pre-v1 format)
		tmp := preStakerStartTimeMetadata{}
		if _, err := MetadataCodec.Unmarshal(bytes, &tmp); err != nil {
			return err
		}

		metadata.UpDuration = tmp.UpDuration
		metadata.LastUpdated = tmp.LastUpdated
		metadata.PotentialReward = tmp.PotentialReward
		metadata.PotentialDelegateeReward = tmp.PotentialDelegateeReward

	default:
		// everything was stored (v1+ format with StakerStartTime)
		if _, err := MetadataCodec.Unmarshal(bytes, metadata); err != nil {
			return err
		}
	}
	metadata.lastUpdated = time.Unix(int64(metadata.LastUpdated), 0)
	return nil
}

type validatorState interface {
	// LoadValidatorMetadata sets the [metadata] of [vdrID] on [netID].
	// GetUptime and SetUptime will return an error if the [vdrID] and
	// [netID] hasn't been loaded. This call will not result in a write to
	// disk.
	LoadValidatorMetadata(
		vdrID ids.NodeID,
		netID ids.ID,
		metadata *validatorMetadata,
	)

	// GetUptime returns the current uptime measurements of [vdrID] on
	// [netID].
	GetUptime(
		vdrID ids.NodeID,
		netID ids.ID,
	) (upDuration time.Duration, lastUpdated time.Time, err error)

	// SetUptime updates the uptime measurements of [vdrID] on [netID].
	// Unless these measurements are deleted first, the next call to
	// WriteUptimes will write this update to disk.
	SetUptime(
		vdrID ids.NodeID,
		netID ids.ID,
		upDuration time.Duration,
		lastUpdated time.Time,
	) error

	// GetDelegateeReward returns the current rewards accrued to [vdrID] on
	// [netID].
	GetDelegateeReward(
		netID ids.ID,
		vdrID ids.NodeID,
	) (amount uint64, err error)

	// SetDelegateeReward updates the rewards accrued to [vdrID] on [netID].
	// Unless these measurements are deleted first, the next call to
	// WriteUptimes will write this update to disk.
	SetDelegateeReward(
		netID ids.ID,
		vdrID ids.NodeID,
		amount uint64,
	) error

	// DeleteValidatorMetadata removes in-memory references to the metadata of
	// [vdrID] on [netID]. If there were staged updates from a prior call to
	// SetUptime or SetDelegateeReward, the updates will be dropped. This call
	// will not result in a write to disk.
	DeleteValidatorMetadata(vdrID ids.NodeID, netID ids.ID)

	// WriteValidatorMetadata writes all staged updates from prior calls to
	// SetUptime or SetDelegateeReward.
	WriteValidatorMetadata(
		dbPrimary database.KeyValueWriter,
		dbNet database.KeyValueWriter,
		codecVersion uint16,
	) error
}

type metadata struct {
	metadata map[ids.NodeID]map[ids.ID]*validatorMetadata // vdrID -> netID -> metadata
	// updatedMetadata tracks the updates since WriteValidatorMetadata was last called
	updatedMetadata map[ids.NodeID]set.Set[ids.ID] // vdrID -> netIDs
}

func newValidatorState() validatorState {
	return &metadata{
		metadata:        make(map[ids.NodeID]map[ids.ID]*validatorMetadata),
		updatedMetadata: make(map[ids.NodeID]set.Set[ids.ID]),
	}
}

func (m *metadata) LoadValidatorMetadata(
	vdrID ids.NodeID,
	netID ids.ID,
	uptime *validatorMetadata,
) {
	chainMetadata, ok := m.metadata[vdrID]
	if !ok {
		chainMetadata = make(map[ids.ID]*validatorMetadata)
		m.metadata[vdrID] = chainMetadata
	}
	chainMetadata[netID] = uptime
}

func (m *metadata) GetUptime(
	vdrID ids.NodeID,
	netID ids.ID,
) (time.Duration, time.Time, error) {
	metadata, exists := m.metadata[vdrID][netID]
	if !exists {
		return 0, time.Time{}, database.ErrNotFound
	}
	return metadata.UpDuration, metadata.lastUpdated, nil
}

func (m *metadata) SetUptime(
	vdrID ids.NodeID,
	netID ids.ID,
	upDuration time.Duration,
	lastUpdated time.Time,
) error {
	metadata, exists := m.metadata[vdrID][netID]
	if !exists {
		return database.ErrNotFound
	}
	metadata.UpDuration = upDuration
	metadata.lastUpdated = lastUpdated

	m.addUpdatedMetadata(vdrID, netID)
	return nil
}

func (m *metadata) GetDelegateeReward(
	netID ids.ID,
	vdrID ids.NodeID,
) (uint64, error) {
	metadata, exists := m.metadata[vdrID][netID]
	if !exists {
		return 0, database.ErrNotFound
	}
	return metadata.PotentialDelegateeReward, nil
}

func (m *metadata) SetDelegateeReward(
	netID ids.ID,
	vdrID ids.NodeID,
	amount uint64,
) error {
	metadata, exists := m.metadata[vdrID][netID]
	if !exists {
		return database.ErrNotFound
	}
	metadata.PotentialDelegateeReward = amount

	m.addUpdatedMetadata(vdrID, netID)
	return nil
}

func (m *metadata) DeleteValidatorMetadata(vdrID ids.NodeID, netID ids.ID) {
	chainMetadata := m.metadata[vdrID]
	delete(chainMetadata, netID)
	if len(chainMetadata) == 0 {
		delete(m.metadata, vdrID)
	}

	chainUpdatedMetadata := m.updatedMetadata[vdrID]
	chainUpdatedMetadata.Remove(netID)
	if chainUpdatedMetadata.Len() == 0 {
		delete(m.updatedMetadata, vdrID)
	}
}

func (m *metadata) WriteValidatorMetadata(
	dbPrimary database.KeyValueWriter,
	dbNet database.KeyValueWriter,
	codecVersion uint16,
) error {
	for vdrID, updatedNets := range m.updatedMetadata {
		for netID := range updatedNets {
			metadata := m.metadata[vdrID][netID]
			metadata.LastUpdated = uint64(metadata.lastUpdated.Unix())

			metadataBytes, err := MetadataCodec.Marshal(codecVersion, metadata)
			if err != nil {
				return err
			}
			db := dbNet
			if netID == constants.PrimaryNetworkID {
				db = dbPrimary
			}
			if err := db.Put(metadata.txID[:], metadataBytes); err != nil {
				return err
			}
		}
		delete(m.updatedMetadata, vdrID)
	}
	return nil
}

func (m *metadata) addUpdatedMetadata(vdrID ids.NodeID, netID ids.ID) {
	updatedNetMetadata, ok := m.updatedMetadata[vdrID]
	if !ok {
		updatedNetMetadata = make(set.Set[ids.ID])
		m.updatedMetadata[vdrID] = updatedNetMetadata
	}
	updatedNetMetadata.Add(netID)
}
