package main

import (
	"crypto"
	"crypto/sha256"
	"fmt"
	"os"
	"runtime"
	"strconv"
	"sync"
	"time"

	"xdx.jelly/xgcl/rsa"
)

var wg sync.WaitGroup

func sign(sk *rsa.PrivateKey, digestData []byte, work, done chan struct{}) {
	for {
		select {
		case <-work:
			_, _ = rsa.PKCS1v15{}.Sign(sk, crypto.SHA256, digestData)
		case <-done:
			wg.Done()
			return
		}
	}

}

func main() {
	msg := []byte("message")
	digestData := sha256.Sum256(msg)

	sk, _, _ := rsa.PKCS1v15{}.GenerateKeyPairRSA(2048)
	threads := runtime.NumCPU()

	loops := 1000
	work := make(chan struct{})
	done := make(chan struct{})

	if len(os.Args) >= 2 {
		threads, _ = strconv.Atoi(os.Args[1])
	}

	if len(os.Args) >= 3 {
		loops, _ = strconv.Atoi(os.Args[2])
	}

	for i := 0; i < threads; i++ {
		wg.Add(1)
		go sign(sk, digestData[:], work, done)
	}

	for i := 0; i < 1000; i++ {
		// warm up
		work <- struct{}{}
	}

	start := time.Now()
	for i := 0; i < threads*loops; i++ {
		work <- struct{}{}
	}
	close(done)
	wg.Wait()
	elapsed := time.Since(start)
	cnt := threads * loops
	fmt.Printf("Sign %d times on %d cores\n", threads*loops, threads)
	fmt.Printf("Used time: %d ms, %d pcs/s\n", elapsed.Milliseconds(), int(float64(cnt)/float64(elapsed.Milliseconds())*1000))

}