package core

import (
	"fmt"
	"git.openprivacy.ca/openprivacy/log"
	"github.com/gtank/merlin"
	ristretto "github.com/gtank/ristretto255"
	"golang.org/x/crypto/sha3"
	"io"
)

// Transcript provides a consistent transcript primitive for our protocols
//
// We have the following goals:
//  - Allow sequential proofs over a common transcript (ensuring a single proof cannot be extracted standalone)
//  - be able to produce a human-readable transcript for auditing.
//
// The design of this API was inspired by Merlin: https://docs.rs/crate/merlin/
type Transcript struct {
	merlinTranscript *merlin.Transcript
	transcript       string
}

// NewTranscript creates a new Transcript with the given Label, the label should be unique to the application
func NewTranscript(label string) *Transcript {
	transcript := new(Transcript)
	transcript.merlinTranscript = merlin.NewTranscript(label)
	return transcript
}

// AddToTranscript appends a value to the transcript with the given label
// This binds the given data to the label.
func (t *Transcript) AddToTranscript(label string, b []byte) {
	op := fmt.Sprintf("%s (%d) %x;", label, len(b), b)
	t.transcript = fmt.Sprintf("%v\n%v", t.transcript, op)
	t.merlinTranscript.AppendMessage([]byte(label), b)
}

// AddElementToTranscript appends a value to the transcript with the given label
// This binds the given data to the label.
func (t *Transcript) AddElementToTranscript(label string, element *ristretto.Element) {
	t.AddToTranscript(label, element.Encode([]byte{}))
}

// OutputTranscriptToAudit outputs a human-readable copy of the transcript so far.
func (t Transcript) OutputTranscriptToAudit() string {
	return t.transcript
}

// NewProtocol provides explicit protocol separation in a transcript (more readable audit scripts and even more explicit
// binding of committed values to a given context)
func (t *Transcript) NewProtocol(label string) {
	op := fmt.Sprintf("---- new-protcol: %s ----", label)
	t.transcript = fmt.Sprintf("%v\n%v", t.transcript, op)
	t.merlinTranscript.AppendMessage([]byte("protocol"), []byte(label))
}

// CommitToTranscript generates a challenge based on the current transcript, it also commits the challenge to the transcript.
func (t *Transcript) CommitToTranscript(label string) []byte {
	b := t.merlinTranscript.ExtractBytes([]byte(label), 64)
	t.transcript = fmt.Sprintf("%v\nextract %v: %v", t.transcript, label, b)
	return b
}

// PRNG defines a psuedorandom number generator
type PRNG struct {
	prng io.Reader
}

// Next returns the next "random" scalar from the PRNG
func (prng *PRNG) Next(buf []byte, next *ristretto.Scalar) error {
	n, err := io.ReadFull(prng.prng, buf)
	if n != 64 || err != nil {
		log.Errorf("could not read prng: %v %v", n, err)
		return fmt.Errorf("error fetching complete output from prng: %v", err)
	}
	next.FromUniformBytes(buf)
	return nil
}

// CommitToPRNG commits the label to the transcript and derives a PRNG from the transcript.
func (t *Transcript) CommitToPRNG(label string) PRNG {
	b := t.merlinTranscript.ExtractBytes([]byte(label), 64)
	prng := sha3.NewShake256()
	prng.Write(b)
	return PRNG{prng: prng}
}

// CommitToGenerator derives a verifiably random generator from the transcript
func (t *Transcript) CommitToGenerator(label string) *ristretto.Element {
	c := t.CommitToTranscript(label)
	return new(ristretto.Element).FromUniformBytes(c)
}

// CommitToGenerators derives a set of verifiably random generators from the transcript
func (t *Transcript) CommitToGenerators(label string, n int) (generators []*ristretto.Element) {
	for i := 0; i < n; i++ {
		generators = append(generators, t.CommitToGenerator(fmt.Sprintf("%v-%d", label, i)))
	}
	return generators
}

// CommitToTranscriptScalar is a convenience method for CommitToTranscript which returns a ristretto Scalar
func (t *Transcript) CommitToTranscriptScalar(label string) *ristretto.Scalar {
	c := t.CommitToTranscript(label)
	s := new(ristretto.Scalar)
	s.FromUniformBytes(c[:])
	return s
}