package crypto

import (
	"crypto/hmac"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/binary"
	"fmt"
	"math/big"

	"github.com/cosmos/go-bip39"
)

var (
	DomainSeparator   = []byte("Nockchain seed")
	PrivateKeyStart   = []byte{4, 178, 67, 11}
	PublicKeyStart    = []byte{234, 230, 92}
	MagicDyckForPoint = []uint64{0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1}
	MagicDyckForT8    = []uint64{0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1}
	Tip5Zero          = [5]uint64{1730770831742798981, 2676322185709933211, 8329210750824781744, 16756092452590401876, 3547445316740171466}
	Tip5One           = [5]uint64{6727110957294540849, 15606243244732609007, 11887284596344881785, 10646863421881571398, 8146872807338919620}
	Tip5ZeroZero      = [5]uint64{4372149332062030091, 17876920912185183887, 13348798570422431948, 8872865212694716527, 3385176510443841516}
)

type MasterKey struct {
	PrivateKey []byte
	PublicKey  []byte
	ChainCode  []byte
}

func (m *MasterKey) DeriveChild(index uint64) (MasterKey, error) {
	idxBytes := make([]byte, 4)
	binary.BigEndian.PutUint32(idxBytes, uint32(index))
	if len(m.PrivateKey) == 0 {
		// Derive public key to child
		if index > 1<<31 {
			return MasterKey{}, fmt.Errorf("public keys can't be hardened")
		}
		data := m.PublicKey
		data = append(data, idxBytes...)

		hash := hmac.New(sha512.New, m.ChainCode)
		hash.Write(data)
		hashed := hash.Sum(nil)
		left := hashed[:32]
		right := hashed[32:]
		leftBigInt := new(big.Int).SetBytes(left)

		pkPoint, err := CheetaPointFromBytes(m.PublicKey)
		if err != nil {
			return MasterKey{}, err
		}

		keyPoint := CheetahAdd(CheetahScaleBig(A_GEN, *leftBigInt), pkPoint)
		for {
			if leftBigInt.Cmp(G_ORDER) < 0 {
				break
			} else {
				data = []byte{0x01}
				data = append(data, right...)
				data = append(data, idxBytes...)

				hash := hmac.New(sha512.New, m.ChainCode)
				hash.Write(data)
				hashed = hash.Sum(nil)
				left = hashed[:32]
				right = hashed[32:]

				leftBigInt = new(big.Int).SetBytes(left)
				keyPoint = CheetahAdd(CheetahScaleBig(A_GEN, *leftBigInt), pkPoint)

			}
		}
		return MasterKey{
			PrivateKey: []byte{},
			PublicKey:  keyPoint.Bytes(),
			ChainCode:  right,
		}, nil
	} else {
		hash := hmac.New(sha512.New, m.ChainCode)
		var data []byte
		if index > 1<<31 {
			//  hardened
			data = []byte{0x00}
			data = append(data, m.PrivateKey...)

			data = append(data, idxBytes...)
		} else {
			data = m.PublicKey
			data = append(data, idxBytes...)
		}
		hash.Write(data)
		hashed := hash.Sum(nil)
		left := hashed[:32]
		right := hashed[32:]

		leftBigInt := new(big.Int).SetBytes(left)
		privBigInt := new(big.Int).SetBytes(m.PrivateKey)
		sum := new(big.Int).Add(leftBigInt, privBigInt)
		keyBigInt := new(big.Int).Mod(sum, G_ORDER)
		for {
			if leftBigInt.Cmp(G_ORDER) < 0 {
				break
			} else {
				data = []byte{0x01}
				data = append(data, right...)
				data = append(data, idxBytes...)

				hash := hmac.New(sha512.New, m.ChainCode)
				hash.Write(data)
				hashed = hash.Sum(nil)
				left = hashed[:32]
				right = hashed[32:]

				leftBigInt = new(big.Int).SetBytes(left)
				sum = new(big.Int).Add(leftBigInt, privBigInt)
				keyBigInt = new(big.Int).Mod(sum, G_ORDER)
			}
		}

		pkPoint := CheetahScaleBig(A_GEN, *keyBigInt)
		return MasterKey{
			PrivateKey: keyBigInt.Bytes(),
			PublicKey:  pkPoint.Bytes(),
			ChainCode:  right,
		}, nil
	}
}

func MasterKeyFromSeed(seed string) (*MasterKey, error) {
	seedBytes, err := bip39.NewSeedWithErrorChecking(seed, "")
	if err != nil {
		return nil, fmt.Errorf("failed to parse seed: %v", err)
	}
	hash := hmac.New(sha512.New, DomainSeparator)
	hash.Write(seedBytes)
	hashedSeed := hash.Sum(nil)

	left := hashedSeed[:32]
	right := hashedSeed[32:]

	leftBigInt := new(big.Int).SetBytes(left)
	for {
		if leftBigInt.Cmp(G_ORDER) < 0 {
			break
		} else {
			hash := hmac.New(sha512.New, DomainSeparator)
			hash.Write(hashedSeed)
			hashedSeed = hash.Sum(nil)
			left = hashedSeed[:32]
			right = hashedSeed[32:]

			leftBigInt = new(big.Int).SetBytes(left)
		}
	}

	pkPoint := CheetahScaleBig(A_GEN, *leftBigInt)
	return &MasterKey{
		PrivateKey: left,
		PublicKey:  pkPoint.Bytes(),
		ChainCode:  right,
	}, nil
}

func MasterKeyFromPrivKey(chainCode, key []byte) (*MasterKey, error) {
	keyBigInt := new(big.Int).SetBytes(key)
	pkPoint := CheetahScaleBig(A_GEN, *keyBigInt)

	return &MasterKey{
		PrivateKey: key,
		PublicKey:  pkPoint.Bytes(),
		ChainCode:  chainCode,
	}, nil
}

func PrivKeyToT8(data []byte) [8]uint64 {
	dataBigInt := new(big.Int).SetBytes(data)
	res := rip(*dataBigInt)
	return [8]uint64(res)
}

func BigIntToT8(data big.Int) [8]uint64 {
	res := rip(data)
	return [8]uint64(res)
}
func SerializeExtend(chainCode []byte, key []byte, version []byte) []byte {
	data := version
	// dep:  depth in chain
	// idx:  index at depth
	// pf:  parent fingerprint
	depth := 0
	idx := []byte{0, 0, 0, 0}
	pf := []byte{0, 0, 0, 0}
	data = append(data, byte(depth%256))
	data = append(data, pf...)
	data = append(data, idx...)
	data = append(data, chainCode...)
	data = append(data, key...)
	return AddChecksum(data)
}
func AddChecksum(data []byte) []byte {
	hash1 := sha256.Sum256(data)
	hash2 := sha256.Sum256(hash1[:])

	checksum := hash2[:4]
	data = append(data, checksum...)

	return data
}

func rip(b big.Int) []uint64 {
	if b.Cmp(big.NewInt(0)) == 0 {
		return []uint64{}
	}
	res := []uint64{new(big.Int).And(&b, big.NewInt(0xFFFFFFFF)).Uint64()}
	rsh := new(big.Int)
	rsh.Div(&b, big.NewInt(4294967296))
	res = append(res, rip(*rsh)...)
	return res
}