package main

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

import (
	"2PPS/lib"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
	"unsafe"

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

//this stores all neccessary information for each client
type clientKeys struct {
	SharedSecret [32]byte
	PirQuery     [][]byte
}

//uses clients publicKey as key
var clientData = make(map[[32]byte]clientKeys)
var topicList []byte
var topicAmount int
var archiveTopicAmount int

var followerPrivateKey *[32]byte
var followerPublicKey *[32]byte
var leaderPublicKey *[32]byte

//needs to be changed at leader/follower/client at the same time
const dataLength = 256
const numThreads = 12

//Maximum Transport Unit
const mtu int = 1100

var dbWriteSize int
var neededSubscriptions int

var round int
var startTime time.Time

func main() {

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

	followerPrivateKey = generatedPrivateKey
	followerPublicKey = generatedPublicKey

	followerConnectionPrivateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		panic(err)
	}

	// Generate a pem block with the private key
	keyPem := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(followerConnectionPrivateKey),
	})

	tml := x509.Certificate{
		// you can add any attr that you need
		NotBefore: time.Now(),
		NotAfter:  time.Now().AddDate(5, 0, 0),
		// you have to generate a different serial number each execution
		SerialNumber: big.NewInt(123123),
		Subject: pkix.Name{
			CommonName:   "New Name",
			Organization: []string{"New Org."},
		},
		BasicConstraintsValid: true,
	}
	cert, err := x509.CreateCertificate(rand.Reader, &tml, &tml, &followerConnectionPrivateKey.PublicKey, followerConnectionPrivateKey)
	if err != nil {
		panic(err)
	}

	// Generate a pem block with the certificate
	certPem := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: cert,
	})

	tlsCert, err := tls.X509KeyPair(certPem, keyPem)
	if err != nil {
		panic(err)
	}

	config := &tls.Config{Certificates: []tls.Certificate{tlsCert}}

	fmt.Println("start leader")

	//listens for leader
	lnLeader, err := tls.Listen("tcp", ":4442", config)
	if err != nil {
		panic(err)
	}
	defer lnLeader.Close()

	leaderConnection, err := lnLeader.Accept()
	if err != nil {
		panic(err)
	}

	//send publicKey to leader
	writeTo(leaderConnection, followerPublicKey[:])

	//receives leader PublicKey
	var tmpLeaderPubKey [32]byte
	_, err = leaderConnection.Read(tmpLeaderPubKey[:])
	if err != nil {
		panic(err)
	}

	neededSubscriptionsBytes := readFrom(leaderConnection, 4)
	neededSubscriptions = byteToInt(neededSubscriptionsBytes)

	dbWriteSizeBytes := readFrom(leaderConnection, 4)
	dbWriteSize = byteToInt(dbWriteSizeBytes)

	leaderPublicKey = &tmpLeaderPubKey

	C.initializeServer(C.int(numThreads))

	//setup ends here

	//locks access to DB
	m := &sync.RWMutex{}
	wg := &sync.WaitGroup{}

	for {

		round++
		fmt.Println("Phase 1 Round", round)

		//create write db for this round
		for i := 0; i < dbWriteSize; i++ {
			C.createDb(C.int(0), C.int(dataLength))
		}

		//receives the virtualAddresses
		virtualAddresses := make([]int, dbWriteSize+1)
		for i := 0; i <= dbWriteSize; i++ {
			virtualAddress := readFrom(leaderConnection, 4)
			virtualAddresses[i] = byteToInt(virtualAddress)
		}

		for i := 0; i < numThreads; i++ {
			wg.Add(1)
			leaderConnection, err := lnLeader.Accept()
			if err != nil {
				panic(err)
			}
			leaderConnection.SetDeadline(time.Time{})
			startTime = time.Now()
			go phase1(i, leaderConnection, m, wg, virtualAddresses)
		}
		wg.Wait()

		//Phase 2
		leaderConnection, err := lnLeader.Accept()
		if err != nil {
			panic(err)
		}
		leaderConnection.SetDeadline(time.Time{})

		phase2(leaderConnection)

		//Phase 3

		if round == 1 {
			//addTestTweets()
		}

		//no tweets -> continue to phase 1 and mb get tweets
		topicList, topicAmount = lib.GetTopicList(0)
		if len(topicList) == 0 {
			continue
		}

		for i := 0; i < numThreads; i++ {
			wg.Add(1)
			leaderConnection, err := lnLeader.Accept()
			if err != nil {
				panic(err)
			}
			leaderConnection.SetDeadline(time.Time{})
			startTime = time.Now()
			go phase3(leaderConnection, wg, m)
		}

		wg.Wait()

		lib.CleanUpdbR(round)
	}
}

