package secrets

import (
	"encoding/hex"
	"encoding/json"
	"fmt"
	"slices"

	"github.com/ethereum/go-ethereum/common"
	"github.com/gagliardetto/solana-go"
	"github.com/pelletier/go-toml/v2"
	"github.com/pkg/errors"

	chainselectors "github.com/smartcontractkit/chain-selectors"
	"github.com/smartcontractkit/smdkg/dkgocr/dkgocrtypes"

	libc "github.com/smartcontractkit/chainlink/system-tests/lib/conversions"
	"github.com/smartcontractkit/chainlink/system-tests/lib/cre"
	"github.com/smartcontractkit/chainlink/system-tests/lib/cre/don/node"
	"github.com/smartcontractkit/chainlink/system-tests/lib/crypto"
)

type NodeEthKey struct {
	JSON     string `toml:"JSON"`
	Password string `toml:"Password"`
	ChainID  int    `toml:"ID"`
}

type NodeSolKey struct {
	JSON     string `toml:"JSON"`
	Password string `toml:"Password"`
	ChainID  string `toml:"ID"`
}

type NodeP2PKey struct {
	JSON     string `toml:"JSON"`
	Password string `toml:"Password"`
}

type NodeDKGRecipientKey struct {
	JSON     string `toml:"JSON"`
	Password string `toml:"Password"`
}
type NodeEthKeyWrapper struct {
	EthKeys []NodeEthKey `toml:"Keys"`
}

type NodeSolKeyWrapper struct {
	SolKeys []NodeSolKey `toml:"Keys"`
}

type NodeSecret struct {
	EthKeys         NodeEthKeyWrapper   `toml:"EVM"`
	SolKeys         NodeSolKeyWrapper   `toml:"Solana"`
	P2PKey          NodeP2PKey          `toml:"P2PKey"`
	DKGRecipientKey NodeDKGRecipientKey `toml:"DKGRecipientKey"`

	// Add more fields as needed to reflect 'Secrets' struct from /core/config/toml/types.go
	// We can't use the original struct, because it's using custom types that serialize secrets to 'xxxxx'
}

func GenerateSecrets(input *cre.GenerateSecretsInput) (cre.NodeIndexToSecretsOverride, error) {
	if input == nil {
		return nil, errors.New("input is nil")
	}
	if err := input.Validate(); err != nil {
		return nil, errors.Wrap(err, "input validation failed")
	}

	overrides := make(cre.NodeIndexToSecretsOverride)

	for i := range input.DonMetadata.NodesMetadata {
		nodeSecret := NodeSecret{}
		if input.EVMKeys != nil {
			nodeSecret.EthKeys = NodeEthKeyWrapper{}
			for chainID, evmKeys := range input.EVMKeys {
				nodeSecret.EthKeys.EthKeys = append(nodeSecret.EthKeys.EthKeys, NodeEthKey{
					JSON:     string(evmKeys.EncryptedJSONs[i]),
					Password: evmKeys.Password,
					ChainID:  chainID,
				})
			}
		}

		if input.DKGRecipientKeys != nil {
			nodeSecret.DKGRecipientKey = NodeDKGRecipientKey{
				JSON:     string(input.DKGRecipientKeys.EncryptedJSONs[i]),
				Password: input.DKGRecipientKeys.Password,
			}
		}

		if input.P2PKeys != nil {
			nodeSecret.P2PKey = NodeP2PKey{
				JSON:     string(input.P2PKeys.EncryptedJSONs[i]),
				Password: input.P2PKeys.Password,
			}
		}

		if input.SolKeys != nil {
			nodeSecret.SolKeys = NodeSolKeyWrapper{}
			for chainID, solKeys := range input.SolKeys {
				nodeSecret.SolKeys.SolKeys = append(nodeSecret.SolKeys.SolKeys, NodeSolKey{
					JSON:     string(solKeys.EncryptedJSONs[i]),
					Password: solKeys.Password,
					ChainID:  chainID,
				})
			}
		}

		nodeSecretString, err := toml.Marshal(nodeSecret)
		if err != nil {
			return nil, errors.Wrap(err, "failed to marshal node secrets")
		}

		overrides[i] = string(nodeSecretString)
	}

	return overrides, nil
}

