// Package bip39 is the Golang implementation of the BIP39 spec.
//
// The official BIP39 spec can be found at
// https://github.com/bitcoin/bips/blob/master/bip-0039.mediawiki
package bip39

import (
	"bytes"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/binary"
	"errors"
	"fmt"
	"math/big"
	"strings"
	"sync"

	"github.com/skycoin/skycoin/src/cipher"
	"github.com/skycoin/skycoin/src/cipher/bip39/wordlists"

	"github.com/skycoin/skycoin/src/cipher/pbkdf2"
)

var (
	// Some bitwise operands for working with big.Ints
	last11BitsMask  = big.NewInt(2047)
	shift11BitsMask = big.NewInt(2048)
	bigOne          = big.NewInt(1)
	bigTwo          = big.NewInt(2)

	// used to isolate the checksum bits from the entropy+checksum byte array
	wordLengthChecksumMasksMapping = map[int]*big.Int{
		12: big.NewInt(15),
		15: big.NewInt(31),
		18: big.NewInt(63),
		21: big.NewInt(127),
		24: big.NewInt(255),
	}
	// used to use only the desired x of 8 available checksum bits.
	// 256 bit (word length 24) requires all 8 bits of the checksum,
	// and thus no shifting is needed for it (we would get a divByZero crash if we did)
	wordLengthChecksumShiftMapping = map[int]*big.Int{
		12: big.NewInt(16),
		15: big.NewInt(8),
		18: big.NewInt(4),
		21: big.NewInt(2),
	}

	// wordList is the set of words to use
	wordList []string

	// wordMap is a reverse lookup map for wordList
	wordMap map[string]int

	// wordListMutex is held when changing wordList
	// TODO -- make wordList non-global
	wordListMutex sync.Mutex
)

var (
	// ErrInvalidEntropyLength is returned when trying to use an entropy set with
	// an invalid size.
	ErrInvalidEntropyLength = errors.New("Entropy length must be [128, 256] and a multiple of 32")

	// ErrChecksumIncorrect is returned when entropy has the incorrect checksum.
	ErrChecksumIncorrect = errors.New("Mnemonic checksum incorrect")

	// ErrSurroundingWhitespace is returned if a mnemonic sentence has extra whitespace around it
	ErrSurroundingWhitespace = errors.New("Mnemonic has extra whitespace around it")

	// ErrInvalidSeparator is returned if the words in a mnemonic sentence are not separated by a single ASCII space
	ErrInvalidSeparator = errors.New("Mnemonic must be separated by single spaces")

	// ErrUnknownWord is returned if a mnemonic sentence contains an unrecognized word
	ErrUnknownWord = errors.New("Mnemonic contains an unrecognized word")

	// ErrInvalidNumberOfWords is returned if a mnemonic sentence does not have 12, 15, 18, 21 or 24 words
	ErrInvalidNumberOfWords = errors.New("Mnemonic must have 12, 15, 18, 21 or 24 words")
)

func init() {
	setWordList(wordlists.English)
}

// setWordList sets the list of words to use for mnemonics. Currently the list
// that is set is used package-wide.
func setWordList(list []string) {
	wordListMutex.Lock()
	defer wordListMutex.Unlock()
	wordList = list
	wordMap = map[string]int{}
	for i, v := range wordList {
		wordMap[v] = i
	}
}

// getWordList gets the list of words to use for mnemonics.
func getWordList() []string {
	wordListMutex.Lock()
	defer wordListMutex.Unlock()
	return wordList
}

// getWordIndex gets word index in wordMap.
func getWordIndex(word string) (int, bool) {
	wordListMutex.Lock()
	defer wordListMutex.Unlock()
	idx, ok := wordMap[word]
	return idx, ok
}

// DefaultMnemonicEntropyBitSize is the default bit size for NewDefaultMnemonic's entropy
// TODO -- make 24 word seeds default? (256 bits entropy - recommended for HD wallets)
const DefaultMnemonicEntropyBitSize = 128

// NewDefaultMnemonic returns a generated mnemonic using entropy with bitSize 128
func NewDefaultMnemonic() (string, error) {
	entropy, err := NewEntropy(DefaultMnemonicEntropyBitSize)
	if err != nil {
		return "", err
	}

	return NewMnemonic(entropy)
}

// MustNewDefaultMnemonic returns a generated mnemonic using entropy with bitSize 128 and panics if there is an error
func MustNewDefaultMnemonic() string {
	seed, err := NewDefaultMnemonic()
	if err != nil {
		panic(err)
	}
	return seed
}

// NewEntropy will create random entropy bytes
// so long as the requested size bitSize is an appropriate size.
//
// bitSize has to be a multiple 32 and be within the inclusive range of {128, 256}
func NewEntropy(bitSize int) ([]byte, error) {
	err := validateEntropyBitSize(bitSize)
	if err != nil {
		return nil, err
	}

	entropy := cipher.RandByte(bitSize / 8)
	return entropy, err
}

