package vassago

import (
	"errors"
	"git.openprivacy.ca/lib.surveillance.resistant.tech/trostle-parrish"
	"math"
	"math/big"
)

// Client implements a Vassago client and can construct PIR requests that a Vassago Server can respond to.
type Client struct {
	pk                    trostle_parrish.TrostleParrish
	numberOfServerBuckets int
}

// NewClient generates a new client with the appropriate encryption parameters for the number of buckets and the underlying data.
func NewClient(numberOfServerBuckets int, ciphertextLength int) Client {
	var client Client
	client.numberOfServerBuckets = numberOfServerBuckets

	// We need homomorphism for up to 2^n * n additions
	t := math.Ceil(math.Log2(math.Pow(2, float64(numberOfServerBuckets)) * float64(numberOfServerBuckets)))

	client.pk = trostle_parrish.Generate(int(t), ciphertextLength)
	return client
}

// GenerateRequest creates a request that can be sent to a VassagoServer to obtain the data from the requested bucket.
func (c Client) GenerateRequest(requestedBucket int) []trostle_parrish.Ciphertext {
	requestVector := make([]trostle_parrish.Ciphertext, c.numberOfServerBuckets)
	for i := 0; i < c.numberOfServerBuckets; i++ {
		if requestedBucket != i {
			requestVector[i] = c.pk.Encrypt(big.NewInt(0))
		} else {
			requestVector[i] = c.pk.Encrypt(big.NewInt(1))
		}
	}
	return requestVector
}

// DecryptResponse takes in response from the VassagoServer to a request and outputs the decrypted bucket contents.
// An error is returned if this fails.
func (c Client) DecryptResponse(response []trostle_parrish.Ciphertext) ([]byte, error) {
	plaintext := make([]byte, len(response))
	for i, res := range response {
		decrypted := c.pk.Decrypt(res)

		// FIXME: Currently we assume we are encoding 1 byte per ciphertext, this could probably be made much more efficient.
		if decrypted.Int64() >= 0 && decrypted.Int64() <= 255 {
				if len(decrypted.Bytes()) > 0 {
					plaintext[i] = decrypted.Bytes()[0]
				} else {
					plaintext[i] = 0x00
				}
		} else {
			return plaintext, errors.New("could not sucessfully decrypt request, please check that the number of server buckets is correct")
		}
	}
	return plaintext, nil
}