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"

//sssssssssssss

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

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

const leader string = "127.0.0.1:4442"

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

//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 followerPrivateKey *[32]byte
var followerPublicKey *[32]byte
var leaderPublicKey *[32]byte

var numThreads int = 12

//has to be dividable by 32
var dataLength int = 32
var numRows int = 2

var round int = 1
var startTime time.Time

var maxTimePerRound time.Duration = 5 * time.Second

func main() {

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

	followerPrivateKey = generatedPrivateKey
	followerPublicKey = generatedPublicKey

	/*
		if len(os.Args) != 4 {
			fmt.Println("try again: numThreads, dataLength, numRows")
			return
		}
		numThreads, _ = strconv.Atoi(os.Args[2])
		dataLength, _ = strconv.Atoi(os.Args[3])
		numRows, _ = strconv.Atoi(os.Args[4])
	*/

	C.initializeServer(C.int(numThreads))

	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", ":4443", config)
	if err != nil {
		panic(err)
	}
	defer lnLeader.Close()

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

	//send publicKey to leader
	_, err = leaderConnection.Write(followerPublicKey[:])
	if err != nil {
		panic(err)
	}

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

	leaderPublicKey = &tmpLeaderPubKey

	//setup ends here

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

	for {

		//phase1
		fmt.Println("phase1")

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

		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)
		}
		wg.Wait()

		//phase2
		fmt.Println("phase2")
		leaderConnection, err := lnLeader.Accept()
		if err != nil {
			panic(err)
		}
		leaderConnection.SetDeadline(time.Time{})

		phase2(leaderConnection)

		//phase3

		fmt.Println("phase3")

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

		wg.Wait()

		lib.CleanUpdbR(round)

		round++

	}
}

