// nolint
//
//nolint:nolintlint
package dummydkg

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"io"

	unsafeRand "math/rand"

	"github.com/smartcontractkit/chainlink/system-tests/lib/cre/capabilities/vault/sanmarinodkg/dkg"

	"github.com/smartcontractkit/tdh2/go/tdh2/lib/group"
	"github.com/smartcontractkit/tdh2/go/tdh2/tdh2"
)

// Prepares the parameters for a new DKG instance.
// Passing the same parameters and seed will always yield the same instance ID and participant keys.
func NewDKGSetup(n int, t int, seed string) (dkg.InstanceID, dkg.ParticipantsConfig, []dkg.ParticipantSecretKey, error) {
	// Setup a deterministic random number generator, which depends on the passed parameter and the seed.
	rand := randFromSeed(fmt.Sprintf("DummyDKG(t=%d, n=%d, seed=%s)", t, n, seed))

	// Derive a unique instance ID for the DKG instance.
	id, err := randomInstanceID(rand)
	if err != nil {
		return "", dkg.ParticipantsConfig{}, nil, fmt.Errorf("failed to generate instance ID: %w", err)
	}

	// Setup the keys for all participants.
	// In a real-world scenario, these keys would be individually generated by each participant.
	sks := make([]dkg.ParticipantSecretKey, 0, n)
	pks := make([]dkg.ParticipantPublicKey, 0, n)
	for i := 0; i < n; i++ {
		sk, pk, err := keygen(rand)
		if err != nil {
			return "", dkg.ParticipantsConfig{}, nil, fmt.Errorf("failed to generate key pair: %w", err)
		}
		sks = append(sks, sk)
		pks = append(pks, pk)
	}

	// Prepare the recipients configuration for the DKG protocol.
	config := dkg.ParticipantsConfig{
		t - 1, // f... the number of faulty participants, not used for dummy implementation
		t,     // t... the minimal number of shares needed to recover the master secret key
		pks,   // the public keys of the participants, implicitly defines the number of participants (n)
	}

	return id, config, sks, nil
}

// Simulates the execution of a DKG protocol and returns its result. This demo implementation generates a DKG result
// locally, but follows the DKG interface definitions.
func NewDKGResult(instanceID dkg.InstanceID, recipientsConfig dkg.ParticipantsConfig, group group.Group) (dkg.Result, error) {
	randStream, err := randStream(randFromSeed(fmt.Sprintf("thd2-randStream-%s", instanceID)))
	if err != nil {
		return nil, fmt.Errorf("failed to initialized randStream: %w", err)
	}

	k := recipientsConfig.T
	n := len(recipientsConfig.PublicKeys)
	_, mpk, mskShares, err := tdh2.GenerateKeys(group, nil, k, n, randStream)
	if err != nil {
		return nil, err
	}

	return &tdh2DKGResult{instanceID, recipientsConfig, mpk, mskShares}, nil
}

// Sets up a new dummy key pair for a participant of the DKG protocol. The public keys of all participants need to be
// passed as configuration for running an actual DKG protocol, whereas the corresponding secret of each participant is
// required to obtain (=decrypt) its master secret key share from the DKG result.
// Note that, for the purpose of this demonstration code, no encryption/decryption is actually performed, but the
// behavior of the dkg.Result interface is accurately reflected.
func keygen(rand io.Reader) (dkg.ParticipantSecretKey, dkg.ParticipantPublicKey, error) {
	sk := make(dkg.ParticipantSecretKey, 32)
	if _, err := io.ReadFull(rand, sk); err != nil {
		return nil, nil, err
	}
	pk := make(dkg.ParticipantPublicKey, 32)
	for i := range pk {
		pk[i] = ^sk[i]
	}
	return sk, pk, nil
}

// Setup a deterministic random number generator based on a seed string.
func randFromSeed(seed string) io.Reader {
	sum := sha256.Sum256([]byte(seed))
	seedInt := int64(binary.LittleEndian.Uint64(sum[:8])) // take first 8 bytes
	return unsafeRand.New(unsafeRand.NewSource(seedInt))
}

// Derives a instance ID (128 bit hex string) for a DKG instance.
func randomInstanceID(rand io.Reader) (dkg.InstanceID, error) {
	buf := make([]byte, 16)
	_, err := io.ReadFull(rand, buf)
	if err != nil {
		return dkg.InstanceID(""), err
	}
	return dkg.InstanceID(fmt.Sprintf("%x", buf)), nil
}

// Internal type that implements the dkg.Result interface for a DKG protocol instance that generates a shared TDH2 key.
type tdh2DKGResult struct {
	instanceID       dkg.InstanceID
	recipientsConfig dkg.ParticipantsConfig
	mpk              *tdh2.PublicKey
	mpkShares        []*tdh2.PrivateShare
}

// Enforce that tdh2DKGResult indeed implements the dkg.Result interface.
var _ dkg.Result = &tdh2DKGResult{}

func (r *tdh2DKGResult) InstanceID() dkg.InstanceID {
	return r.instanceID
}

func (r *tdh2DKGResult) RecipientsConfig() dkg.ParticipantsConfig {
	return r.recipientsConfig
}

func (r *tdh2DKGResult) MasterPublicKey() (dkg.MasterPublicKey, error) {
	mpk, err := r.mpk.Marshal()
	if err != nil {
		return nil, err
	}
	return dkg.MasterPublicKey(mpk), nil
}

// THD2 does not use the concept of master public key shares independently, but rather the values are included in the
// tdh2.PublicKey struct (hArray field). We could unmarshal here and return the shares, but leave this unimplemented for
// now, as the MasterPublicKey() method is sufficient for the purpose of this demo code.
func (r *tdh2DKGResult) MasterPublicKeyShares() ([]dkg.MasterPublicKeyShare, error) {
	return nil, errors.New("not implemented, master public keys shares are returned via MasterPublicKey()")
}

func (r *tdh2DKGResult) MasterSecretKeyShare(sk dkg.ParticipantSecretKey) (dkg.MasterSecretKeyShare, error) {
	pk := make(dkg.ParticipantPublicKey, 32)
	for i := range pk {
		pk[i] = ^sk[i]
	}

	for i, storedPk := range r.recipientsConfig.PublicKeys {
		if bytes.Equal(storedPk, pk) {
			mskShare := r.mpkShares[i]
			mskShareBytes, err := mskShare.Marshal()
			if err != nil {
				return nil, err
			}
			return dkg.MasterSecretKeyShare(mskShareBytes), nil
		}
	}
	return nil, errors.New("failed to decrypt master secret key share (invalid decryption key?)")
}

// Based on tdh2/tdh2_test.go.
// A call to randStream(...) returns a cipher.Stream used for providing randomness in the context of TDH2.
func randStream(rand io.Reader) (cipher.Stream, error) {
	block, err := aes.NewCipher(make([]byte, 16))
	if err != nil {
		return nil, err
	}
	iv := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(rand, iv); err != nil {
		return nil, err
	}
	return cipher.NewCTR(block, iv), nil
}
