package testnet

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/ghodss/yaml"
	"github.com/pkg/errors"
	fastssz "github.com/prysmaticlabs/fastssz"
	"github.com/prysmaticlabs/prysm/v3/config/params"
	"github.com/prysmaticlabs/prysm/v3/io/file"
	ethpb "github.com/prysmaticlabs/prysm/v3/proto/prysm/v1alpha1"
	"github.com/prysmaticlabs/prysm/v3/runtime/interop"
	"github.com/sirupsen/logrus"
	"github.com/urfave/cli/v2"
)

var (
	generateGenesisStateFlags = struct {
		DepositJsonFile string
		ChainConfigFile string
		ConfigName      string
		NumValidators   uint64
		GenesisTime     uint64
		OutputSSZ       string
		OutputJSON      string
		OutputYaml      string
	}{}
	log           = logrus.WithField("prefix", "genesis")
	outputSSZFlag = &cli.StringFlag{
		Name:        "output-ssz",
		Destination: &generateGenesisStateFlags.OutputSSZ,
		Usage:       "Output filename of the SSZ marshaling of the generated genesis state",
		Value:       "",
	}
	outputYamlFlag = &cli.StringFlag{
		Name:        "output-yaml",
		Destination: &generateGenesisStateFlags.OutputYaml,
		Usage:       "Output filename of the YAML marshaling of the generated genesis state",
		Value:       "",
	}
	outputJsonFlag = &cli.StringFlag{
		Name:        "output-json",
		Destination: &generateGenesisStateFlags.OutputJSON,
		Usage:       "Output filename of the JSON marshaling of the generated genesis state",
		Value:       "",
	}
	generateGenesisStateCmd = &cli.Command{
		Name:   "generate-genesis",
		Usage:  "Generate a beacon chain genesis state",
		Action: cliActionGenerateGenesisState,
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:        "chain-config-file",
				Destination: &generateGenesisStateFlags.ChainConfigFile,
				Usage:       "The path to a YAML file with chain config values",
			},
			&cli.StringFlag{
				Name:        "deposit-json-file",
				Destination: &generateGenesisStateFlags.DepositJsonFile,
				Usage:       "Path to deposit_data.json file generated by the staking-deposit-cli tool for optionally specifying validators in genesis state",
			},
			&cli.StringFlag{
				Name:        "config-name",
				Usage:       "Config kind to be used for generating the genesis state. Default: mainnet. Options include mainnet, interop, minimal, prater, ropsten, sepolia. --chain-config-file will override this flag.",
				Destination: &generateGenesisStateFlags.ConfigName,
				Value:       params.MainnetName,
			},
			&cli.Uint64Flag{
				Name:        "num-validators",
				Usage:       "Number of validators to deterministically generate in the genesis state",
				Destination: &generateGenesisStateFlags.NumValidators,
				Required:    true,
			},
			&cli.Uint64Flag{
				Name:        "genesis-time",
				Destination: &generateGenesisStateFlags.GenesisTime,
				Usage:       "Unix timestamp seconds used as the genesis time in the genesis state. If unset, defaults to now()",
			},
			outputSSZFlag,
			outputYamlFlag,
			outputJsonFlag,
		},
	}
)

// Represents a json object of hex string and uint64 values for
// validators on Ethereum. This file can be generated using the official staking-deposit-cli.
type depositDataJSON struct {
	PubKey                string `json:"pubkey"`
	Amount                uint64 `json:"amount"`
	WithdrawalCredentials string `json:"withdrawal_credentials"`
	DepositDataRoot       string `json:"deposit_data_root"`
	Signature             string `json:"signature"`
}

func cliActionGenerateGenesisState(cliCtx *cli.Context) error {
	if generateGenesisStateFlags.GenesisTime == 0 {
		log.Info("No genesis time specified, defaulting to now()")
	}
	outputJson := generateGenesisStateFlags.OutputJSON
	outputYaml := generateGenesisStateFlags.OutputYaml
	outputSSZ := generateGenesisStateFlags.OutputSSZ
	noOutputFlag := outputSSZ == "" && outputJson == "" && outputYaml == ""
	if noOutputFlag {
		return fmt.Errorf(
			"no %s, %s, %s flag(s) specified. At least one is required",
			outputJsonFlag.Name,
			outputYamlFlag.Name,
			outputSSZFlag.Name,
		)
	}
	if err := setGlobalParams(); err != nil {
		return fmt.Errorf("could not set config params: %v", err)
	}
	genesisState, err := generateGenesis(cliCtx.Context)
	if err != nil {
		return fmt.Errorf("could not generate genesis state: %v", err)
	}
	if outputJson != "" {
		if err := writeToOutputFile(outputJson, genesisState, json.Marshal); err != nil {
			return err
		}
	}
	if outputYaml != "" {
		if err := writeToOutputFile(outputYaml, genesisState, yaml.Marshal); err != nil {
			return err
		}
	}
	if outputSSZ != "" {
		marshalFn := func(o interface{}) ([]byte, error) {
			marshaler, ok := o.(fastssz.Marshaler)
			if !ok {
				return nil, errors.New("not a marshaler")
			}
			return marshaler.MarshalSSZ()
		}
		if err := writeToOutputFile(outputSSZ, genesisState, marshalFn); err != nil {
			return err
		}
	}
	log.Info("Command completed")
	return nil
}