func phase1(id int, leaderWorkerConnection net.Conn, m sync.Mutex, wg *sync.WaitGroup) {

	gotClient := make([]byte, 1)

	for {
		_, err := leaderWorkerConnection.Read(gotClient)
		if err != nil {
			panic(err)
		}

		//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 {
			panic(err)
		}
		clientPublicKey = &tmpClientPublicKey

		//gets dpfQuery from leader
		dpfLengthBytes := make([]byte, 4)
		_, err = leaderWorkerConnection.Read(dpfLengthBytes)
		if err != nil {
			panic(err)
		}
		dpfLength := byteToInt(dpfLengthBytes)

		dpfQueryBEncrypted := make([]byte, dpfLength)

		_, err = leaderWorkerConnection.Read(dpfQueryBEncrypted)
		if err != nil {
			panic(err)
		}

		//auditing starts here

		//gets seeed from leader
		seed := make([]byte, 16)
		_, err = leaderWorkerConnection.Read(seed)
		if err != nil {
			panic(err)
		}

		//receive client audit result
		auditLengthBytes := make([]byte, 4)
		_, err = leaderWorkerConnection.Read(auditLengthBytes)
		if err != nil {
			panic(err)
		}

		auditLength := 200 //byteToInt(auditLengthBytes)
		clientAuditEncrypted := make([]byte, auditLength)

		_, err = leaderWorkerConnection.Read(clientAuditEncrypted)
		if err != nil {
			panic(err)
		}

		//decrypts the clients audit result
		var decryptNonce [24]byte
		copy(decryptNonce[:], clientAuditEncrypted[:24])
		clientAuditB, ok := box.Open(nil, clientAuditEncrypted[24:], &decryptNonce, clientPublicKey, followerPrivateKey)
		if !ok {
			panic("clientAuditB decryption not ok")
		}

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

		vector := make([]byte, dbSize*16)
		//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(i))
			//fmt.Println(i, pos)
			v := C.evalDPF(C.ctx[id], (*C.uchar)(&dpfQueryB[0]), pos, C.int(ds), (*C.uchar)(&dataShare[0]))
			fmt.Println("")
			fmt.Println("v", i, v)
			copy(vector[i*16:(i+1)*16], C.GoBytes(unsafe.Pointer(&v), 16))
			for j := 0; j < ds; j++ {
				db[i][j] = db[i][j] ^ dataShare[j]
			}
		}

		//prepare for audit
		mVal := make([]byte, 16)
		cVal := make([]byte, 16)

		C.serverSetupProof(C.ctx[id], (*C.uchar)(&seed[0]), C.dbSize, (*C.uchar)(&vector[0]), (*C.uchar)(&mVal[0]), (*C.uchar)(&cVal[0]))

		//compute audit query
		auditResultB := make([]byte, 96)
		C.serverComputeQuery(C.ctx[id], (*C.uchar)(&seed[0]), (*C.uchar)(&mVal[0]), (*C.uchar)(&cVal[0]), (*C.uchar)(&clientAuditB[0]), (*C.uchar)(&auditResultB[0]))

		//encrypt follower audit result
		var nonce [24]byte
		//fill nonce with randomness
		_, err = rand.Read(nonce[:])
		if err != nil {
			panic("couldn't get randomness for nonce!")
		}
		auditResultBEncrypted := box.Seal(nonce[:], auditResultB, &nonce, leaderPublicKey, followerPrivateKey)

		//gets audit result from leader
		auditResultAEncryptedLengthBytes := make([]byte, 4)
		_, err = leaderWorkerConnection.Read(auditResultAEncryptedLengthBytes)
		if err != nil {
			panic(err)
		}

		auditResultAEncryptedLength := byteToInt(auditResultAEncryptedLengthBytes)

		auditResultAEncrypted := make([]byte, auditResultAEncryptedLength)
		_, err = leaderWorkerConnection.Read(auditResultAEncrypted)
		if err != nil {
			panic(err)
		}

		//send follower audit result to leader
		auditResultBEncryptedLengthBytes := intToByte(len(auditResultBEncrypted))
		_, err = leaderWorkerConnection.Write(auditResultBEncryptedLengthBytes)
		if err != nil {
			panic(err)
		}

		_, err = leaderWorkerConnection.Write(auditResultBEncrypted)
		if err != nil {
			panic(err)
		}

		//decrypts the audit result from leader
		copy(decryptNonce[:], auditResultAEncrypted[:24])
		auditResultA, ok := box.Open(nil, auditResultAEncrypted[24:], &decryptNonce, leaderPublicKey, followerPrivateKey)
		if !ok {
			panic("auditResultA decryption not ok")
		}

		//compute audit result
		auditResult := int(C.serverVerifyProof((*C.uchar)(&auditResultA[0]), (*C.uchar)(&auditResultB[0])))

		if auditResult == 0 {
			//fmt.Println("audit failed")
		} /*else {
			fmt.Println("audit passed")
		}
		*/

		//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 := make([]byte, 16)
	_, err := leaderWorkerConnection.Read(seedLeader)
	if err != nil {
		panic(err)
	}

	//receive data from leader
	tmpdbLeader := make([][]byte, dbSize)
	for i := range tmpdbLeader {
		tmpdbLeader[i] = make([]byte, dataLength)
	}
	for i := 0; i < dbSize; i++ {
		_, err = leaderWorkerConnection.Read(tmpdbLeader[i])
		if err != nil {
			panic(err)
		}
	}

	//writes seed to leader
	_, err = leaderWorkerConnection.Write(seedFollower)
	if err != nil {
		panic(err)
	}

	//write data to leader
	for i := 0; i < dbSize; i++ {
		_, err = leaderWorkerConnection.Write(tmpdbFollower[i])
		if err != nil {
			panic(err)
		}
	}
	//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++ {
		_, err = leaderWorkerConnection.Write(C.GoBytes(unsafe.Pointer(ciphersFollowers[i]), 16))
		if err != nil {
			panic(err)
		}
	}

	//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][0] == 0 {
			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 {
							topic = topic + string(letter)
						}
					} else if position == 1 {
						text = text + string(letter)
					}
				}
			}

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

			tweets = append(tweets, tweet)
		}
	}

	//fmt.Println("tweets recovered: ", tweets)

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

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

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

	gotClient := make([]byte, 1)

	for {
		_, err := leaderWorkerConnection.Read(gotClient)
		if err != nil {
			panic(err)
		}

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

		subPhase := make([]byte, 1)
		_, err = leaderWorkerConnection.Read(subPhase)
		if err != nil {
			panic(err)
		}

		var clientPublicKey [32]byte
		_, err = leaderWorkerConnection.Read(clientPublicKey[:])
		if err != nil {
			panic(err)
		}

		//gets the client data
		clientKeys := clientData[clientPublicKey]

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

		getSendTweets(clientKeys, nil, leaderWorkerConnection)

		wantsArchive := make([]byte, 1)
		_, err = leaderWorkerConnection.Read(wantsArchive)
		if err != nil {
			panic(err)
		}

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

		}

		//saves clientKeys
		clientData[clientPublicKey] = clientKeys
	}
}

