package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

var query string = `
query {
	me {
		id
		canonicalName
		sshKeys {
			cursor
			results {
				fingerprint
			}
		}
	}
}
`

var errors sync.Map

func doReqs(wg *sync.WaitGroup, tok string, nsuccess, nfail *int64, duration *int64) {
	defer wg.Done()
	start := time.Now()
	defer func() {
		elapsed := time.Now().Sub(start)
		atomic.AddInt64(duration, elapsed.Milliseconds())
	}()

	client := http.Client{}
	body, err := json.Marshal(&struct{
		Query string `json:"query"`
	}{
		Query: query,
	})
	if err != nil {
		panic(err)
	}
	req, err := http.NewRequest("POST", "http://localhost:5100/query",
		bytes.NewReader(body))
	if err != nil {
		panic(err)
	}
	req.Header.Add("Authorization", "Bearer " + tok)
	req.Header.Add("Content-Type", "application/json")
	if resp, err := client.Do(req); err != nil {
		atomic.AddInt64(nfail, 1)
		errors.Store(fmt.Sprintf("%v", err), nil)
	} else {
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			atomic.AddInt64(nsuccess, 1)
		} else {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return
			}
			errors.Store(string(body), nil)
			atomic.AddInt64(nfail, 1)
		}
	}
}

func main() {
	batch, _ := strconv.Atoi(os.Args[2])
	log.Printf("Sending batch of %d requests", batch)
	var nsuccess, nfail, duration int64
	var wg sync.WaitGroup
	wg.Add(batch)
	for i := 0; i < batch; i += 1 {
		go doReqs(&wg, os.Args[1], &nsuccess, &nfail, &duration)
	}
	time.Sleep(1 * time.Second)
	wg.Wait()
	log.Printf("%d success, %d failure; %dms avg",
		nsuccess, nfail, duration / int64(batch))
	errors.Range(func(key, value interface{}) bool {
		body, _ := key.(string)
		log.Printf("Sample error: %s", body)
		return true
	})
}