func AddKeysToTopology(topology *cre.Topology, keys *cre.GenerateKeysOutput) (*cre.Topology, error) {
	if topology == nil {
		return nil, errors.New("topology is nil")
	}

	if keys == nil {
		return nil, errors.New("keys is nil")
	}

	if len(keys.P2PKeys) != len(topology.DonsMetadata) {
		return nil, fmt.Errorf("number of P2P keys does not match the number of DONs. Expected %d, got %d", len(topology.DonsMetadata), len(keys.P2PKeys))
	}

	if len(keys.DKGRecipientKeys) != len(topology.DonsMetadata) {
		return nil, fmt.Errorf("number of DKG recipient keys does not match the number of DONs. Expected %d, got %d", len(topology.DonsMetadata), len(keys.DKGRecipientKeys))
	}

	if len(keys.EVMKeys) != len(topology.DonsMetadata) {
		return nil, fmt.Errorf("number of EVM keys does not match the number of DONs. Expected %d, got %d", len(topology.DonsMetadata), len(keys.EVMKeys))
	}

	for _, donMetadata := range topology.DonsMetadata {
		if _, ok := keys.P2PKeys[donMetadata.ID]; !ok {
			return nil, fmt.Errorf("no P2P keys found for DON %d", donMetadata.ID)
		}

		p2pKeys := keys.P2PKeys[donMetadata.ID]
		if len(p2pKeys.PeerIDs) != len(donMetadata.NodesMetadata) {
			return nil, fmt.Errorf("number of P2P keys for DON %d does not match the number of nodes. Expected %d, got %d", donMetadata.ID, len(donMetadata.NodesMetadata), len(p2pKeys.PeerIDs))
		}
		for idx, nodeMetadata := range donMetadata.NodesMetadata {
			nodeMetadata.Labels = append(nodeMetadata.Labels, &cre.Label{
				Key:   node.NodeP2PIDKey,
				Value: p2pKeys.PeerIDs[idx],
			})
		}

		if _, ok := keys.DKGRecipientKeys[donMetadata.ID]; !ok {
			return nil, fmt.Errorf("no DKG recipient keys found for DON %d", donMetadata.ID)
		}
		dkgRecipientKeys := keys.DKGRecipientKeys[donMetadata.ID]
		if len(dkgRecipientKeys.PubKeys) != len(donMetadata.NodesMetadata) {
			return nil, fmt.Errorf("number of DKG recipient keys for DON %d does not match the number of nodes. Expected %d, got %d", donMetadata.ID, len(donMetadata.NodesMetadata), len(dkgRecipientKeys.PubKeys))
		}
		for idx, nodeMetadata := range donMetadata.NodesMetadata {
			nodeMetadata.Labels = append(nodeMetadata.Labels, &cre.Label{
				Key:   cre.NodeDKGRecipientKey,
				Value: common.Bytes2Hex(dkgRecipientKeys.PubKeys[idx]),
			})
		}

		if _, ok := keys.EVMKeys[donMetadata.ID]; !ok {
			return nil, fmt.Errorf("no EVM keys found for DON %d", donMetadata.ID)
		}

		chainIDsToEVMKeys := keys.EVMKeys[donMetadata.ID]

		// Now add the EVM addresses to the node metadata
		for chainID, evmKeys := range chainIDsToEVMKeys {
			chainSelector, selectorErr := chainselectors.SelectorFromChainId(libc.MustSafeUint64(int64(chainID)))
			if selectorErr != nil {
				return nil, errors.Wrapf(selectorErr, "failed to get chain selector for chain ID %d", chainID)
			}
			if len(evmKeys.PublicAddresses) != len(donMetadata.NodesMetadata) {
				return nil, fmt.Errorf("number of EVM keys for DON %d and chain ID %d does not match the number of nodes. Expected %d, got %d", donMetadata.ID, chainID, len(donMetadata.NodesMetadata), len(evmKeys.PublicAddresses))
			}

			for idx, nodeMetadata := range donMetadata.NodesMetadata {
				nodeMetadata.Labels = append(nodeMetadata.Labels, &cre.Label{
					Key:   node.AddressKeyFromSelector(chainSelector),
					Value: evmKeys.PublicAddresses[idx].Hex(),
				})
			}
		}

		chainIDsToSolKeys, ok := keys.SolKeys[donMetadata.ID]
		if !ok {
			continue
		}

		for chainID, solKeys := range chainIDsToSolKeys {
			selector, ok := chainselectors.SolanaChainIdToChainSelector()[chainID]
			if !ok {
				return nil, errors.New("selector not found for solana chainID " + chainID)
			}

			if len(solKeys.PublicAddresses) != len(donMetadata.NodesMetadata) {
				return nil, fmt.Errorf("number of sol keys for DON %d and chainID %s doesn't match the number of Nodes. Expected %d got %d",
					donMetadata.ID, chainID, len(donMetadata.NodesMetadata), len(solKeys.PublicAddresses))
			}
			for idx, nodeMetadata := range donMetadata.NodesMetadata {
				nodeMetadata.Labels = append(nodeMetadata.Labels, &cre.Label{
					Key:   node.AddressKeyFromSelector(selector),
					Value: solKeys.PublicAddresses[idx].String(),
				})
			}
		}
	}

	return topology, nil
}

