package privacypass

import (
	"crypto/hmac"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"git.openprivacy.ca/cwtch.im/tapir/primitives/core"
	"git.openprivacy.ca/openprivacy/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.SetUniformBytes(b)

	Ht := sha3.Sum512(t.t)
	T, _ := new(ristretto.Element).SetUniformBytes(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.Bytes()...))
	mac := hmac.New(sha3.New512, key[:])
	mac.Write(data)
	return SpentToken{t.t, mac.Sum(nil)}
}

// 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, ristretto.NewGeneratorElement().Bytes())
	transcript.AddToTranscript(BatchProofY, Y.Bytes())
	transcript.AddToTranscript(BatchProofPVector, []byte(fmt.Sprintf("%v", blindedTokens)))
	transcript.AddToTranscript(BatchProofQVector, []byte(fmt.Sprintf("%v", signedTokens)))
	prng := transcript.CommitToPRNG("w")
	M := ristretto.NewIdentityElement()
	Z := ristretto.NewIdentityElement()
	buf := make([]byte, 64)
	c := new(ristretto.Scalar)
	for i := range blindedTokens {
		err := prng.Next(buf, c)
		if err != nil {
			log.Errorf("error verifying batch proof: %v", err)
			return false
		}
		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, ristretto.NewGeneratorElement(), 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
}

// MarshalJSON - in order to store tokens in a serialized form we need to expose the private, unexported value
// `t`. Note that `r` is not needed to spend the token, and as such we effectively destroy it when we serialize.
// Ideally, go would let us do this with an annotation, alas.
func (t Token) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct {
		T []byte `json:"t"`
		W *ristretto.Element
	}{
		T: t.t,
		W: t.W,
	})
}