func phase1(id int, leaderWorkerConnection net.Conn, m *sync.RWMutex, wg *sync.WaitGroup, virtualAddresses []int) {

	for {
		gotClient := readFrom(leaderWorkerConnection, 1)

		//this worker is done
		if gotClient[0] == 0 {
			wg.Done()
			return
		}

		//setup the worker-specific db
		dbSize := int(C.dbSize)
		db := make([][]byte, dbSize)
		for i := 0; i < dbSize; i++ {
			db[i] = make([]byte, int(C.db[i].dataSize))
		}

		//gets clients publicKey
		var clientPublicKey *[32]byte
		var tmpClientPublicKey [32]byte
		_, err := leaderWorkerConnection.Read(tmpClientPublicKey[:])
		if err != nil {
			fmt.Println("no error handling")
			panic(err)
		}
		clientPublicKey = &tmpClientPublicKey

		m.RLock()
		clientKeys := clientData[tmpClientPublicKey]
		m.RUnlock()

		clientKeys, pirQuery, errorBool := handlePirQuery(clientKeys, leaderWorkerConnection, 0, tmpClientPublicKey, true)
		if errorBool {
			continue
		}
		getSendVirtualAddress(pirQuery[0], virtualAddresses, clientKeys.SharedSecret, leaderWorkerConnection)

		m.Lock()
		clientData[*clientPublicKey] = clientKeys
		m.Unlock()

		//gets dpfQuery from leader
		dpfLengthBytes, errorBool := readFromWError(leaderWorkerConnection, 4)
		if errorBool {
			continue
		}
		dpfLength := byteToInt(dpfLengthBytes)

		dpfQueryBEncrypted, errorBool := readFromWError(leaderWorkerConnection, dpfLength)
		if errorBool {
			continue
		}

		//decrypt dpfQueryB for sorting into db
		var decryptNonce [24]byte
		copy(decryptNonce[:], dpfQueryBEncrypted[:24])
		dpfQueryB, ok := box.Open(nil, dpfQueryBEncrypted[24:], &decryptNonce, clientPublicKey, followerPrivateKey)
		if !ok {
			panic("dpfQueryB decryption not ok")
		}

		ds := int(C.db[0].dataSize)
		dataShareFollower := make([]byte, ds)
		pos := C.getUint128_t(C.int(virtualAddresses[dbWriteSize]))
		C.evalDPF(C.ctx[id], (*C.uchar)(&dpfQueryB[0]), pos, C.int(ds), (*C.uchar)(&dataShareFollower[0]))

		writeTo(leaderWorkerConnection, dataShareFollower)

		dataShareLeader, errorBool := readFromWError(leaderWorkerConnection, ds)
		if errorBool {
			continue
		}

		auditXOR := make([]byte, ds)
		passedAudit := true
		for i := 0; i < ds; i++ {
			auditXOR[i] = dataShareLeader[i] ^ dataShareFollower[i]

			//client tried to write to a position that is not a virtuallAddress
			if auditXOR[i] != 0 {
				passedAudit = false
			}
		}

		if passedAudit {
			//run dpf, xor into local db
			for i := 0; i < dbSize; i++ {
				ds := int(C.db[i].dataSize)
				dataShare := make([]byte, ds)
				pos := C.getUint128_t(C.int(virtualAddresses[i]))
				C.evalDPF(C.ctx[id], (*C.uchar)(&dpfQueryB[0]), pos, C.int(ds), (*C.uchar)(&dataShare[0]))
				for j := 0; j < ds; j++ {
					db[i][j] = db[i][j] ^ dataShare[j]
				}
			}

			//xor the worker's DB into the main DB
			for i := 0; i < dbSize; i++ {
				m.Lock()
				C.xorIn(C.int(i), (*C.uchar)(&db[i][0]))
				m.Unlock()
			}
		}
	}
}

