package main

/*
#cgo CFLAGS: -O2
#cgo LDFLAGS: -lcrypto -lm
#include "../c/dpf.h"
#include "../c/okvClient.h"
#include "../c/dpf.c"
#include "../c/okvClient.c"
*/
import "C"

//sssssssssss
import (
	lib "2PPS/lib"
	"bytes"
	"crypto/rand"
	"crypto/sha256"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"math/big"
	"net"
	"sync"
	"time"
	"unsafe"

	"golang.org/x/crypto/nacl/box"
)

type tweet struct {
	Topics []string
	Text   string
}

const leader string = "127.0.0.1:4441"

//needs to be changed at leader/follower/client at the same time
const neededSubscriptions = 1
const numClients = 2

var topicList []string
var archiveTopicList []string

//todo! expand this for multiple clients
var archiveInterests = make([]int, 1)
var sharedSecret [numClients][2][32]byte = createSharedSecret()

var wantsArchive = make([]byte, 1)

var dataLen int = 32
var numThreads int = 12
var numRows int = 2

var leaderPublicKey *[32]byte
var followerPublicKey *[32]byte
var clientPrivateKey *[32]byte
var clientPublicKey *[32]byte

func main() {

	//creates test tweets
	tweets := make([][]byte, numClients)
	for i := range tweets {
		tweets[i] = make([]byte, dataLen)
	}

	for i := 0; i < numClients; i++ {
		var tweet []byte
		if i == 0 {
			topics := []byte("house; mouse")
			text := []byte("I am a house in a mouse;")
			tweet = append(tweet, topics...)
			tweet = append(tweet, text...)
		}

		length := len(tweet)
		for i := length; i < dataLen; i++ {
			tweet = append(tweet, []byte(";")[0])
		}
		tweets[i] = tweet
	}

	wg := &sync.WaitGroup{}

	for i := 0; i < numClients; i++ {
		wg.Add(1)
		go client(tweets[i], i)
	}
	wg.Wait()
}