// EntropyFromMnemonic takes a mnemonic generated by this library,
// and returns the input entropy used to generate the given mnemonic.
// An error is returned if the given mnemonic is invalid.
func EntropyFromMnemonic(mnemonic string) ([]byte, error) {
	words, err := splitMnemonicWords(mnemonic)
	if err != nil {
		return nil, err
	}

	// Decode the words into a big.Int.
	b := big.NewInt(0)
	for _, v := range words {
		index, found := wordMap[v]
		if !found {
			// This should have been caught by splitMnemonicWords()
			panic(fmt.Sprintf("word %q not found in reverse map", v))
		}
		var wordBytes [2]byte
		binary.BigEndian.PutUint16(wordBytes[:], uint16(index))
		b = b.Mul(b, shift11BitsMask)
		b = b.Or(b, big.NewInt(0).SetBytes(wordBytes[:]))
	}

	// Build and add the checksum to the big.Int.
	checksum := big.NewInt(0)
	checksumMask := wordLengthChecksumMasksMapping[len(words)]
	checksum = checksum.And(b, checksumMask)

	b.Div(b, big.NewInt(0).Add(checksumMask, bigOne))

	// The entropy is the underlying bytes of the big.Int. Any upper bytes of
	// all 0's are not returned so we pad the beginning of the slice with empty
	// bytes if necessary.
	entropy := b.Bytes()
	entropy = padByteSlice(entropy, len(words)/3*4)

	// Generate the checksum and compare with the one we got from the mneomnic.
	entropyChecksumBytes := computeChecksum(entropy)
	entropyChecksum := big.NewInt(int64(entropyChecksumBytes[0]))
	if l := len(words); l != 24 {
		checksumShift := wordLengthChecksumShiftMapping[l]
		entropyChecksum.Div(entropyChecksum, checksumShift)
	}

	if checksum.Cmp(entropyChecksum) != 0 {
		return nil, ErrChecksumIncorrect
	}

	return entropy, nil
}

// NewMnemonic will return a string consisting of the mnemonic words for
// the given entropy.
// If the provide entropy is invalid, an error will be returned.
func NewMnemonic(entropy []byte) (string, error) {
	// Compute some lengths for convenience.
	entropyBitLength := len(entropy) * 8
	checksumBitLength := entropyBitLength / 32
	sentenceLength := (entropyBitLength + checksumBitLength) / 11

	// Validate that the requested size is supported.
	err := validateEntropyBitSize(entropyBitLength)
	if err != nil {
		return "", err
	}

	// Add checksum to entropy.
	entropy = addChecksum(entropy)

	// Break entropy up into sentenceLength chunks of 11 bits.
	// For each word AND mask the rightmost 11 bits and find the word at that index.
	// Then bitshift entropy 11 bits right and repeat.
	// Add to the last empty slot so we can work with LSBs instead of MSB.

	// Entropy as an int so we can bitmask without worrying about bytes slices.
	entropyInt := new(big.Int).SetBytes(entropy)

	// Slice to hold words in.
	words := make([]string, sentenceLength)

	// Throw away big.Int for AND masking.
	word := big.NewInt(0)

	for i := sentenceLength - 1; i >= 0; i-- {
		// Get 11 right most bits and bitshift 11 to the right for next time.
		word.And(entropyInt, last11BitsMask)
		entropyInt.Div(entropyInt, shift11BitsMask)

		// Get the bytes representing the 11 bits as a 2 byte slice.
		wordBytes := padByteSlice(word.Bytes(), 2)

		// Convert bytes to an index and add that word to the list.
		words[i] = wordList[binary.BigEndian.Uint16(wordBytes)]
	}

	return strings.Join(words, " "), nil
}

// NewSeed creates a hashed seed output given the mnemonic string and a password.
// An error is returned if the mnemonic is not valid.
func NewSeed(mnemonic string, password string) ([]byte, error) {
	if err := ValidateMnemonic(mnemonic); err != nil {
		return nil, err
	}

	return newSeed(mnemonic, password), nil
}

// newSeed creates a hashed seed output given a provided string and password.
// No checking is performed to validate that the string provided is a valid mnemonic.
func newSeed(mnemonic, password string) []byte {
	return pbkdf2.Key([]byte(mnemonic), []byte("mnemonic"+password), 2048, 64, sha512.New)
}

// ValidateMnemonic returns an error if a mnemonic is invalid. It can be invalid
// for these reasons:
// - Number of words not a multiple of 3 and not at least 12 or at most 24 words
// - Words are not separated by exactly one ascii space
// - Mnemonic string has leading or trailing whitespace
// - Any word is not present in the wordlist
// - The mnemonic checksum is incorrect
// Note: this only works on ascii wordlists. Japanese, Chinese etc are not supported yet.
func ValidateMnemonic(mnemonic string) error {
	words, err := splitMnemonicWords(mnemonic)
	if err != nil {
		return err
	}

	if !isMnemonicChecksumValid(words) {
		return ErrChecksumIncorrect
	}

	return nil
}