// secrets struct mirrors `Secrets` struct in "github.com/smartcontractkit/chainlink/v2/core/config/toml"
// we use a copy to avoid depending on the core config package, we consider it safe, because that struct changes very rarely
type secrets struct {
	EVM             ethKeys         `toml:",omitempty"` // choose EVM as the TOML field name to align with relayer config convention
	P2PKey          p2PKey          `toml:",omitempty"`
	Solana          solKeys         `toml:",omitempty"`
	DKGRecipientKey dkgRecipientKey `toml:",omitempty"`
}

type p2PKey struct {
	JSON     *string
	Password *string
}

type dkgRecipientKey struct {
	JSON     *string
	Password *string
}

type ethKeys struct {
	Keys []*ethKey
}

type ethKey struct {
	JSON     *string
	ID       *int
	Password *string
}

type solKeys struct {
	Keys []*solKey
}

type solKey struct {
	JSON     *string
	ID       *string
	Password *string
}

// struct required for reading "address" from this bit of encrypted JSON:
// JSON = '{"address":"e753ac0b6e175ce3a939c55433a0109c5a6f8777"}'
type evmJSON struct {
	Address string `json:"address"`
}

var publicEVMAddressFromEncryptedJSON = func(jsonString string) (string, error) {
	var eJSON evmJSON
	err := json.Unmarshal([]byte(jsonString), &eJSON)
	if err != nil {
		return "", errors.Wrap(err, "failed to unmarshal evm json")
	}
	return eJSON.Address, nil
}

// struct required for reading "address" from this bit of encrypted JSON:
// JSON = '{"publicKey":"22b4b2618de6dc8254d76276d51f6a9d53471d5b2465c8cae237f21425b10b7d"}'
type solJSON struct {
	PublicKey string `json:"publicKey"`
}

var publicSolKeyFromEncryptedJSON = func(jsonString string) (solana.PublicKey, error) {
	var eJSON solJSON
	err := json.Unmarshal([]byte(jsonString), &eJSON)
	if err != nil {
		return solana.PublicKey{}, errors.Wrap(err, "failed to unmarshal Solana json")
	}

	b, err := hex.DecodeString(eJSON.PublicKey)
	if err != nil {
		return solana.PublicKey{}, errors.Wrap(err, "invalid hex string for public key solana")
	}

	return solana.PublicKeyFromBytes(b), nil
}

// struct required for reading "peerID" from this bit of encrypted JSON:
// JSON = '{"keyType":"P2P","publicKey":"f3c458c9064bdde449a3904ba8d3f8f5ebf79623077430325252c3368f920199","peerID":"p2p_12D3KooWSDvtYVF3FoyGeMrmDxYeJZMzbEyMHRwmf5GUSqgJhST2"}'
type p2pJSON struct {
	PeerID string `json:"peerID"`
}