func client(tweet []byte, clientNumber int) {

	/*
		if len(os.Args) != 4 {
			fmt.Println("try again with: numThreads, dataLength, numRows")
			return
		}

		//input when executing is follower amount
		serverAmount, _ = strconv.Atoi(os.Args[1])
		serverAmount++
		dataLen, _ = strconv.Atoi(os.Args[2])
		numThreads, _ = strconv.Atoi(os.Args[3])
		dbSize, _ = strconv.Atoi(os.Args[4])
	*/

	generatedPublicKey, generatedPrivateKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		panic(err)
	}
	clientPrivateKey = generatedPrivateKey
	clientPublicKey = generatedPublicKey

	C.initializeCipher()

	//initializes the connection to the leader
	conf := &tls.Config{
		InsecureSkipVerify: true,
	}

	leaderConn, err := tls.Dial("tcp", leader, conf)
	if err != nil {
		panic(err)
	}
	leaderConn.SetDeadline(time.Time{})

	//receives topics first so client can participate asap
	receiveTopicLists(leaderConn)

	//gets the public keys of both servers
	var tmpLeaderPubKey [32]byte
	_, err = leaderConn.Read(tmpLeaderPubKey[:])
	if err != nil {
		panic(err)
	}
	leaderPublicKey = &tmpLeaderPubKey

	var tmpFollowerPubKey [32]byte
	_, err = leaderConn.Read(tmpFollowerPubKey[:])
	if err != nil {
		panic(err)
	}
	followerPublicKey = &tmpFollowerPubKey

	//sends own public key
	_, err = leaderConn.Write(clientPublicKey[:])
	if err != nil {
		panic(err)
	}

	//setup ends above
	//while client is active he is always connected and has to participate

	for {
		//gets current phase
		phase := make([]byte, 1)
		_, err = leaderConn.Read(phase)
		if err != nil {
			panic(err)
		}

		fmt.Println("Phase: ", phase[0])

		if phase[0] == 1 {

			//todo! put into tweet creation
			roundAsBytes := make([]byte, 4)
			_, err = leaderConn.Read(roundAsBytes)
			if err != nil {
				panic(err)
			}

			//prep the query
			dataSize := len(tweet)
			querySize := make([]byte, 4)
			cQuerySize := C.int(byteToInt(querySize))

			var dpfQueryA *C.uchar
			var dpfQueryB *C.uchar

			//change to clientNumber
			C.prepQuery(C.int(0), C.int(numRows), (*C.uchar)(&tweet[0]), C.int(dataSize), &cQuerySize, &dpfQueryA, &dpfQueryB)

			intQuerySize := int(cQuerySize) //byteToInt(querySize)

			//write the query
			queryAPlaintext := C.GoBytes(unsafe.Pointer(dpfQueryA), C.int(intQuerySize))

			//encrypts queryA and appends it to message
			var nonce [24]byte
			//fill nonce with randomness
			_, err = rand.Read(nonce[:])
			if err != nil {
				panic("couldn't get randomness for nonce!")
			}
			dpfQueryAEncrypted := box.Seal(nonce[:], queryAPlaintext, &nonce, leaderPublicKey, clientPrivateKey)

			//encrypts queryB and appends it to message
			queryBPlaintext := C.GoBytes(unsafe.Pointer(dpfQueryB), C.int(intQuerySize))

			//fill nonce with randomness
			_, err = rand.Read(nonce[:])
			if err != nil {
				panic("couldn't get randomness for nonce!")
			}
			dpfQueryBEncrypted := box.Seal(nonce[:], queryBPlaintext, &nonce, followerPublicKey, clientPrivateKey)

			//writes the dpfQuery to the leader
			dpfLengthBytes := intToByte(len(dpfQueryAEncrypted))
			_, err = leaderConn.Write(dpfLengthBytes)
			if err != nil {
				panic(err)
			}

			_, err = leaderConn.Write(dpfQueryAEncrypted)
			if err != nil {
				panic(err)
			}

			_, err = leaderConn.Write(dpfQueryBEncrypted)
			if err != nil {
				panic(err)
			}

			//auditing starts here

			//read seed from leader(in preparation for auditing)
			seed := make([]byte, 16)
			_, err = leaderConn.Read(seed)
			if err != nil {
				panic(err)
			}

			//prepare message for auditor, box it, and send to server A
			//prepare the auditor message
			outputsA := (*C.uchar)(C.malloc(C.ulong(160)))
			outputsB := (*C.uchar)(C.malloc(C.ulong(160)))

			C.prepAudit((*C.uchar)(&seed[0]), outputsA, outputsB, dpfQueryA, dpfQueryB)

			auditPlaintextLeader := C.GoBytes(unsafe.Pointer(outputsA), C.int(160))
			auditPlaintextFollower := C.GoBytes(unsafe.Pointer(outputsB), C.int(160))

			//encrypt messages for auditing
			//fill nonce with randomness
			_, err = rand.Read(nonce[:])
			if err != nil {
				panic("couldn't get randomness for nonce!")
			}
			auditCiphertextLeader := box.Seal(nonce[:], auditPlaintextLeader, &nonce, leaderPublicKey, clientPrivateKey)

			//fill nonce with randomness
			_, err = rand.Read(nonce[:])
			if err != nil {
				panic("couldn't get randomness for nonce!")
			}
			auditCiphertextFollower := box.Seal(nonce[:], auditPlaintextFollower, &nonce, followerPublicKey, clientPrivateKey)

			//send boxed audits to leader
			auditLengthBytes := intToByte(len(auditCiphertextLeader))

			_, err = leaderConn.Write(auditLengthBytes)
			if err != nil {
				panic(err)
			}

			_, err = leaderConn.Write(auditCiphertextLeader)
			if err != nil {
				panic(err)
			}

			_, err = leaderConn.Write(auditCiphertextFollower)
			if err != nil {
				panic(err)
			}

			C.free(unsafe.Pointer(outputsA))
			C.free(unsafe.Pointer(outputsB))
			C.free(unsafe.Pointer(dpfQueryA))
			C.free(unsafe.Pointer(dpfQueryB))
		} else if phase[0] == 3 {
			/*
				possible Values
				0 : new client
					leader expects sharedSecrets, expects pirQuery
				1 : update needed
					leader sends topicList, performs local update of sharedSecret, expects pirQuery
				2 : no update needed
					nothing
			*/
			subPhase := make([]byte, 1)
			_, err := leaderConn.Read(subPhase)
			if err != nil {
				panic(err)
			}

			var encryptedQueryLeader, encryptedQueryFollower []byte
			//first time participating
			if subPhase[0] == 0 {
				receiveTopicLists(leaderConn)
				encryptedQueryLeader, encryptedQueryFollower = createPIRQuery(int(subPhase[0]), clientNumber)
				sendQuerys(encryptedQueryLeader, encryptedQueryFollower, leaderConn, false)

			}
			//updates the topic list and what client is interested in
			if subPhase[0] == 1 {

				receiveTopicLists(leaderConn)

				//updates local secret
				for index := 0; index < 2; index++ {
					sharedSecret[clientNumber][index] = sha256.Sum256(sharedSecret[clientNumber][index][:])
				}

				encryptedQueryLeader, encryptedQueryFollower = createPIRQuery(int(subPhase[0]), clientNumber)
				sendQuerys(encryptedQueryLeader, encryptedQueryFollower, leaderConn, false)

			}

			receiveTweets(sharedSecret[clientNumber], leaderConn, clientNumber, false)

			if len(archiveTopicList) > 0 {
				wantsArchive[0] = 0 //archive test
			} else {
				wantsArchive[0] = 0
			}

			_, err = leaderConn.Write(wantsArchive)
			if err != nil {
				panic(err)
			}

			if wantsArchive[0] == 1 && len(archiveTopicList) > 0 {
				encryptedQueryLeader, encryptedQueryFollower = createPIRQuery(-1, clientNumber)
				sendQuerys(encryptedQueryLeader, encryptedQueryFollower, leaderConn, true)
				receiveTweets(sharedSecret[clientNumber], leaderConn, clientNumber, true)
			}

		} else {
			panic("somethin went wrong")
		}
	}
}