// splitMnemonicWords attempts to verify that the provided mnemonic is valid.
// Validity is determined by both the number of words being appropriate,
// and that all the words in the mnemonic are present in the word list.
func splitMnemonicWords(mnemonic string) ([]string, error) {
	// Make sure no leading/trailing whitespace
	if mnemonic != strings.TrimSpace(mnemonic) {
		return nil, ErrSurroundingWhitespace
	}

	// Create a list of all the words in the mnemonic sentence
	words := strings.Split(mnemonic, " ")

	// Detect duplicate whitespace
	for _, w := range words {
		if w == "" {
			return nil, ErrInvalidSeparator
		}
	}

	numOfWords := len(words)

	// The number of words should be 12, 15, 18, 21 or 24
	if numOfWords%3 != 0 || numOfWords < 12 || numOfWords > 24 {
		return nil, ErrInvalidNumberOfWords
	}

	// Check if all words belong in the wordlist
	for _, word := range words {
		if _, ok := wordMap[word]; !ok {
			return nil, ErrUnknownWord
		}
	}

	return words, nil
}

// isMnemonicChecksumValid validates the checksum value of a mnemonic
func isMnemonicChecksumValid(words []string) bool {
	if len(words)%3 != 0 || len(words) < 12 || len(words) > 24 {
		panic("invalid number of words") // caller should validate words before passing to this function
	}

	var (
		entropyBitSize   = len(words) * 11
		checksumBitSize  = entropyBitSize % 32
		fullByteSize     = (entropyBitSize-checksumBitSize)/8 + 1
		checksumByteSize = fullByteSize - (fullByteSize % 4)
	)

	// Convert word indices to a big.Int representing the entropy.
	checksummedEntropy := big.NewInt(0)
	modulo := big.NewInt(2048)
	for _, v := range words {
		index := big.NewInt(int64(wordMap[v]))
		checksummedEntropy.Mul(checksummedEntropy, modulo)
		checksummedEntropy.Add(checksummedEntropy, index)
	}

	// Calculate the unchecksummed entropy so we can validate that the checksum is
	// correct.
	checksumModulo := big.NewInt(0).Exp(bigTwo, big.NewInt(int64(checksumBitSize)), nil)
	rawEntropy := big.NewInt(0).Div(checksummedEntropy, checksumModulo)

	// Convert big.Ints to byte padded byte slices.
	rawEntropyBytes := padByteSlice(rawEntropy.Bytes(), checksumByteSize)
	checksummedEntropyBytes := padByteSlice(checksummedEntropy.Bytes(), fullByteSize)

	// Validate that the checksum is correct.
	newChecksummedEntropyBytes := padByteSlice(addChecksum(rawEntropyBytes), fullByteSize)
	return bytes.Equal(checksummedEntropyBytes, newChecksummedEntropyBytes)
}

// Appends to data the first (len(data) / 32)bits of the result of sha256(data)
// Currently only supports data up to 32 bytes
func addChecksum(data []byte) []byte {
	// Get first byte of sha256
	hash := computeChecksum(data)
	firstChecksumByte := hash[0]

	// len() is in bytes so we divide by 4
	checksumBitLength := uint(len(data) / 4)

	// For each bit of check sum we want we shift the data one the left
	// and then set the (new) right most bit equal to checksum bit at that index
	// staring from the left
	dataBigInt := new(big.Int).SetBytes(data)
	for i := uint(0); i < checksumBitLength; i++ {
		// Bitshift 1 left
		dataBigInt.Mul(dataBigInt, bigTwo)

		// Set rightmost bit if leftmost checksum bit is set
		if uint8(firstChecksumByte&(1<<(7-i))) > 0 {
			dataBigInt.Or(dataBigInt, bigOne)
		}
	}

	return dataBigInt.Bytes()
}

func computeChecksum(data []byte) []byte {
	hasher := sha256.New()
	hasher.Write(data) // nolint: errcheck
	return hasher.Sum(nil)
}

// validateEntropyBitSize ensures that entropy is the correct size for being a
// mnemonic.
func validateEntropyBitSize(bitSize int) error {
	if (bitSize%32) != 0 || bitSize < 128 || bitSize > 256 {
		return ErrInvalidEntropyLength
	}
	return nil
}

// padByteSlice returns a byte slice of the given size with contents of the
// given slice left padded and any empty spaces filled with 0's.
func padByteSlice(slice []byte, length int) []byte {
	offset := length - len(slice)
	if offset <= 0 {
		return slice
	}
	newSlice := make([]byte, length)
	copy(newSlice[offset:], slice)
	return newSlice
}