func phase2(leaderWorkerConnection net.Conn) {

	//gets current seed
	seedFollower := make([]byte, 16)
	C.readSeed((*C.uchar)(&seedFollower[0]))

	//get data
	dbSize := int(C.dbSize)
	tmpdbFollower := make([][]byte, dbSize)
	for i := range tmpdbFollower {
		tmpdbFollower[i] = make([]byte, dataLength)
	}
	for i := 0; i < dbSize; i++ {
		C.readData(C.int(i), (*C.uchar)(&tmpdbFollower[i][0]))
	}

	//receive seed from leader
	seedLeader := readFrom(leaderWorkerConnection, 16)

	//receive data from leader
	tmpdbLeader := make([][]byte, dbSize)
	for i := range tmpdbLeader {
		tmpdbLeader[i] = make([]byte, dataLength)
	}
	for i := 0; i < dbSize; i++ {
		tmpdbLeader[i] = readFrom(leaderWorkerConnection, dataLength)
	}

	//writes seed to leader
	writeTo(leaderWorkerConnection, seedFollower)

	//write data to leader
	for i := 0; i < dbSize; i++ {
		writeTo(leaderWorkerConnection, tmpdbFollower[i])
	}
	//put together the db
	tmpdb := make([][]byte, dbSize)
	for i := range tmpdb {
		tmpdb[i] = make([]byte, dataLength)
	}

	//get own Ciphers
	ciphersFollowers := make([]*C.uchar, dbSize)
	for i := 0; i < dbSize; i++ {
		ciphersFollowers[i] = (*C.uchar)(C.malloc(16))
	}
	for i := 0; i < dbSize; i++ {
		C.getCipher(0, C.int(i), ciphersFollowers[i])
	}

	//receive ciphers from leader
	ciphersLeader := make([]byte, dbSize*16)
	for i := 0; i < dbSize; i++ {
		_, err := leaderWorkerConnection.Read(ciphersLeader[i*16:])
		if err != nil {
			panic(err)
		}
	}

	//send own Ciphers to leader
	for i := 0; i < dbSize; i++ {
		writeTo(leaderWorkerConnection, C.GoBytes(unsafe.Pointer(ciphersFollowers[i]), 16))
	}

	//put in ciphers from leader
	for i := 0; i < dbSize; i++ {
		C.putCipher(0, C.int(i), (*C.uchar)(&ciphersLeader[i*16]))
	}

	for i := 0; i < dbSize; i++ {
		C.decryptRow(C.int(i), (*C.uchar)(&tmpdb[i][0]), (*C.uchar)(&tmpdbLeader[i][0]), (*C.uchar)(&tmpdbFollower[i][0]), (*C.uchar)(&seedLeader[0]), (*C.uchar)(&seedFollower[0]))
	}

	var tweets []lib.Tweet
	for i := 0; i < dbSize; i++ {
		//discard cover message
		if tmpdb[i][1] == 0 {
			continue
		} else if -1 == strings.Index(string(tmpdb[i]), ";;") {
			continue
		} else {
			//reconstruct tweet
			var position int = 0
			var topics []string
			var topic string
			var text string
			for _, letter := range tmpdb[i] {
				if string(letter) == ";" {
					if topic != "" {
						topics = append(topics, topic)
						topic = ""
					}
					position++
				} else {
					if position == 0 {
						if string(letter) == "," {
							topics = append(topics, topic)
							topic = ""
						} else {
							//if topics are
							//int
							//topic = topic + fmt.Sprint(int(letter))
							//string
							topic = topic + string(letter)
						}
					} else if position == 1 {
						text = text + string(letter)
					}
				}
			}

			tweet := lib.Tweet{"", -1, topics, text, round}

			if text != "" {
				tweets = append(tweets, tweet)
			}
		}
	}

	//sort into read db
	lib.NewEntries(tweets, 0)

	//reset write db after the tweets were moved to read db
	C.resetDb()

	//gets current dbWriteSize from leader
	dbWriteSizeBytes := readFrom(leaderWorkerConnection, 4)
	dbWriteSize = byteToInt(dbWriteSizeBytes)

	lib.CleanUpdbR(round)

}

func addTestTweets() {
	//creates test tweets
	tweets := make([]lib.Tweet, 5)

	for i := range tweets {
		j := i
		if i == 1 {
			j = 0
		}
		text := "Text " + strconv.Itoa(i)
		var topics []string
		topics = append(topics, "Topic "+strconv.Itoa(j))
		tweets[i] = lib.Tweet{"", -1, topics, text, i}
	}

	lib.NewEntries(tweets, 0)
}

