//go:build include_onchain_tests
// +build include_onchain_tests

package rln

import (
	"context"
	"crypto/ecdsa"
	"crypto/rand"
	"encoding/hex"
	"math/big"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/waku-org/go-zerokit-rln/rln"

	"github.com/ethereum/go-ethereum/accounts/abi/bind"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/ethereum/go-ethereum/ethclient"
	"github.com/stretchr/testify/suite"
	"github.com/waku-org/go-waku/waku/v2/protocol/rln/contracts"
	"github.com/waku-org/go-waku/waku/v2/protocol/rln/group_manager"
	"github.com/waku-org/go-waku/waku/v2/protocol/rln/group_manager/dynamic"
	"github.com/waku-org/go-waku/waku/v2/protocol/rln/keystore"
	"github.com/waku-org/go-waku/waku/v2/protocol/rln/web3"
	"github.com/waku-org/go-waku/waku/v2/timesource"
	"github.com/waku-org/go-waku/waku/v2/utils"
)

const keystorePassword = "test"

func TestWakuRLNRelayDynamicSuite(t *testing.T) {
	suite.Run(t, new(WakuRLNRelayDynamicSuite))
}

type WakuRLNRelayDynamicSuite struct {
	suite.Suite
	web3Config *web3.Config
	u1PrivKey  *ecdsa.PrivateKey
	u2PrivKey  *ecdsa.PrivateKey
}

func TempFileName(prefix, suffix string) string {
	randBytes := make([]byte, 16)
	rand.Read(randBytes)
	return filepath.Join(os.TempDir(), prefix+hex.EncodeToString(randBytes)+suffix)
}

func (s *WakuRLNRelayDynamicSuite) SetupTest() {

	clientAddr := os.Getenv("GANACHE_NETWORK_RPC_URL")
	if clientAddr == "" {
		clientAddr = "ws://localhost:8545"
	}

	backend, err := ethclient.Dial(clientAddr)
	s.Require().NoError(err)
	defer backend.Close()

	chainID, err := backend.ChainID(context.TODO())
	s.Require().NoError(err)

	// TODO: obtain account list from ganache mnemonic or from eth_accounts

	s.u1PrivKey, err = crypto.ToECDSA(common.FromHex("0x156ec84a451d8a2d0062993242b6c4e863647f5544ff8030f23578d4142f43f8"))
	s.Require().NoError(err)
	s.u2PrivKey, err = crypto.ToECDSA(common.FromHex("0xa00da43843ad6b5161ddbace48f293ac3f82f8a8257af34de4c32900bb6e9a97"))
	s.Require().NoError(err)

	// Deploying contracts
	auth, err := bind.NewKeyedTransactorWithChainID(s.u1PrivKey, chainID)
	s.Require().NoError(err)

	poseidonHasherAddr, _, _, err := contracts.DeployPoseidonHasher(auth, backend)
	s.Require().NoError(err)

	registryAddress, tx, rlnRegistry, err := contracts.DeployRLNRegistry(auth, backend, poseidonHasherAddr)
	s.Require().NoError(err)
	txReceipt, err := bind.WaitMined(context.TODO(), backend, tx)
	s.Require().NoError(err)
	s.Require().Equal(txReceipt.Status, types.ReceiptStatusSuccessful)

	tx, err = rlnRegistry.NewStorage(auth)
	s.Require().NoError(err)
	txReceipt, err = bind.WaitMined(context.TODO(), backend, tx)
	s.Require().NoError(err)
	s.Require().Equal(txReceipt.Status, types.ReceiptStatusSuccessful)

	s.web3Config = web3.NewConfig(clientAddr, registryAddress)
	err = s.web3Config.Build(context.TODO())
	s.Require().NoError(err)
}

func (s *WakuRLNRelayDynamicSuite) generateCredentials(rlnInstance *rln.RLN) *rln.IdentityCredential {
	identityCredential, err := rlnInstance.MembershipKeyGen()
	s.Require().NoError(err)
	return identityCredential
}

func (s *WakuRLNRelayDynamicSuite) register(appKeystore *keystore.AppKeystore, identityCredential *rln.IdentityCredential, privKey *ecdsa.PrivateKey) rln.MembershipIndex {
	membershipFee, err := s.web3Config.RLNContract.MEMBERSHIPDEPOSIT(&bind.CallOpts{Context: context.TODO()})
	s.Require().NoError(err)

	auth, err := bind.NewKeyedTransactorWithChainID(privKey, s.web3Config.ChainID)
	s.Require().NoError(err)

	auth.Value = membershipFee
	auth.Context = context.TODO()
	tx, err := s.web3Config.RegistryContract.Register(auth, s.web3Config.RLNContract.StorageIndex, rln.Bytes32ToBigInt(identityCredential.IDCommitment))
	s.Require().NoError(err)

	txReceipt, err := bind.WaitMined(context.TODO(), s.web3Config.ETHClient, tx)
	s.Require().NoError(err)
	s.Require().Equal(txReceipt.Status, types.ReceiptStatusSuccessful)

	evt, err := s.web3Config.RLNContract.ParseMemberRegistered(*txReceipt.Logs[0])
	s.Require().NoError(err)

	membershipIndex := rln.MembershipIndex(uint(evt.Index.Int64()))

	membershipCredential := keystore.MembershipCredentials{
		IdentityCredential:     identityCredential,
		TreeIndex:              membershipIndex,
		MembershipContractInfo: keystore.NewMembershipContractInfo(s.web3Config.ChainID, s.web3Config.RegistryContract.Address),
	}

	err = appKeystore.AddMembershipCredentials(membershipCredential, keystorePassword)
	s.Require().NoError(err)

	return membershipIndex
}

