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"
	mr "math/rand"
	"net"
	"sort"
	"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 = 1

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 dbWriteSize 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...)
		}
		tweet = append(tweet, []byte(";")[0])
		length := len(tweet)
		for i := length; i < dataLen; i++ {
			//todo! replace with random chars
			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 {

			//gets current dbWriteSize from leader
			dbWriteSizeBytes := make([]byte, 4)
			_, err = leaderConn.Read(dbWriteSizeBytes)
			if err != nil {
				panic(err)
			}
			dbWriteSize = byteToInt(dbWriteSizeBytes)

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

			//request virtualAddress from leader via pirQuery
			encryptedQueryLeader, encryptedQueryFollower := createAuditPIRQuery(clientNumber)
			sendQuerys(encryptedQueryLeader, encryptedQueryFollower, leaderConn, false)
			pos := receiveTweets(sharedSecret[clientNumber], leaderConn, clientNumber, false, true)
			fmt.Println(pos)

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

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

			//change
			C.prepQuery(C.int(pos), C.int(dbWriteSize), (*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)
			}

			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, 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, false)
			}

		} 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 {
		tmptopicsOfInterest = addFakeInterests(neededSubscriptions, len(tmpTopicList), tmptopicsOfInterest, false)
	}

	for topic, position := range tmptopicsOfInterest {
		//change used to be positon-1
		topicsOfInterestAsBytes[topic][position] = 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, doAuditing bool) int {

	tmpNeededSubscriptions := neededSubscriptions
	if getArchive {
		tmpNeededSubscriptions = len(archiveInterests)
	}
	if doAuditing {
		tmpNeededSubscriptions = 1
	}
	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][:]...)
			}
		}

		//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))

		if doAuditing {
			return byteToInt(tweets)
		}
	}
	return 0
}

//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
}

func createAuditPIRQuery(clientNumber int) ([]byte, []byte) {
	//pirQuery [serverAmount][dbWriteSize]byte
	pirQuerys := make([][]byte, 2)
	for i := range pirQuerys {
		pirQuerys[i] = make([]byte, dbWriteSize)
	}

	//for leader
	//pirQuery will be filled with random bits
	for index := range pirQuerys[0] {
		bit := mr.Intn(2)
		pirQuerys[0][index] = byte(bit)
	}
	copy(pirQuerys[1], pirQuerys[0])

	//the positon the virtual address will be taken from
	pos := mr.Intn(dbWriteSize)
	pirQuerys[0][pos] = 1
	pirQuerys[1][pos] = 0

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

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

	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
}

//generates a topicOfInterest array with random values
func addFakeInterests(length, max int, topicsOfInterest []int, doAuditing bool) []int {
	fakeTopicsOfInterest := make([]int, length)
	maxInt := max

	//fills the array with unique random ascending values ranging from 0 to max
	for i := 0; i < length; i++ {
		fakeTopicsOfInterest[i] = mr.Intn(maxInt)

		for j := 0; j < i; j++ {
			if fakeTopicsOfInterest[i] == fakeTopicsOfInterest[j] {
				i--
				break
			}
		}
	}

	if doAuditing {
		sort.Ints(fakeTopicsOfInterest)
		return fakeTopicsOfInterest
	}

	//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) == neededSubscriptions {
			break
		}
	}

	sort.Ints(topicsOfInterest)

	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
}