package zkp

import (
	"crypto/ecdsa"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"strings"

	"github.com/ethereum/go-ethereum/crypto"
	"qoobing.com/gomod/log"
)

type DLogProof struct {
	ChallengeResponse *big_Int `json:"cr"`
	RandomPoint       string   `json:"pkr"`
}

func (dlogproof *DLogProof) ToString() string {
	var b, _ = json.Marshal(dlogproof)
	return string(b)
}

func DLogProofFromString(jsonstr string) (*DLogProof, error) {
	var proof = DLogProof{}
	if err := json.Unmarshal([]byte(jsonstr), &proof); err != nil {
		return nil, err
	}
	return &proof, nil
}

/*
 * Prove:
 *   step 1. generate random bignumber: skr, let pkr = skr * G
 *   step 2. caculate challenge: c = SHA256(pkr || G || pk)
 *   step 3. caculate challenge response: cr = skr - c * sk
 * result:
 *   (pkr, cr)
 */
func DLogProve(sk *big.Int) (proof *DLogProof, err error) {
	// Step 0. get pk
	cv := crypto.S256()
	Gx := cv.Params().Gx
	Gy := cv.Params().Gy
	x, y := cv.ScalarBaseMult(sk.Bytes())
	var pk = ecdsa.PublicKey{Curve: cv, X: x, Y: y}

	// Step 1. generate random bignumber: skr, let pkr = skr * G
	randomkey, err := crypto.GenerateKey()
	if err != nil {
		log.Warningf("Gernerate random Q1 failed: '%s'", err.Error())
		return nil, errors.New("generate random key/number failed")
	}
	var skr = randomkey.D
	var pkr = randomkey.PublicKey
	//TODO: sk_t_rand_commitment.zeroize();

	// Step 2. caculate challenge: c = SHA256(pkr || G || pk)
	challenge := fmt.Sprintf(
		"%064x,%064x@%064x,%064x@%064x,%064x",
		pkr.X, pkr.Y, Gx, Gy, pk.X, pk.Y)
	chash256 := sha256.Sum256([]byte(challenge))
	var c = new(big.Int).SetBytes(chash256[:])
	log.Debugf("challenge='%s',c='%x'", challenge, c)

	// Step 3. caculate challenge response: cr = skr - c * sk
	cvN := cv.Params().N
	tmp := new(big.Int).Mul(c, sk)
	var cr = new(big.Int).Sub(skr, tmp)
	cr = new(big.Int).Mod(cr, cvN)

	// return
	proof = &DLogProof{
		RandomPoint:       fmt.Sprintf("%x,%x", pkr.X, pkr.Y),
		ChallengeResponse: (*big_Int)(cr),
	}
	return proof, nil
}

/*
 * Verify:
 *   step 1. caculate challenge: c = SHA256(pkr || G || pk)
 *   step 2. caculate proof verify: pkv = cr * G + c * pk
 * result:
 *   pkr =?= pkv
 */
func DLogVerify(pkstr string, proof *DLogProof) bool {
	// Step 0. get pk & pkr
	var pk, err1 = HexStringToS256Point(pkstr)
	if err1 != nil {
		//errors.New("input pk is invalid")
		return false
	}
	var pkr, err2 = HexStringToS256Point(proof.RandomPoint)
	if err2 != nil {
		//errors.New("input pkr is invalid")
		return false
	}

	// Step 1. caculate challenge: c = SHA256(pkr || G || pk)
	cv := crypto.S256()
	Gx := cv.Params().Gx
	Gy := cv.Params().Gy
	challenge := fmt.Sprintf(
		"%064x,%064x@%064x,%064x@%064x,%064x",
		pkr.X, pkr.Y, Gx, Gy, pk.X, pk.Y)
	chash256 := sha256.Sum256([]byte(challenge))
	var c = new(big.Int).SetBytes(chash256[:])
	var cr = (*big.Int)(proof.ChallengeResponse)
	log.Debugf("challenge='%s',c='%x',cr='%x'", challenge, c, cr)

	// Step 2. caculate proof verify: pkv = cr * G + c * pk
	var pkv = ecdsa.PublicKey{Curve: crypto.S256()}
	pkv.X, pkv.Y = cv.ScalarBaseMult(cr.Bytes())
	tempX, tempY := cv.ScalarMult(pk.X, pk.Y, c.Bytes())
	tempX, tempY = cv.Add(pkv.X, pkv.Y, tempX, tempY)
	log.Debugf("pkv='%x,%x'", tempX, tempY)

	// return
	if pkr.X.Cmp(tempX) != 0 || pkr.Y.Cmp(tempY) != 0 {
		log.Debugf("pkr='%x,%x'", pkr.X, pkr.Y)
		log.Debugf("pkv='%x,%x'", tempX, tempY)
		return false
	}
	return true
}

func HexStringToS256Point(pstr string) (*ecdsa.PublicKey, error) {
	var pk = ecdsa.PublicKey{Curve: crypto.S256()}
	if x2y2 := strings.Split(pstr, ","); len(x2y2) != 2 {
		return nil, errors.New("input pk is invalid")
	} else if x, err := hex.DecodeString(x2y2[0]); err != nil {
		return nil, errors.New("input pk.x is invalid")
	} else if y, err := hex.DecodeString(x2y2[1]); err != nil {
		return nil, errors.New("input pk.y is invalid")
	} else {
		pk.X = new(big.Int).SetBytes(x)
		pk.Y = new(big.Int).SetBytes(y)
	}
	return &pk, nil
}