package starkware

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"math/big"
	"strings"
)

func GenerateKRfc6979(msgHash, priKey *big.Int, seed int) *big.Int {
	msgHash = big.NewInt(0).Set(msgHash) // copy
	bitMod := msgHash.BitLen() % 8
	if bitMod <= 4 && bitMod >= 1 && msgHash.BitLen() > 248 {
		msgHash.Mul(msgHash, big.NewInt(16))
	}
	var extra []byte
	if seed > 0 {
		buf := new(bytes.Buffer)
		var data interface{}
		if seed < 256 {
			data = uint8(seed)
		} else if seed < 65536 {
			data = uint16(seed)
		} else if seed < 4294967296 {
			data = uint32(seed)
		} else {
			data = uint64(seed)
		}
		_ = binary.Write(buf, binary.BigEndian, data)
		extra = buf.Bytes()
	}
	return generateSecret(EcOrder, priKey, sha256.New, msgHash.Bytes(), extra)
}

type Signature struct {
	R string `json:"r"`
	S string `json:"s"`
}

func Sign(messageHash string, privateKey string) Signature {
	if strings.HasPrefix(privateKey, "0x") {
		privateKey = privateKey[2:]
	}
	privKey, _ := new(big.Int).SetString(privateKey, 16)
	msgHash, _ := new(big.Int).SetString(messageHash, 10)
	seed := 0
	EcGen := pedersenCfg.ConstantPoints[1]
	alpha := pedersenCfg.ALPHA
	nBit := big.NewInt(0).Exp(big.NewInt(2), NElementBitsEcdsa, nil)
	for {
		k := GenerateKRfc6979(msgHash, privKey, seed)
		//	Update seed for next iteration in case the value of k is bad.
		if seed == 0 {
			seed = 1
		} else {
			seed += 1
		}
		// Cannot fail because 0 < k < EC_ORDER and EC_ORDER is prime.
		x := ecMult(k, EcGen, alpha, FieldPrime)[0]
		// !(1 <= x < 2 ** N_ELEMENT_BITS_ECDSA)
		if !(x.Cmp(one) > 0 && x.Cmp(nBit) < 0) {
			continue
		}
		// msg_hash + r * priv_key
		x1 := big.NewInt(0).Add(msgHash, big.NewInt(0).Mul(x, privKey))
		// (msg_hash + r * priv_key) % EC_ORDER == 0
		if big.NewInt(0).Mod(x1, EcOrder).Cmp(zero) == 0 {
			continue
		}
		// w = div_mod(k, msg_hash + r * priv_key, EC_ORDER)
		w := divMod(k, x1, EcOrder)
		// not (1 <= w < 2 ** N_ELEMENT_BITS_ECDSA)
		if !(w.Cmp(one) > 0 && w.Cmp(nBit) < 0) {
			continue
		}
		s1 := divMod(one, w, EcOrder)
		return Signature{
			R: "0x" + fmt.Sprintf("%x", x),
			S: "0x" + fmt.Sprintf("%x", s1),
		}
	}
}