package main

import (
	"context"
	"fmt"
	"os"
	"sort"

	"github.com/open-policy-agent/opa/metrics"

	"github.com/open-policy-agent/opa/rego"
)

func printMetrics(m metrics.Metrics) {
	all := m.All()
	var keys []string
	for k := range all {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	for _, k := range keys {
		fmt.Printf("%s: %+v\n", k, all[k].(map[string]interface{})["mean"])
	}
}

func main() {
	ctx := context.Background()

	// Raw input data that will be used in the first evaluation
	input := map[string]interface{}{"x": 2}

	m := metrics.New()
	bh := metrics.New()
	ah := metrics.New()

	iterations := 5000

	for i := 0; i < iterations; i++ {
		// Create a simple query over the input.
		r := rego.New(
			rego.Module("test.rego", `
			package test
			p {
				input.x = 1
			}
			`),
			rego.Query("data.test.p"),
			rego.Input(input),
			rego.Metrics(m),
		)

		// Run evaluation.
		_, err := r.Eval(ctx)

		if err != nil {
			fmt.Printf("error: %s", err.Error())
			os.Exit(1)
		}

		// Update the "before" histogram
		for k, v := range m.All() {
			bh.Histogram(k).Update(v.(int64))
		}
		m.Clear()

		// Update input
		input["x"] = 1

		r.Update(
			rego.Input(input),
		)

		// Evaluate again
		_, err = r.Eval(ctx)

		if err != nil {
			fmt.Printf("error: %s", err.Error())
			os.Exit(1)
		}

		// Update the "after" histogram
		for k, v := range m.All() {
			ah.Histogram(k).Update(v.(int64))
		}
		m.Clear()
	}

	fmt.Println("iterations:", iterations)
	fmt.Println("--- Baseline (mean)")
	printMetrics(bh)
	fmt.Println("--- With Update() (mean)")
	printMetrics(ah)
}