tutus-chain/pkg/tutustest/government.go

package tutustest

import (
	"math/big"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"git.marketally.com/tutus-one/tutus-chain/pkg/core/native/nativenames"
	"git.marketally.com/tutus-one/tutus-chain/pkg/core/transaction"
	"git.marketally.com/tutus-one/tutus-chain/pkg/util"
	"git.marketally.com/tutus-one/tutus-chain/pkg/vm/stackitem"
)

// GovernmentHelper provides utilities for testing Tutus government contracts.
// It wraps common operations like registering citizens, granting roles, and
// setting up cross-contract test scenarios.
type GovernmentHelper struct {
	*Executor
	t testing.TB

	// Contract invokers for common government contracts
	Vita     *ContractInvoker
	Lex      *ContractInvoker
	Eligere  *ContractInvoker
	Scire    *ContractInvoker
	Salus    *ContractInvoker
	VTS      *ContractInvoker
	Annos    *ContractInvoker
	Tribute  *ContractInvoker
	RoleReg  *ContractInvoker
	Treasury *ContractInvoker
}

// NewGovernmentHelper creates a helper for testing government contracts.
// It initializes invokers for all major government contracts.
func NewGovernmentHelper(t testing.TB, e *Executor) *GovernmentHelper {
	g := &GovernmentHelper{
		Executor: e,
		t:        t,
	}

	// Initialize contract invokers with committee authority
	g.Vita = e.CommitteeInvoker(e.NativeHash(t, nativenames.Vita))
	g.Lex = e.CommitteeInvoker(e.NativeHash(t, nativenames.Lex))
	g.Eligere = e.CommitteeInvoker(e.NativeHash(t, nativenames.Eligere))
	g.Scire = e.CommitteeInvoker(e.NativeHash(t, nativenames.Scire))
	g.Salus = e.CommitteeInvoker(e.NativeHash(t, nativenames.Salus))
	g.VTS = e.CommitteeInvoker(e.NativeHash(t, nativenames.VTS))
	g.Annos = e.CommitteeInvoker(e.NativeHash(t, nativenames.Annos))
	g.Tribute = e.CommitteeInvoker(e.NativeHash(t, nativenames.Tribute))
	g.RoleReg = e.CommitteeInvoker(e.NativeHash(t, nativenames.RoleRegistry))
	g.Treasury = e.CommitteeInvoker(e.NativeHash(t, nativenames.Treasury))

	return g
}

// Citizen represents a registered Vita holder for testing.
type Citizen struct {
	Account    SingleSigner
	VitaID     uint64
	BirthTime  uint64
	Registered bool
}

// RegisterCitizen registers a new Vita token for the given account.
// birthTimestamp is the citizen's birth date (Unix timestamp in milliseconds).
// Returns the Citizen with VitaID populated.
func (g *GovernmentHelper) RegisterCitizen(account SingleSigner, birthTimestamp uint64) *Citizen {
	citizen := &Citizen{
		Account:   account,
		BirthTime: birthTimestamp,
	}

	// Get current token count to predict the new VitaID
	g.Vita.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) {
		count, err := stack[0].TryInteger()
		require.NoError(t, err)
		citizen.VitaID = count.Uint64()
	}, "totalSupply")

	// Register the Vita token
	g.Vita.Invoke(g.t, true, "register", account.ScriptHash(), birthTimestamp)
	citizen.Registered = true

	return citizen
}

// RegisterAdultCitizen registers a citizen who is 25 years old (adult).
// Useful for tests requiring voting age or adult status.
func (g *GovernmentHelper) RegisterAdultCitizen(account SingleSigner) *Citizen {
	// 25 years ago in milliseconds
	birthTime := uint64(time.Now().AddDate(-25, 0, 0).UnixMilli())
	return g.RegisterCitizen(account, birthTime)
}

// RegisterChildCitizen registers a citizen who is 10 years old (child).
// Useful for tests verifying age restrictions.
func (g *GovernmentHelper) RegisterChildCitizen(account SingleSigner) *Citizen {
	// 10 years ago in milliseconds
	birthTime := uint64(time.Now().AddDate(-10, 0, 0).UnixMilli())
	return g.RegisterCitizen(account, birthTime)
}

// RegisterElderCitizen registers a citizen who is 70 years old (elder).
// Useful for tests involving retirement age.
func (g *GovernmentHelper) RegisterElderCitizen(account SingleSigner) *Citizen {
	// 70 years ago in milliseconds
	birthTime := uint64(time.Now().AddDate(-70, 0, 0).UnixMilli())
	return g.RegisterCitizen(account, birthTime)
}

// NewCitizenAccount creates a new account and registers it as a citizen.
// Returns both the citizen and their account for signing transactions.
func (g *GovernmentHelper) NewCitizenAccount() *Citizen {
	acc := g.Vita.NewAccount(g.t).(SingleSigner)
	return g.RegisterAdultCitizen(acc)
}

