package privacypass

import (
	"crypto/hmac"
	"crypto/rand"
	"cwtch.im/tapir/primitives/core"
	"fmt"
	"git.openprivacy.ca/openprivacy/libricochet-go/log"
	ristretto "github.com/gtank/ristretto255"

	"golang.org/x/crypto/sha3"
)

// Token is an implementation of PrivacyPass
// Davidson A, Goldberg I, Sullivan N, Tankersley G, Valsorda F. Privacy pass: Bypassing internet challenges anonymously. Proceedings on Privacy Enhancing Technologies. 2018 Jun 1;2018(3):164-80.
type Token struct {
	t []byte
	r *ristretto.Scalar
	W *ristretto.Element
}

// GetT returns the underlying bytes for token for use in constraint proofs.
func (t Token) GetT() []byte {
	return t.t
}

// BlindedToken encapsulates a Blinded Token
type BlindedToken struct {
	P *ristretto.Element
}

// SignedToken encapsulates a Signed (Blinded) Token
type SignedToken struct {
	Q *ristretto.Element
}

// SpentToken encapsulates the parameters needed to spend a Token
type SpentToken struct {
	T   []byte
	MAC []byte
}

// TokenPaymentHandler defines an interface with external payment processors
type TokenPaymentHandler interface {
	MakePayment()
	// Next Token
	NextToken(data []byte, hostname string) (SpentToken, error)
}

// GenBlindedToken initializes the Token
// GenToken() & Blind()
func (t *Token) GenBlindedToken() BlindedToken {
	t.t = make([]byte, 32)
	rand.Read(t.t)
	t.r = new(ristretto.Scalar)
	b := make([]byte, 64)
	rand.Read(b)
	t.r.FromUniformBytes(b)

	Ht := sha3.Sum512(t.t)
	T := new(ristretto.Element).FromUniformBytes(Ht[:])
	P := new(ristretto.Element).ScalarMult(t.r, T)
	return BlindedToken{P}
}

// unblindSignedToken unblinds a token that has been signed by a server
func (t *Token) unblindSignedToken(token SignedToken) {
	t.W = new(ristretto.Element).ScalarMult(new(ristretto.Scalar).Invert(t.r), token.Q)
}

// SpendToken binds the token with data and then redeems the token
func (t *Token) SpendToken(data []byte) SpentToken {
	key := sha3.Sum256(append(t.t, t.W.Encode(nil)...))
	mac := hmac.New(sha3.New512, key[:])
	return SpentToken{t.t, mac.Sum(data)}
}

// GenerateBlindedTokenBatch generates a batch of blinded tokens (and their unblinded equivalents)
func GenerateBlindedTokenBatch(num int) (tokens []*Token, blindedTokens []BlindedToken) {
	for i := 0; i < num; i++ {
		tokens = append(tokens, new(Token))
		blindedTokens = append(blindedTokens, tokens[i].GenBlindedToken())
	}
	return
}

// verifyBatchProof verifies a given batch proof (see also UnblindSignedTokenBatch)
func verifyBatchProof(dleq DLEQProof, Y *ristretto.Element, blindedTokens []BlindedToken, signedTokens []SignedToken, transcript *core.Transcript) bool {
	transcript.NewProtocol(BatchProofProtocol)
	transcript.AddToTranscript(BatchProofX, new(ristretto.Element).Base().Encode(nil))
	transcript.AddToTranscript(BatchProofY, Y.Encode(nil))
	transcript.AddToTranscript(BatchProofPVector, []byte(fmt.Sprintf("%v", blindedTokens)))
	transcript.AddToTranscript(BatchProofQVector, []byte(fmt.Sprintf("%v", signedTokens)))
	prng := transcript.CommitToPRNG("w")
	M := new(ristretto.Element).Zero()
	Z := new(ristretto.Element).Zero()
	for i := range blindedTokens {
		c := prng.Next()
		M = new(ristretto.Element).Add(new(ristretto.Element).ScalarMult(c, blindedTokens[i].P), M)
		Z = new(ristretto.Element).Add(new(ristretto.Element).ScalarMult(c, signedTokens[i].Q), Z)
	}
	return VerifyDiscreteLogEquivalenceProof(dleq, new(ristretto.Element).Base(), Y, M, Z, transcript)
}

// UnblindSignedTokenBatch taking in a set of tokens, their blinded & signed counterparts, a server public key (Y), a DLEQ proof and a transcript
// verifies that the signing procedure has taken place correctly and unblinds the tokens.
func UnblindSignedTokenBatch(tokens []*Token, blindedTokens []BlindedToken, signedTokens []SignedToken, Y *ristretto.Element, proof DLEQProof, transcript *core.Transcript) bool {
	verified := verifyBatchProof(proof, Y, blindedTokens, signedTokens, transcript)
	if !verified {
		log.Debugf("Failed to unblind tokens: %v", transcript.OutputTranscriptToAudit())
		return false
	}
	for i, t := range tokens {
		t.unblindSignedToken(signedTokens[i])
	}
	return true
}