func (s *WakuRLNRelayDynamicSuite) TestDynamicGroupManagement() {
	// Create a RLN instance
	rlnInstance, err := rln.NewRLN()
	s.Require().NoError(err)

	rt, err := group_manager.NewMerkleRootTracker(5, rlnInstance)
	s.Require().NoError(err)

	u1Credentials := s.generateCredentials(rlnInstance)
	appKeystore, err := keystore.New(s.tmpKeystorePath(), dynamic.RLNAppInfo, utils.Logger())
	s.Require().NoError(err)

	membershipIndex := s.register(appKeystore, u1Credentials, s.u1PrivKey)

	gm, err := dynamic.NewDynamicGroupManager(s.web3Config.ETHClientAddress, s.web3Config.RegistryContract.Address, &membershipIndex, appKeystore, keystorePassword, prometheus.DefaultRegisterer, rlnInstance, rt, utils.Logger())
	s.Require().NoError(err)

	// initialize the WakuRLNRelay
	rlnRelay := &WakuRLNRelay{
		Details: group_manager.Details{
			RootTracker:  rt,
			GroupManager: gm,
			RLN:          rlnInstance,
		},
		log:          utils.Logger(),
		nullifierLog: make(map[rln.MerkleNode][]rln.ProofMetadata),
	}

	err = rlnRelay.Start(context.TODO())
	s.Require().NoError(err)

	u2Credentials := s.generateCredentials(rlnInstance)
	appKeystore2, err := keystore.New(s.tmpKeystorePath(), dynamic.RLNAppInfo, utils.Logger())
	s.Require().NoError(err)

	membershipIndex = s.register(appKeystore2, u2Credentials, s.u2PrivKey)

	time.Sleep(1 * time.Second)

	treeCommitment, err := rlnInstance.GetLeaf(membershipIndex)
	s.Require().NoError(err)
	s.Require().Equal(u2Credentials.IDCommitment, treeCommitment)
}

func (s *WakuRLNRelayDynamicSuite) TestInsertKeyMembershipContract() {
	// Create a RLN instance
	rlnInstance, err := rln.NewRLN()
	s.Require().NoError(err)

	credentials1 := s.generateCredentials(rlnInstance)
	credentials2 := s.generateCredentials(rlnInstance)
	credentials3 := s.generateCredentials(rlnInstance)

	appKeystore, err := keystore.New(s.tmpKeystorePath(), dynamic.RLNAppInfo, utils.Logger())
	s.Require().NoError(err)

	s.register(appKeystore, credentials1, s.u1PrivKey)

	// Batch Register
	auth, err := bind.NewKeyedTransactorWithChainID(s.u2PrivKey, s.web3Config.ChainID)
	s.Require().NoError(err)

	membershipFee, err := s.web3Config.RLNContract.MEMBERSHIPDEPOSIT(&bind.CallOpts{Context: context.TODO()})
	s.Require().NoError(err)

	auth.Value = membershipFee.Mul(big.NewInt(2), membershipFee)
	auth.Context = context.TODO()

	tx, err := s.web3Config.RegistryContract.Register1(auth, s.web3Config.RLNContract.StorageIndex, []*big.Int{rln.Bytes32ToBigInt(credentials2.IDCommitment), rln.Bytes32ToBigInt(credentials3.IDCommitment)})
	s.Require().NoError(err)

	txReceipt, err := bind.WaitMined(context.TODO(), s.web3Config.ETHClient, tx)
	s.Require().NoError(err)
	s.Require().Equal(txReceipt.Status, types.ReceiptStatusSuccessful)
}