var publicP2PAddressFromEncryptedJSON = func(jsonString string) (string, error) {
	var pJSON p2pJSON
	err := json.Unmarshal([]byte(jsonString), &pJSON)
	if err != nil {
		return "", errors.Wrap(err, "failed to unmarshal p2p json")
	}
	return pJSON.PeerID, nil
}

type dkgRecipientJSON struct {
	PublicKey string `json:"publicKey"`
}

var publicDKGRecipientKeyFromEncryptedJSON = func(jsonString string) (dkgocrtypes.P256ParticipantPublicKey, error) {
	var dJSON dkgRecipientJSON
	err := json.Unmarshal([]byte(jsonString), &dJSON)
	if err != nil {
		return nil, errors.Wrap(err, "failed to unmarshal dkg recipient key json")
	}
	return hex.DecodeString(dJSON.PublicKey)
}

func KeysOutputFromConfig(nodeSets []*cre.CapabilitiesAwareNodeSet) (*cre.GenerateKeysOutput, error) {
	output := &cre.GenerateKeysOutput{
		EVMKeys:          make(cre.DonsToEVMKeys),
		SolKeys:          make(cre.DonsToSolKeys),
		P2PKeys:          make(cre.DonsToP2PKeys),
		DKGRecipientKeys: make(cre.DonsToDKGRecipientKeys),
	}

	p2pKeysFoundPerDon := make(map[uint32]int)
	dkgRecipientKeysFoundPerDon := make(map[uint32]int)
	evmKeysFoundPerDon := make(map[uint32]int)
	for donIdx, nodeSet := range nodeSets {
		donIdxUint32 := uint32(donIdx) // #nosec G115: ignore as this will NEVER happen, we don't have zillions of DONs
		p2pKeys := crypto.P2PKeys{}
		dkgRecipientKeys := crypto.DKGRecipientKeys{}
		evmKeysPerChainID := make(cre.ChainIDToEVMKeys)
		solKeysPerChainID := make(cre.ChainIDToSolKeys)
		for nodeIdx, nodeSpec := range nodeSet.NodeSpecs {
			if nodeSpec.Node.TestSecretsOverrides != "" {
				var sSecrets secrets
				unmarshallErr := toml.Unmarshal([]byte(nodeSpec.Node.TestSecretsOverrides), &sSecrets)
				if unmarshallErr != nil {
					return nil, errors.Wrapf(unmarshallErr, "failed to unmarshal secrets for node %d in DON %d", nodeIdx, donIdx)
				}

				// For simplicity we will allow importing only both P2P keys and EVM keys, not just one of them
				if sSecrets.P2PKey.JSON == nil || sSecrets.P2PKey.Password == nil {
					return nil, fmt.Errorf("P2P key or password is nil for node %d in DON %d", nodeIdx, donIdx)
				}
				p2pKeys.EncryptedJSONs = append(p2pKeys.EncryptedJSONs, []byte(*sSecrets.P2PKey.JSON))
				p2pKeys.Password = *sSecrets.P2PKey.Password
				peerID, peerIDErr := publicP2PAddressFromEncryptedJSON(*sSecrets.P2PKey.JSON)
				if peerIDErr != nil {
					return nil, errors.Wrapf(peerIDErr, "failed to get public p2p address for node %d in DON %d from encrypted JSON", nodeIdx, donIdx)
				}
				p2pKeys.PeerIDs = append(p2pKeys.PeerIDs, peerID)
				p2pKeysFoundPerDon[donIdxUint32]++

				if sSecrets.DKGRecipientKey.JSON != nil {
					dkgRecipientKeys.EncryptedJSONs = append(dkgRecipientKeys.EncryptedJSONs, []byte(*sSecrets.DKGRecipientKey.JSON))
					dkgRecipientKeys.Password = *sSecrets.DKGRecipientKey.Password
					dkgRecipientPubKey, err := publicDKGRecipientKeyFromEncryptedJSON(*sSecrets.DKGRecipientKey.JSON)
					if err != nil {
						return nil, errors.Wrapf(err, "failed to get public DKG recipient key for node %d in DON %d from encrypted JSON", nodeIdx, donIdx)
					}
					dkgRecipientKeys.PubKeys = append(dkgRecipientKeys.PubKeys, dkgRecipientPubKey)
					dkgRecipientKeysFoundPerDon[donIdxUint32]++
				}

				if len(sSecrets.EVM.Keys) == 0 {
					return nil, fmt.Errorf("EVM keys is nil for node %d in DON %d", nodeIdx, donIdx)
				}

				for _, evmKey := range sSecrets.EVM.Keys {
					if evmKey.JSON == nil || evmKey.Password == nil || evmKey.ID == nil {
						return nil, fmt.Errorf("EVM key or password or ID is nil for node %d in DON %d", nodeIdx, donIdx)
					}

					publicEVMAddress, publicEVMAddressErr := publicEVMAddressFromEncryptedJSON(*evmKey.JSON)
					if publicEVMAddressErr != nil {
						return nil, errors.Wrapf(publicEVMAddressErr, "failed to get public evm address for node %d in DON %d from encrypted JSON", nodeIdx, donIdx)
					}

					// if the DON doesn't support the chain, we skip it; if slice is empty, it means that the DON supports all chains
					if len(nodeSet.SupportedChains) > 0 && !slices.Contains(nodeSet.SupportedChains, libc.MustSafeUint64(int64(*evmKey.ID))) {
						continue
					}

					if _, ok := evmKeysPerChainID[*evmKey.ID]; !ok {
						evmKeysPerChainID[*evmKey.ID] = &crypto.EVMKeys{}
					}

					evmKeysPerChainID[*evmKey.ID].EncryptedJSONs = append(evmKeysPerChainID[*evmKey.ID].EncryptedJSONs, []byte(*evmKey.JSON))
					evmKeysPerChainID[*evmKey.ID].PublicAddresses = append(evmKeysPerChainID[*evmKey.ID].PublicAddresses, common.HexToAddress(publicEVMAddress))
					evmKeysPerChainID[*evmKey.ID].Password = *evmKey.Password
				}
				evmKeysFoundPerDon[donIdxUint32]++

				for _, solKey := range sSecrets.Solana.Keys {
					if solKey.JSON == nil || solKey.Password == nil || solKey.ID == nil {
						return nil, fmt.Errorf("solana key or password or id is nil for node %d in don %d", nodeIdx, donIdx)
					}
					// if the DON doesn't support solana chain, we skip it
					if !slices.Contains(nodeSet.SupportedSolChains, *solKey.ID) {
						continue
					}

					publicSolAddr, addrErr := publicSolKeyFromEncryptedJSON(*solKey.JSON)
					if addrErr != nil {
						return nil, errors.Wrapf(addrErr, "failed to get public Solana address for node %d in DON %d from encrypted JSON", nodeIdx, donIdx)
					}
					if _, ok := solKeysPerChainID[*solKey.ID]; !ok {
						solKeysPerChainID[*solKey.ID] = &crypto.SolKeys{}
					}

					solKeysPerChainID[*solKey.ID].EncryptedJSONs = append(solKeysPerChainID[*solKey.ID].EncryptedJSONs, []byte(*solKey.JSON))
					solKeysPerChainID[*solKey.ID].PublicAddresses = append(solKeysPerChainID[*solKey.ID].PublicAddresses,
						publicSolAddr)
					solKeysPerChainID[*solKey.ID].Password = *solKey.Password
				}
			}
		}
		// +1 because we use 1-based indexing in the CRE
		donIndexToUse := libc.MustSafeUint64FromInt(donIdx + 1)
		output.P2PKeys[donIndexToUse] = &p2pKeys
		output.DKGRecipientKeys[donIndexToUse] = &dkgRecipientKeys
		output.EVMKeys[donIndexToUse] = evmKeysPerChainID
		output.SolKeys[donIndexToUse] = solKeysPerChainID
	}

	anyFound := false
	// Validate that we found keys for all nodes in all DONs
	for donIdx, nodeSet := range nodeSets {
		donIdxUint32 := uint32(donIdx) // #nosec G115
		if p2pKeysFoundPerDon[donIdxUint32] != 0 && len(nodeSet.NodeSpecs) != p2pKeysFoundPerDon[donIdxUint32] {
			return nil, fmt.Errorf("number of P2P keys found for DON %d does not match the number of nodes. Expected %d, got %d", donIdx, len(nodeSet.NodeSpecs), p2pKeysFoundPerDon[donIdxUint32])
		}
		if evmKeysFoundPerDon[donIdxUint32] != 0 && len(nodeSet.NodeSpecs) != evmKeysFoundPerDon[donIdxUint32] {
			return nil, fmt.Errorf("number of EVM keys found for DON %d does not match the number of nodes. Expected %d, got %d", donIdx, len(nodeSet.NodeSpecs), evmKeysFoundPerDon[donIdxUint32])
		}
		if p2pKeysFoundPerDon[donIdxUint32] != 0 && evmKeysFoundPerDon[donIdxUint32] != 0 {
			anyFound = true
		}
	}

	if !anyFound {
		// If no keys were found for any DON, we can return empty output
		return nil, nil
	}

	return output, nil
}

