package main
import (
    "fmt"
    "net/rpc"
    "log"
    "time"
    "sync"
    "math/rand"
    "os"
    "flag"
    "strconv"
)

type putStat struct {
    mu sync.Mutex
    successPut int
    failPut int
}

type getStat struct {
    mu sync.Mutex
    successGet int
    failGet int
}

type getRequest struct {
    client *rpc.Client
    key int
}

type putRequest struct {
    client *rpc.Client
    key int
    value string
}

 // keep track of success and failure of put(key, value)
func countPutStat(putInfo *putStat, isAdd bool) {
    if isAdd == true {
        putInfo.successPut += 1
    } else {
        putInfo.failPut += 1
    }
}

// keep track of success and failure of get(key)
func countGetStat(getInfo *getStat, value string) {
    if value != "" {
        getInfo.successGet += 1
    } else {
        getInfo.failGet += 1
    }
}

//for now hard code the probability of get and put
func getOrPut(getProb float64, putProb float64) int {
    var method int
    rand.Seed(time.Now().UnixNano())
    prob := rand.Float64()
    if prob <= 0.4 {
        method = 1 // put
    } else {
        method = 0 // get
    }
    return method
}

// poplulate getRequest struct
func setGetRequest(numOp int, clients [3]*rpc.Client, getRequestInfo *getRequest) (*getRequest){
    rand.Seed(time.Now().UnixNano())
    key := rand.Intn(numOp)
    switch node := key % 3; node { // indicate which node to fire request to by using a specific client of that node
        case 0:
            getRequestInfo.client = clients[0]
        case 1:
            getRequestInfo.client = clients[1]
        case 2:
            getRequestInfo.client = clients[2]
    }
    getRequestInfo.key = key
    return getRequestInfo
}

// poplulate putRequest struct
func setPutRequest(numOp int, clients [3]*rpc.Client, value string, putRequestInfo *putRequest) {
    rand.Seed(time.Now().UnixNano())
    key := rand.Intn(numOp)
    switch node := key % 3; node { // indicate which node to fire request to by using a specific client of that node
        case 0:
            putRequestInfo.client = clients[0]
        case 1:
            putRequestInfo.client = clients[1]
        case 2:
            putRequestInfo.client = clients[2]
    }
    putRequestInfo.key = key
    putRequestInfo.value = value
}

func put(client *rpc.Client, key int, value string, putInfo *putStat, wg *sync.WaitGroup, timeRecords chan<- time.Duration) {
    defer timeTrack(time.Now(), timeRecords)
    defer wg.Done()
    var isAdd bool // return true if put put new (key, value)
                   // return false if put update value of the already existing key
    client.Call("DHT.Put", &PutArgs{key, value}, &isAdd)
    putInfo.mu.Lock()
    countPutStat(putInfo, isAdd)
    putInfo.mu.Unlock()
}

func get(client *rpc.Client, key int, getInfo *getStat, wg *sync.WaitGroup, timeRecords chan<- time.Duration) string {
    defer timeTrack(time.Now(), timeRecords)
    defer wg.Done()
    var value string
    client.Call("DHT.Get", &GetArgs{key}, &value)
    getInfo.mu.Lock()
    countGetStat(getInfo, value)
    getInfo.mu.Unlock()
    return value
}

// time each operation
func timeTrack(start time.Time, timeRecords chan<- time.Duration) {
    elapsed := time.Since(start)
    timeRecords <- elapsed
}

func getCommandLine() int {
    flag.Parse()
    s := flag.Arg(0)
    i, err := strconv.Atoi(s)
    if err != nil {
        fmt.Println(err)
        os.Exit(2)
    }
    return i
}

func calLatency(durations []time.Duration, numOp float64) float64 {
    var latency float64
    for _, elem := range durations {
        latency += elem.Seconds()
    }
    return latency/numOp
}


// from medium article: https://medium.com/@kpbird/golang-generate-fixed-size-random-string-dd6dbd5e63c0
func RandStringRunes(n int) string {
    var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
    b := make([]rune, n)
    for i := range b {
        b[i] = letterRunes[rand.Intn(len(letterRunes))]
    }
    return string(b)
}

func main() {
    // connect to all 3 servers
    client0, err := rpc.DialHTTP("tcp", "142.93.123.124" + ":5000")
    client1, _ := rpc.DialHTTP("tcp", "142.93.204.111" + ":5001")
    client2, _ := rpc.DialHTTP("tcp", "142.93.126.155" + ":5002")

    if err != nil {
            log.Fatal("dialing:", err)
    }
    clients := [3]*rpc.Client{client0, client1, client2}
    getRequestInfo := new(getRequest)
    putRequestInfo := new(putRequest)
    getInfo := new(getStat)
    putInfo := new(putStat)
    timeRecords := make(chan time.Duration)

    // get command line argument for number of operations
    numOp := getCommandLine()
    var wg sync.WaitGroup
    start := time.Now()
    wg.Add(numOp)
    for i := 0; i < numOp; i++ {
        method := getOrPut(0.6, 0.4) // 0 = get, 1 = put
        if method == 0 { // get
            setGetRequest(numOp, clients, getRequestInfo)
            go get(getRequestInfo.client, getRequestInfo.key, getInfo, &wg, timeRecords)
        } else { // put
            setPutRequest(numOp, clients, RandStringRunes(20), putRequestInfo)
            go put(putRequestInfo.client, putRequestInfo.key, putRequestInfo.value, putInfo, &wg, timeRecords)
        }
    }
    wg.Wait()
    elapsed := time.Since(start)
    fmt.Println("The whole operation took:", elapsed.Seconds(), "s")
    fmt.Println("Throughput:", float64(numOp)/elapsed.Seconds(), "operations/s")
    var durations []time.Duration
    for i := 0; i < numOp; i++ {
        durations = append(durations, <-timeRecords)
    }
    var latency = calLatency(durations, float64(numOp))
    fmt.Println("Latency:", latency, "s/operation")

    // print map of each server after put
    var m = map[int]string{}
    err = clients[0].Call("DHT.PrintDHT", 0, &m)
    err = clients[1].Call("DHT.PrintDHT", 0, &m)
    err = clients[2].Call("DHT.PrintDHT", 0, &m)

    fmt.Println("Get sucess:", getInfo.successGet)
    fmt.Println("Get failure:", getInfo.failGet)
    fmt.Println("Put success:", putInfo.successPut)
    fmt.Println("Put failure:", putInfo.failPut)
}