func getSendTweets(clientKeys clientKeys, archiveQuerys [][]byte, leaderWorkerConnection net.Conn) {
	tmpNeededSubscriptions := neededSubscriptions
	if archiveQuerys != nil {
		tmpNeededSubscriptions = len(archiveQuerys)
	}
	for i := 0; i < tmpNeededSubscriptions; i++ {
		//gets all requested tweets
		var tweets []byte
		if archiveQuerys == nil {
			tweets = lib.GetTweets(clientKeys.PirQuery[i], dataLength, 0)
		} else {
			tweets = lib.GetTweets(archiveQuerys[i], dataLength, 1)
		}

		//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[:]...)
		}

		fmt.Println(expandedSharedSecret)

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

		lib.Xor(tweets, expandedSharedSecret[:])

		//sends tweets to leader
		tweetsLengthBytes := intToByte(len(tweets))

		_, err := leaderWorkerConnection.Write(tweetsLengthBytes)
		if err != nil {
			panic(err)
		}

		_, err = leaderWorkerConnection.Write(tweets)
		if err != nil {
			panic(err)
		}

	}
}

func handlePirQuery(clientKeys clientKeys, leaderWorkerConnection net.Conn, subPhase int, clientPublicKey [32]byte) (clientKeys, [][]byte) {
	archiveNeededSubscriptions := make([]byte, 4)
	if subPhase == -1 {
		_, err := leaderWorkerConnection.Read(archiveNeededSubscriptions)
		if err != nil {
			panic(err)
		}
	}

	//gets the msg length
	msgLengthBytes := make([]byte, 4)
	_, err := leaderWorkerConnection.Read(msgLengthBytes)
	if err != nil {
		panic(err)
	}
	msgLength := byteToInt(msgLengthBytes)

	message := make([]byte, msgLength)

	//gets the message
	_, err = leaderWorkerConnection.Read(message)
	if err != nil {
		panic(err)
	}

	var decryptNonce [24]byte
	copy(decryptNonce[:], message[:24])
	decrypted, ok := box.Open(nil, message[24:], &decryptNonce, &clientPublicKey, followerPrivateKey)
	if !ok {
		panic("pirQuery decryption not ok")
	}
	//gets sharedSecret
	if subPhase == 0 {
		//bs!
		var newSharedSecret [32]byte
		for index := 0; index < 32; index++ {
			newSharedSecret[index] = decrypted[index]
		}
		clientKeys.SharedSecret = newSharedSecret
		decrypted = decrypted[32:]

		//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
	pirQueryFlattened := decrypted
	tmpNeededSubscriptions := neededSubscriptions
	tmpTopicAmount := topicAmount
	if subPhase == -1 {
		tmpNeededSubscriptions = byteToInt(archiveNeededSubscriptions)
		_, tmpTopicAmount = lib.GetTopicList(1)
	}
	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
}

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
}