func setGlobalParams() error {
	chainConfigFile := generateGenesisStateFlags.ChainConfigFile
	if chainConfigFile != "" {
		log.Infof("Specified a chain config file: %s", chainConfigFile)
		return params.LoadChainConfigFile(chainConfigFile, nil)
	}
	cfg, err := params.ByName(generateGenesisStateFlags.ConfigName)
	if err != nil {
		return fmt.Errorf("unable to find config using name %s: %v", generateGenesisStateFlags.ConfigName, err)
	}
	return params.SetActive(cfg.Copy())
}

func generateGenesis(ctx context.Context) (*ethpb.BeaconState, error) {
	genesisTime := generateGenesisStateFlags.GenesisTime
	numValidators := generateGenesisStateFlags.NumValidators
	depositJsonFile := generateGenesisStateFlags.DepositJsonFile
	if depositJsonFile != "" {
		expanded, err := file.ExpandPath(depositJsonFile)
		if err != nil {
			return nil, err
		}
		inputJSON, err := os.Open(expanded) // #nosec G304
		if err != nil {
			return nil, err
		}
		defer func() {
			if err := inputJSON.Close(); err != nil {
				log.WithError(err).Printf("Could not close file %s", depositJsonFile)
			}
		}()
		log.Printf("Generating genesis state from input JSON deposit data %s", depositJsonFile)
		return genesisStateFromJSONValidators(ctx, inputJSON, genesisTime)
	}
	if numValidators == 0 {
		return nil, fmt.Errorf(
			"expected --num-validators > 0 to have been provided",
		)
	}
	// If no JSON input is specified, we create the state deterministically from interop keys.
	genesisState, _, err := interop.GenerateGenesisState(ctx, genesisTime, numValidators)
	if err != nil {
		return nil, err
	}
	return genesisState, err
}

func genesisStateFromJSONValidators(ctx context.Context, r io.Reader, genesisTime uint64) (*ethpb.BeaconState, error) {
	enc, err := io.ReadAll(r)
	if err != nil {
		return nil, err
	}
	var depositJSON []*depositDataJSON
	if err := json.Unmarshal(enc, &depositJSON); err != nil {
		return nil, err
	}
	depositDataList := make([]*ethpb.Deposit_Data, len(depositJSON))
	depositDataRoots := make([][]byte, len(depositJSON))
	for i, val := range depositJSON {
		data, dataRootBytes, err := depositJSONToDepositData(val)
		if err != nil {
			return nil, err
		}
		depositDataList[i] = data
		depositDataRoots[i] = dataRootBytes
	}
	beaconState, _, err := interop.GenerateGenesisStateFromDepositData(ctx, genesisTime, depositDataList, depositDataRoots)
	if err != nil {
		return nil, err
	}
	return beaconState, nil
}

func depositJSONToDepositData(input *depositDataJSON) (depositData *ethpb.Deposit_Data, dataRoot []byte, err error) {
	pubKeyBytes, err := hex.DecodeString(strings.TrimPrefix(input.PubKey, "0x"))
	if err != nil {
		return
	}
	withdrawalbytes, err := hex.DecodeString(strings.TrimPrefix(input.WithdrawalCredentials, "0x"))
	if err != nil {
		return
	}
	signatureBytes, err := hex.DecodeString(strings.TrimPrefix(input.Signature, "0x"))
	if err != nil {
		return
	}
	dataRootBytes, err := hex.DecodeString(strings.TrimPrefix(input.DepositDataRoot, "0x"))
	if err != nil {
		return
	}
	depositData = &ethpb.Deposit_Data{
		PublicKey:             pubKeyBytes,
		WithdrawalCredentials: withdrawalbytes,
		Amount:                input.Amount,
		Signature:             signatureBytes,
	}
	dataRoot = dataRootBytes
	return
}

func writeToOutputFile(
	fPath string,
	data interface{},
	marshalFn func(o interface{}) ([]byte, error),
) error {
	encoded, err := marshalFn(data)
	if err != nil {
		return err
	}
	if err := file.WriteFile(fPath, encoded); err != nil {
		return err
	}
	log.Printf("Done writing genesis state to %s", fPath)
	return nil
}