func (s *WakuRLNRelayDynamicSuite) TestMerkleTreeConstruction() {
	// Create a RLN instance
	rlnInstance, err := rln.NewRLN()
	s.Require().NoError(err)

	credentials1 := s.generateCredentials(rlnInstance)
	credentials2 := s.generateCredentials(rlnInstance)

	err = rlnInstance.InsertMembers(0, []rln.IDCommitment{credentials1.IDCommitment, credentials2.IDCommitment})
	s.Require().NoError(err)

	//  get the Merkle root
	expectedRoot, err := rlnInstance.GetMerkleRoot()
	s.Require().NoError(err)

	// register the members to the contract
	appKeystore, err := keystore.New(s.tmpKeystorePath(), dynamic.RLNAppInfo, utils.Logger())
	s.Require().NoError(err)

	membershipIndex := s.register(appKeystore, credentials1, s.u1PrivKey)
	membershipIndex = s.register(appKeystore, credentials2, s.u1PrivKey)

	rlnInstance, rootTracker, err := GetRLNInstanceAndRootTracker(s.tmpRLNDBPath())
	s.Require().NoError(err)
	// mount the rln relay protocol in the on-chain/dynamic mode
	gm, err := dynamic.NewDynamicGroupManager(s.web3Config.ETHClientAddress, s.web3Config.RegistryContract.Address, &membershipIndex, appKeystore, keystorePassword, prometheus.DefaultRegisterer, rlnInstance, rootTracker, utils.Logger())
	s.Require().NoError(err)

	rlnRelay := New(group_manager.Details{
		RLN:          rlnInstance,
		RootTracker:  rootTracker,
		GroupManager: gm,
	}, timesource.NewDefaultClock(), prometheus.DefaultRegisterer, utils.Logger())
	s.Require().NoError(err)

	err = rlnRelay.Start(context.TODO())
	s.Require().NoError(err)

	// wait for the event to reach the group handler
	time.Sleep(2 * time.Second)

	// rln pks are inserted into the rln peer's Merkle tree and the resulting root
	// is expected to be the same as the calculatedRoot i.e., the one calculated outside of the mountRlnRelayDynamic proc
	calculatedRoot, err := rlnRelay.RLN.GetMerkleRoot()
	s.Require().NoError(err)
	s.Require().Equal(expectedRoot, calculatedRoot)
}

func (s *WakuRLNRelayDynamicSuite) TestCorrectRegistrationOfPeers() {
	// Creating an RLN instance (just for generating membership keys)
	rlnInstance, err := rln.NewRLN()
	s.Require().NoError(err)

	// Node 1 ============================================================

	// Register credentials1 in contract and keystore1
	credentials1 := s.generateCredentials(rlnInstance)
	appKeystore, err := keystore.New(s.tmpKeystorePath(), dynamic.RLNAppInfo, utils.Logger())
	s.Require().NoError(err)
	membershipGroupIndex := s.register(appKeystore, credentials1, s.u1PrivKey)

	// mount the rln relay protocol in the on-chain/dynamic mode
	rootInstance, rootTracker, err := GetRLNInstanceAndRootTracker(s.tmpRLNDBPath())
	s.Require().NoError(err)
	gm1, err := dynamic.NewDynamicGroupManager(s.web3Config.ETHClientAddress, s.web3Config.RegistryContract.Address, &membershipGroupIndex, appKeystore, keystorePassword, prometheus.DefaultRegisterer, rootInstance, rootTracker, utils.Logger())
	s.Require().NoError(err)

	rlnRelay1 := New(group_manager.Details{
		GroupManager: gm1,
		RootTracker:  rootTracker,
		RLN:          rootInstance,
	}, timesource.NewDefaultClock(), prometheus.DefaultRegisterer, utils.Logger())

	err = rlnRelay1.Start(context.TODO())
	s.Require().NoError(err)

	// Node 2 ============================================================

	// Register credentials2 in contract and keystore2
	credentials2 := s.generateCredentials(rlnInstance)
	appKeystore2, err := keystore.New(s.tmpKeystorePath(), dynamic.RLNAppInfo, utils.Logger())
	s.Require().NoError(err)

	membershipGroupIndex = s.register(appKeystore2, credentials2, s.u2PrivKey)

	// mount the rln relay protocol in the on-chain/dynamic mode
	rootInstance, rootTracker, err = GetRLNInstanceAndRootTracker(s.tmpRLNDBPath())
	s.Require().NoError(err)
	gm2, err := dynamic.NewDynamicGroupManager(s.web3Config.ETHClientAddress, s.web3Config.RegistryContract.Address, &membershipGroupIndex, appKeystore2, keystorePassword, prometheus.DefaultRegisterer, rootInstance, rootTracker, utils.Logger())
	s.Require().NoError(err)

	rlnRelay2 := New(group_manager.Details{
		GroupManager: gm2,
		RootTracker:  rootTracker,
		RLN:          rootInstance,
	}, timesource.NewDefaultClock(), prometheus.DefaultRegisterer, utils.Logger())
	err = rlnRelay2.Start(context.TODO())
	s.Require().NoError(err)

	// ==================================
	// the two nodes should be registered into the contract
	// since nodes are spun up sequentially
	// the first node has index 0 whereas the second node gets index 1
	idx1 := rlnRelay1.GroupManager.MembershipIndex()
	idx2 := rlnRelay2.GroupManager.MembershipIndex()
	s.Require().NoError(err)
	s.Require().Equal(rln.MembershipIndex(0), idx1)
	s.Require().Equal(rln.MembershipIndex(1), idx2)
}

func (s *WakuRLNRelayDynamicSuite) tmpKeystorePath() string {
	keystoreDir, err := os.MkdirTemp("", "keystore_dir")
	s.Require().NoError(err)
	return filepath.Join(keystoreDir, "keystore.json")
}

func (s *WakuRLNRelayDynamicSuite) tmpRLNDBPath() string {
	dbPath, err := os.MkdirTemp("", "rln_db")
	s.Require().NoError(err)
	return dbPath
}