func phase3(leaderWorkerConnection net.Conn, wg *sync.WaitGroup, m *sync.RWMutex) {

	for {
		gotClient, errorBool := readFromWError(leaderWorkerConnection, 1)
		if errorBool {
			continue
		}

		//this worker is done
		if gotClient[0] == 0 {
			wg.Done()
			return
		}

		subPhase, errorBool := readFromWError(leaderWorkerConnection, 1)
		if errorBool {
			continue
		}

		var clientPublicKey [32]byte
		_, err := leaderWorkerConnection.Read(clientPublicKey[:])
		if err != nil {
			fmt.Println("no error handling")
			panic(err)
		}

		//gets the client data
		m.RLock()
		clientKeys := clientData[clientPublicKey]
		m.RUnlock()

		if subPhase[0] == 0 || subPhase[0] == 1 {
			clientKeys, _, errorBool = handlePirQuery(clientKeys, leaderWorkerConnection, int(subPhase[0]), clientPublicKey, false)
			if errorBool {
				continue
			}
		}

		getSendTweets(clientKeys, nil, leaderWorkerConnection, clientPublicKey)

		wantsArchive, errorBool := readFromWError(leaderWorkerConnection, 1)
		if errorBool {
			continue
		}

		if wantsArchive[0] == 1 {
			_, archiveQuerys, errorBool := handlePirQuery(clientKeys, leaderWorkerConnection, -1, clientPublicKey, false)
			if errorBool {
				continue
			}
			getSendTweets(clientKeys, archiveQuerys, leaderWorkerConnection, clientPublicKey)
		}

		//saves clientKeys
		m.Lock()
		clientData[clientPublicKey] = clientKeys
		m.Unlock()
	}
}

//gets tweet from db and sends them to leader
func getSendTweets(clientKeys clientKeys, archiveQuerys [][]byte, leaderWorkerConnection net.Conn, pubKey [32]byte) {
	tmpNeededSubscriptions := neededSubscriptions
	if tmpNeededSubscriptions > topicAmount {
		tmpNeededSubscriptions = topicAmount
	}
	if archiveQuerys != nil {
		tmpNeededSubscriptions = len(archiveQuerys)
		if tmpNeededSubscriptions > archiveTopicAmount {
			tmpNeededSubscriptions = archiveTopicAmount
		}
	}
	for i := 0; i < tmpNeededSubscriptions; i++ {
		//gets all requested tweets
		var tweets []byte
		if archiveQuerys == nil {
			tweets = lib.GetTweets(clientKeys.PirQuery[i], dataLength, 0, pubKey)
		} else {
			tweets = lib.GetTweets(archiveQuerys[i], dataLength, 1, pubKey)
		}
		//expand sharedSecret so it is of right length
		expandBy := len(tweets) / 32
		var expandedSharedSecret []byte
		for i := 0; i < expandBy; i++ {
			expandedSharedSecret = append(expandedSharedSecret, clientKeys.SharedSecret[:]...)
		}

		//Xor's sharedSecret with all tweets
		lib.Xor(expandedSharedSecret[:], tweets)

		//sends tweets to leader
		writeTo(leaderWorkerConnection, tweets)
	}
}

//returns true if client connection is lost
func handlePirQuery(clientKeys clientKeys, leaderWorkerConnection net.Conn, subPhase int, clientPublicKey [32]byte, doAuditing bool) (clientKeys, [][]byte, bool) {

	archiveNeededSubscriptions := make([]byte, 4)
	if subPhase == -1 {
		var errorBool bool
		archiveNeededSubscriptions, errorBool = readFromWError(leaderWorkerConnection, 4)
		if errorBool {
			return clientKeys, nil, true
		}
	}

	//gets the msg length
	msgLengthBytes, errorBool := readFromWError(leaderWorkerConnection, 4)
	if errorBool {
		return clientKeys, nil, true
	}
	msgLength := byteToInt(msgLengthBytes)

	//gets the message
	message, errorBool := readFromWError(leaderWorkerConnection, msgLength)
	if errorBool {
		return clientKeys, nil, true
	}

	var decryptNonce [24]byte
	copy(decryptNonce[:], message[:24])
	decrypted, ok := box.Open(nil, message[24:], &decryptNonce, &clientPublicKey, followerPrivateKey)
	if !ok {
		fmt.Println("pirQuery decryption not ok")
		return clientKeys, nil, true
	}

	//gets sharedSecret
	if subPhase == 0 {
		var newSharedSecret [32]byte
		for index := 0; index < 32; index++ {
			newSharedSecret[index] = decrypted[index]
		}
		clientKeys.SharedSecret = newSharedSecret
		decrypted = decrypted[32:]

		if doAuditing {
			result := make([][]byte, 1)
			result[0] = decrypted
			return clientKeys, result, false
		}

		//follower updates sharedSecret
	} else if subPhase == 1 {
		sharedSecret := clientKeys.SharedSecret
		sharedSecret = sha256.Sum256(sharedSecret[:])
		clientKeys.SharedSecret = sharedSecret
	}

	//follower expects pirQuery

	//transforms byteArray to ints of wanted topics
	tmpNeededSubscriptions := neededSubscriptions
	if tmpNeededSubscriptions > topicAmount {
		tmpNeededSubscriptions = topicAmount
	}
	tmpTopicAmount := topicAmount
	if subPhase == -1 {
		tmpNeededSubscriptions = byteToInt(archiveNeededSubscriptions)
		_, archiveTopicAmount = lib.GetTopicList(1)
		if tmpNeededSubscriptions > archiveTopicAmount {
			tmpNeededSubscriptions = archiveTopicAmount
		}
		tmpTopicAmount = archiveTopicAmount
	}

	pirQueryFlattened := decrypted
	pirQuerys := make([][]byte, tmpNeededSubscriptions)
	for i := range pirQuerys {
		pirQuerys[i] = make([]byte, tmpTopicAmount)
	}
	for i := 0; i < tmpNeededSubscriptions; i++ {
		pirQuerys[i] = pirQueryFlattened[i*tmpTopicAmount : (i+1)*tmpTopicAmount]
	}

	//sets the pirQuery for the client in case whe are not archiving
	if subPhase != -1 {
		clientKeys.PirQuery = pirQuerys
	}

	return clientKeys, pirQuerys, false
}