//creates and sends the pirQuerys for each server
func createPIRQuery(subPhase int, clientNumber int) ([]byte, []byte) {
	//later this will be taken from gui, this is only for testing
	topicsOfInterest := make([]int, 1)
	topicsOfInterest[0] = 1

	archiveInterests[0] = 1

	tmptopicsOfInterest := make([]int, len(topicsOfInterest))
	copy(tmptopicsOfInterest, topicsOfInterest)

	tmpNeededSubscriptions := neededSubscriptions

	tmpTopicList := make([]string, len(topicList))
	copy(tmpTopicList, topicList)
	if wantsArchive[0] == 1 && subPhase == -1 {
		tmpNeededSubscriptions = len(archiveInterests)
		copy(tmptopicsOfInterest, archiveInterests) //archiveInterests from gui
		copy(tmpTopicList, archiveTopicList)
	}

	topicsOfInterestAsBytes := make([][]byte, tmpNeededSubscriptions)
	for i := range topicsOfInterestAsBytes {
		topicsOfInterestAsBytes[i] = make([]byte, len(tmpTopicList))
	}

	//creates fake topicsOfInterest if client is boooring
	if len(tmptopicsOfInterest) < tmpNeededSubscriptions && subPhase != -1 {
		//this should be done when client is idling
		tmptopicsOfInterest = addFakeInterests(tmptopicsOfInterest)
	}

	for topic, position := range tmptopicsOfInterest {
		topicsOfInterestAsBytes[topic][position-1] = 1
	}

	//this for N servers for one topic
	/* interested in topic 3, 6 servers
	this needs to be repeated for every topic client is interested in

	wanted
	[0, 0, 1]

	random creation of 1 pirQuery
	[0, 1, 0]

	manual creation of 1 pirQuery
	[0, 1, 1]

	xor result
	[0, 0, 1]

	*/

	//pirQuery [serverAmount][topicsofinterest][topicAmount]byte
	pirQuerys := make([][][]byte, 2)
	for i := range pirQuerys {
		pirQuerys[i] = make([][]byte, len(tmptopicsOfInterest))
		for j := range pirQuerys[i] {
			pirQuerys[i][j] = make([]byte, len(tmpTopicList))
		}
	}

	//for leader
	//pirQuery will be filled with random bits
	for topic := range tmptopicsOfInterest {
		for index := range tmpTopicList {
			bit, err := rand.Int(rand.Reader, big.NewInt(2))
			if err != nil {
				panic(err)
			}
			pirQuerys[0][topic][index] = byte(bit.Int64())
		}
	}

	//creating last manually with result and wanted
	//if position random result correct -> 0, not correct -> 1
	for topic := range tmptopicsOfInterest {
		for index := range tmpTopicList {
			if pirQuerys[0][topic][index] == topicsOfInterestAsBytes[topic][index] {
				pirQuerys[1][topic][index] = 0
			} else {
				pirQuerys[1][topic][index] = 1
			}
		}
	}

	//flattens the querys to be able to send them more efficently
	messagesFlattened := make([][]byte, 2)
	//adds the sharedSecret to the first pirQuery when first time participating
	if subPhase == 0 {
		for server := 0; server < 2; server++ {
			messagesFlattened[server] = append(messagesFlattened[server], sharedSecret[clientNumber][server][:]...)
		}
	}

	for server := 0; server < 2; server++ {
		for topic := range pirQuerys[server] {
			messagesFlattened[server] = append(messagesFlattened[server], pirQuerys[server][topic][:]...)
		}
	}

	var nonce [24]byte
	_, err := rand.Read(nonce[:])
	if err != nil {
		panic("couldn't get randomness for nonce!")
	}
	encryptedQueryLeader := box.Seal(nonce[:], messagesFlattened[0], &nonce, leaderPublicKey, clientPrivateKey)

	_, err = rand.Read(nonce[:])
	if err != nil {
		panic("couldn't get randomness for nonce!")
	}
	encryptedQueryFollower := box.Seal(nonce[:], messagesFlattened[1], &nonce, followerPublicKey, clientPrivateKey)

	return encryptedQueryLeader, encryptedQueryFollower
}

