package trostle_parrish

import (
	"crypto/rand"
	"math/big"
)


// TrostleParrish implements a somewhat-homomorphic cryptosystem outlined by [1] and [2]
// [1] Trostle, Jonathan, and Andy Parrish. "Efficient computationally private information retrieval from anonymity or trapdoor groups." International Conference on Information Security. Springer, Berlin, Heidelberg, 2010.
// [2] Mayberry, Travis, Erik-Oliver Blass, and Agnes Hui Chan. "PIRMAP: Efficient private information retrieval for MapReduce." International Conference on Financial Cryptography and Data Security. Springer, Berlin, Heidelberg, 2013.
type TrostleParrish struct {
	m, b, T, N *big.Int

	// We precompute these values.
	bInv       *big.Int
	twon       *big.Int
}

// More readable type definitions
type Ciphertext *big.Int

// Generate sets up a Trostle-Parish cryptosystem
// This system is "somewhat-homomorphic"
// 2^t = number of additions to support
// n = max plaintext bit size
func Generate(t int, n int) (pk TrostleParrish) {
	for {

		// choose secret m and b
		pk.m, _ = rand.Prime(rand.Reader, (t*2)+n)
		pk.b, _ = rand.Int(rand.Reader, pk.m)

		pk.N = big.NewInt(int64(n))
		pk.twon = new(big.Int).Exp(big.NewInt(2), pk.N, pk.m)

		// T is our parameter that determines how many addition operations we can safely do
		tbi := big.NewInt(int64(n/t))
		pk.T = new(big.Int).Exp(big.NewInt(2), tbi, pk.m)

		// Check that b has an inverse mod m
		pk.bInv = new(big.Int).ModInverse(pk.b, pk.m)
		for pk.bInv != nil {
			return
		}
	}
}

// Encrypt takes in a Plaintext(x) and a keyset (pk)
func (pk TrostleParrish) Encrypt(x *big.Int) Ciphertext {
	r, _ := rand.Int(rand.Reader, pk.T)
	// r * 2^n
	rtwon := new(big.Int).Mul(r, pk.twon)
	// r * 2^n + x
	rtwonx := new(big.Int).Add(rtwon, x)
	// b * (r * 2^n + x)
	bmod := new(big.Int).Mul(pk.b, rtwonx)
	//  b * (r * 2^n + x) mod m
	return new(big.Int).Mod(bmod, pk.m)
}

// Add two ciphertexts together and return the result.
// Note that the number of meaningful additions is limited by the encryption parameters
func Add(a Ciphertext, b Ciphertext) Ciphertext {
	return new(big.Int).Add(a, b)
}

// Decrypt takes in a Ciphertext(c) and Keyset(pk) and outputs the Plaintext
func (pk TrostleParrish) Decrypt(c Ciphertext) *big.Int {
	// b^-1 * c
	plaintext := new(big.Int).Mul(pk.bInv, c)
	//  b^-1 * c mod m
	modm := new(big.Int).Mod(plaintext, pk.m)
	//  b^-1 * c mod m mod 2^n
	return new(big.Int).Mod(modm, pk.twon)
}