// VerifyVitaOwnership checks that the given account owns a Vita token.
func (g *GovernmentHelper) VerifyVitaOwnership(account util.Uint160) bool {
	var hasVita bool
	g.Vita.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) {
		balance, err := stack[0].TryInteger()
		require.NoError(t, err)
		hasVita = balance.Cmp(big.NewInt(0)) > 0
	}, "balanceOf", account)
	return hasVita
}

// GetVitaID returns the Vita token ID for the given owner, or -1 if not found.
func (g *GovernmentHelper) GetVitaID(owner util.Uint160) int64 {
	var vitaID int64 = -1
	g.Vita.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) {
		if stack[0].Type() != stackitem.AnyT {
			tokens := stack[0].Value().([]stackitem.Item)
			if len(tokens) > 0 {
				id, err := tokens[0].TryInteger()
				require.NoError(t, err)
				vitaID = id.Int64()
			}
		}
	}, "tokensOf", owner)
	return vitaID
}

// SuspendVita suspends a citizen's Vita token (requires committee + Lex restriction).
func (g *GovernmentHelper) SuspendVita(vitaID uint64, reason string) {
	// First, create a Lex liberty restriction (required for due process)
	// This would normally require a judicial order
	g.Vita.Invoke(g.t, true, "suspend", vitaID, reason)
}

// TransferVTS transfers VTS tokens between accounts.
func (g *GovernmentHelper) TransferVTS(from, to SingleSigner, amount int64) util.Uint256 {
	vtsInvoker := g.VTS.WithSigners(from)
	return vtsInvoker.Invoke(g.t, true, "transfer", from.ScriptHash(), to.ScriptHash(), amount, nil)
}

// GetVTSBalance returns the VTS balance for an account.
func (g *GovernmentHelper) GetVTSBalance(account util.Uint160) *big.Int {
	var balance *big.Int
	g.VTS.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) {
		var err error
		balance, err = stack[0].TryInteger()
		require.NoError(t, err)
	}, "balanceOf", account)
	return balance
}

// MintVTS mints VTS tokens to an account (committee only).
func (g *GovernmentHelper) MintVTS(to util.Uint160, amount int64) {
	g.VTS.Invoke(g.t, true, "mint", to, amount)
}

// CreateProposal creates a new Eligere proposal.
// Returns the proposal ID.
func (g *GovernmentHelper) CreateProposal(proposer *Citizen, title, description string, category int64) uint64 {
	eligereInvoker := g.Eligere.WithSigners(proposer.Account)

	var proposalID uint64
	eligereInvoker.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) {
		id, err := stack[0].TryInteger()
		require.NoError(t, err)
		proposalID = id.Uint64()
	}, "createProposal", title, description, category)

	return proposalID
}

// Vote casts a vote on a proposal.
func (g *GovernmentHelper) Vote(voter *Citizen, proposalID uint64, vote int64) {
	eligereInvoker := g.Eligere.WithSigners(voter.Account)
	eligereInvoker.Invoke(g.t, true, "vote", proposalID, vote)
}

// PrepareVitaRegistration prepares a Vita registration transaction without executing.
// Useful for batching or testing transaction ordering.
func (g *GovernmentHelper) PrepareVitaRegistration(account SingleSigner, birthTimestamp uint64) *transaction.Transaction {
	return g.Vita.PrepareInvoke(g.t, "register", account.ScriptHash(), birthTimestamp)
}

// BatchRegisterCitizens registers multiple citizens in a single block.
func (g *GovernmentHelper) BatchRegisterCitizens(accounts []SingleSigner) []*Citizen {
	citizens := make([]*Citizen, len(accounts))
	txs := make([]*transaction.Transaction, len(accounts))
	birthTime := uint64(time.Now().AddDate(-25, 0, 0).UnixMilli())

	// Get starting VitaID
	var startID uint64
	g.Vita.InvokeAndCheck(g.t, func(t testing.TB, stack []stackitem.Item) {
		count, err := stack[0].TryInteger()
		require.NoError(t, err)
		startID = count.Uint64()
	}, "totalSupply")

	// Prepare all transactions
	for i, acc := range accounts {
		citizens[i] = &Citizen{
			Account:   acc,
			BirthTime: birthTime,
			VitaID:    startID + uint64(i),
		}
		txs[i] = g.Vita.PrepareInvoke(g.t, "register", acc.ScriptHash(), birthTime)
	}

	// Execute all in one block
	g.Vita.AddNewBlock(g.t, txs...)

	// Verify all succeeded
	for i, tx := range txs {
		g.Vita.CheckHalt(g.t, tx.Hash(), stackitem.Make(true))
		citizens[i].Registered = true
	}

	return citizens
}