func sendQuerys(encryptedQueryLeader, encryptedQueryFollower []byte, leaderConn net.Conn, getArchive bool) {
	encryptedLength := len(encryptedQueryLeader)

	//sends the pirQuerysLength to the leader
	_, err := leaderConn.Write(intToByte(encryptedLength))
	if err != nil {
		panic(err)
	}

	//sends the pirQuerys to the leader
	_, err = leaderConn.Write(encryptedQueryLeader)
	if err != nil {
		panic(err)
	}

	_, err = leaderConn.Write(encryptedQueryFollower)
	if err != nil {
		panic(err)
	}

	if getArchive {
		leaderConn.Write(intToByte(len(archiveInterests)))
		if err != nil {
			panic(err)
		}
	}
}

func receiveTweets(sharedSecret [2][32]byte, leaderConn net.Conn, clientNumber int, getArchive bool) {

	tmpNeededSubscriptions := neededSubscriptions
	if getArchive {
		tmpNeededSubscriptions = len(archiveInterests)
	}
	for i := 0; i < tmpNeededSubscriptions; i++ {
		//client receives tweets
		tweetsLengthBytes := make([]byte, 4)

		_, err := leaderConn.Read(tweetsLengthBytes)
		if err != nil {
			panic(err)
		}
		tweetsLength := byteToInt(tweetsLengthBytes)

		tweets := make([]byte, tweetsLength)
		_, err = leaderConn.Read(tweets)
		if err != nil {
			panic(err)
		}

		//expand sharedSecret so it is of right length
		expandBy := len(tweets) / 32
		expandedSharedSecrets := make([][]byte, 2)

		for i := 0; i < 2; i++ {
			for j := 0; j < expandBy; j++ {
				expandedSharedSecrets[i] = append(expandedSharedSecrets[i], sharedSecret[i][:]...)
			}
			//fmt.Println(expandedSharedSecrets[i])
		}

		//xors the received messge into the message to display
		for i := 0; i < 2; i++ {
			lib.Xor(expandedSharedSecrets[i][:], tweets)
		}

		//tweets can be displayed
		fmt.Println("final result: ", string(tweets))
	}
}