func GenerateKeys(input *cre.GenerateKeysInput) (*cre.GenerateKeysOutput, error) {
	if input == nil {
		return nil, errors.New("input is nil")
	}

	if err := input.Validate(); err != nil {
		return nil, errors.Wrap(err, "input validation failed")
	}

	if input.Out != nil {
		return input.Out, nil
	}

	output := &cre.GenerateKeysOutput{
		EVMKeys:          make(cre.DonsToEVMKeys),
		SolKeys:          make(cre.DonsToSolKeys),
		P2PKeys:          make(cre.DonsToP2PKeys),
		DKGRecipientKeys: make(cre.DonsToDKGRecipientKeys),
	}

	for _, donMetadata := range input.Topology.DonsMetadata {
		if input.GenerateP2PKeys {
			p2pKeys, err := crypto.GenerateP2PKeys(input.Password, len(donMetadata.NodesMetadata))
			if err != nil {
				return nil, errors.Wrap(err, "failed to generate P2P keys")
			}
			output.P2PKeys[donMetadata.ID] = p2pKeys
		}
		if input.GenerateDKGRecipientKeys {
			dkgKeys, err := crypto.GenerateDKGRecipientKeys(input.Password, len(donMetadata.NodesMetadata))
			if err != nil {
				return nil, errors.Wrap(err, "failed to generate DKG recipient keys")
			}
			output.DKGRecipientKeys[donMetadata.ID] = dkgKeys
		}

		if len(input.GenerateEVMKeysForChainIDs) > 0 {
			for _, chainID := range input.GenerateEVMKeysForChainIDs {
				// if the DON doesn't support the chain, we skip it; if slice is empty, it means that the DON supports all chains
				if len(donMetadata.SupportedChains) > 0 && !slices.Contains(donMetadata.SupportedChains, libc.MustSafeUint64(int64(chainID))) {
					continue
				}
				evmKeys, err := crypto.GenerateEVMKeys(input.Password, len(donMetadata.NodesMetadata))
				if err != nil {
					return nil, errors.Wrap(err, "failed to generate EVM keys")
				}
				if _, ok := output.EVMKeys[donMetadata.ID]; !ok {
					output.EVMKeys[donMetadata.ID] = make(cre.ChainIDToEVMKeys)
				}
				output.EVMKeys[donMetadata.ID][chainID] = evmKeys
			}
		}

		for _, chainID := range input.GenerateSolKeysForChainIDs {
			solKeys, err := crypto.GenerateSolKeys(input.Password, len(donMetadata.NodesMetadata))
			if err != nil {
				return nil, errors.Wrap(err, "failed to generate Sol keys")
			}
			if _, ok := output.SolKeys[donMetadata.ID]; !ok {
				output.SolKeys[donMetadata.ID] = make(cre.ChainIDToSolKeys)
			}

			output.SolKeys[donMetadata.ID][chainID] = solKeys
		}
	}

	return output, nil
}