func getSendVirtualAddress(pirQuery []byte, virtualAddresses []int, sharedSecret [32]byte, leaderWorkerConnection net.Conn) {
	//xores all requested addresses into virtuallAddress
	virtualAddress := make([]byte, 4)

	for index, num := range pirQuery {
		if num == 1 {
			currentAddress := intToByte(virtualAddresses[index])
			for i := 0; i < 4; i++ {
				virtualAddress[i] = virtualAddress[i] ^ currentAddress[i]
			}
		}
	}

	//xores the sharedSecret
	for i := 0; i < 4; i++ {
		virtualAddress[i] = virtualAddress[i] ^ sharedSecret[i]
	}

	writeTo(leaderWorkerConnection, virtualAddress)
}

//sends the array to the connection
func writeTo(connection net.Conn, array []byte) {
	remainingLength := len(array)
	for remainingLength > 0 {
		if remainingLength >= mtu {
			_, err := connection.Write(array[:mtu])
			if err != nil {
				panic(err)
			}
			array = array[mtu:]
			remainingLength -= mtu
		} else {
			_, err := connection.Write(array)
			if err != nil {
				panic(err)
			}
			remainingLength = 0
		}
	}
}

//reads an array which is returned and of size "size" from the connection
//returned bool is one if connection to client was lost
func readFrom(connection net.Conn, size int) []byte {
	var array []byte
	remainingSize := size
	for remainingSize > 0 {
		var err error
		toAppend := make([]byte, mtu)
		if remainingSize > mtu {
			_, err = connection.Read(toAppend)
			array = append(array, toAppend...)
			remainingSize -= mtu
		} else {
			_, err = connection.Read(toAppend[:remainingSize])
			array = append(array, toAppend[:remainingSize]...)
			remainingSize = 0

		}
		if err != nil {
			panic(err)
		}
	}

	return array
}

//reads an array which is returned and of size "size" from the connection
//returns true if connection to client is lost
func readFromWError(connection net.Conn, size int) ([]byte, bool) {
	var array []byte
	remainingSize := size + 1
	for remainingSize > 0 {
		var err error
		toAppend := make([]byte, mtu)
		if remainingSize > mtu {
			_, err = connection.Read(toAppend)
			array = append(array, toAppend...)
			remainingSize -= mtu
		} else {
			_, err = connection.Read(toAppend[:remainingSize])
			array = append(array, toAppend[:remainingSize]...)
			remainingSize = 0
		}
		if err != nil {
			panic(err)
		}
	}
	if array[0] == 1 {
		return nil, true
	}
	return array[1:], false

	/*

		array := make([]byte, size+1)
		_, err := connection.Read(array)
		if err != nil {
			panic(err)
		}
		if array[0] == 1 {
			return nil, true
		}
		return array[1:], false

	*/
}

func transformBytesToStringArray(topicsAsBytes []byte) []string {
	var topics []string
	var topic string
	var position int = 0
	for _, letter := range topicsAsBytes {
		if string(letter) == "," {
			topics[position] = topic
			topic = ""
			position++
		} else {
			topic = topic + string(letter)
		}
	}
	return topics
}

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

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
}