//creates a shared secret for each server
func createSharedSecret() [numClients][2][32]byte {
	var tmpSharedSecret [numClients][2][32]byte
	for i := 0; i < numClients; i++ {
		for j := 0; j < 2; j++ {
			_, err := rand.Read(tmpSharedSecret[i][j][:])
			if err != nil {
				panic("couldn't get randomness for sharedSecret!")
			}
		}
	}

	return tmpSharedSecret
}

//generates a topicOfInterest Array with random values
//todo! fix
func addFakeInterests(topicsOfInterest []int) []int {
	length := neededSubscriptions
	fakeTopicsOfInterest := make([]int, length)

	//fills the array with unique random ascending values in range with len(topicList)
	for index := 0; index < length; index++ {
		min := (index * (len(topicList) / length)) + 1
		max := ((index + 1) * (len(topicList) / length))
		//fmt.Println("max", max, "list", len(topicList))
		if max == len(topicList)-1 {
			max++
		}
		//fmt.Println("i", index, min, max, len(topicList))
		bigNumber, err := rand.Int(rand.Reader, big.NewInt(int64(max-min+1)))
		if err != nil {
			panic(err)
		}
		var number int = int(bigNumber.Int64()) + min

		fakeTopicsOfInterest[index] = number
	}
	//adds unique and new random numbers to topicOfInterests until length is satisfied
	for _, number := range fakeTopicsOfInterest {
		if !inList(number, topicsOfInterest) {
			topicsOfInterest = append(topicsOfInterest, number)
		}
		if len(topicsOfInterest) == length {
			break
		}
	}

	return topicsOfInterest
}

func inList(number int, list []int) bool {
	for _, element := range list {
		if element == number {
			return true
		}
	}
	return false
}

func receiveTopicLists(leaderConn net.Conn) {
	for i := 0; i < 2; i++ {
		topicListLength := make([]byte, 4)
		_, err := leaderConn.Read(topicListLength)
		if err != nil {
			panic(err)
		}

		recTopicList := make([]byte, byteToInt(topicListLength))
		_, err = leaderConn.Read(recTopicList[:])
		if err != nil {
			panic(err)
		}

		var tmpTopicList []string
		arrayReader := bytes.NewReader(recTopicList[:])
		json.NewDecoder(arrayReader).Decode(&tmpTopicList)
		if i == 0 {
			topicList = tmpTopicList
		} else {
			archiveTopicList = tmpTopicList
		}
	}

}

func intToByte(myInt int) (retBytes []byte) {
	retBytes = make([]byte, 4)
	retBytes[3] = byte((myInt >> 24) & 0xff)
	retBytes[2] = byte((myInt >> 16) & 0xff)
	retBytes[1] = byte((myInt >> 8) & 0xff)
	retBytes[0] = byte(myInt & 0xff)
	return
}

func byteToInt(myBytes []byte) (x int) {
	x = int(myBytes[3])<<24 + int(myBytes[2])<<16 + int(myBytes[1])<<8 + int(myBytes[0])
	return
}