import benchmark
import context
import os
import time
import x.executor

const default_drain_rounds = 16
const default_drain_jobs = 32
const default_sync_rounds = 32
const default_timeout_rounds = 8
const default_timeout_ms = 2
const default_producers = 16

fn main() {
	run_benchmarks() or {
		eprintln('x.executor benchmark failed: ${err.msg()}')
		exit(1)
	}
}

fn run_benchmarks() ! {
	drain_rounds := env_int('XEXECUTOR_BENCH_DRAIN_ROUNDS', default_drain_rounds, 1, 200)
	drain_jobs := env_int('XEXECUTOR_BENCH_DRAIN_JOBS', default_drain_jobs, 1, 512)
	sync_rounds := env_int('XEXECUTOR_BENCH_SYNC_ROUNDS', default_sync_rounds, 1, 500)
	timeout_rounds := env_int('XEXECUTOR_BENCH_TIMEOUT_ROUNDS', default_timeout_rounds, 1, 100)
	timeout_ms := env_int('XEXECUTOR_BENCH_TIMEOUT_MS', default_timeout_ms, 1, 100)
	producers := env_int('XEXECUTOR_BENCH_PRODUCERS', default_producers, 1, 256)

	println('x.executor cautious benchmark')
	println('Override sizes with XEXECUTOR_BENCH_* environment variables.')

	mut checksum := 0
	mut b := benchmark.start()
	checksum += bench_try_post_drain(drain_rounds, drain_jobs)!
	b.measure('try_post + drain_pending: rounds=${drain_rounds}, jobs=${drain_jobs}, checksum=${checksum}')

	checksum += bench_run_sync(sync_rounds)!
	b.measure('run_sync: rounds=${sync_rounds}, checksum=${checksum}')

	checksum += bench_bounded_timeout(timeout_rounds, timeout_ms)!
	b.measure('bounded admission timeout: rounds=${timeout_rounds}, timeout_ms=${timeout_ms}, checksum=${checksum}')

	checksum += bench_many_producers(producers)!
	b.measure('multi-producer post_with_context: producers=${producers}, checksum=${checksum}')
}

fn env_int(name string, default_value int, min_value int, max_value int) int {
	raw := os.getenv(name)
	if raw == '' {
		return default_value
	}
	mut value := raw.int()
	if value < min_value {
		value = min_value
	}
	if value > max_value {
		value = max_value
	}
	return value
}

fn bench_try_post_drain(rounds int, jobs int) !int {
	mut total := 0
	for _ in 0 .. rounds {
		mut ex := executor.new(queue_size: jobs)!
		done := chan int{cap: jobs}
		for _ in 0 .. jobs {
			ex.try_post(fn [done] () ! {
				done <- 1
			})!
		}
		drained := ex.drain_pending(jobs)!
		total += drained
		for _ in 0 .. jobs {
			total += <-done
		}
		ex.stop()
		ran_after_stop := ex.run_one()!
		if ran_after_stop {
			return error('try_post drain benchmark ran an unexpected job after stop')
		}
		ex.wait()!
	}
	return total
}

fn bench_run_sync(rounds int) !int {
	mut ex := executor.new(queue_size: rounds)!
	run_result := chan string{cap: 1}
	runner := spawn fn [mut ex, run_result] () {
		ex.run() or {
			run_result <- err.msg()
			return
		}
		run_result <- 'ok'
	}()

	done := chan int{cap: rounds}
	for _ in 0 .. rounds {
		ex.run_sync(fn [done] () ! {
			done <- 1
		})!
	}

	mut total := 0
	for _ in 0 .. rounds {
		total += <-done
	}
	ex.stop()
	run_msg := <-run_result
	if run_msg != 'ok' {
		runner.wait()
		return error('run_sync benchmark owner loop returned an error')
	}
	ex.wait()!
	runner.wait()
	return total
}

fn bench_bounded_timeout(rounds int, timeout_ms int) !int {
	timeout := time.Duration(timeout_ms) * time.millisecond
	mut total := 0
	for _ in 0 .. rounds {
		mut ex := executor.new(queue_size: 1)!
		ex.try_post(fn () ! {})!
		ex.post_with_timeout(timeout, fn () ! {}) or {
			if err.msg() != 'executor: timeout' {
				return err
			}
			total++
		}
		ran := ex.run_one()!
		if !ran {
			return error('bounded timeout benchmark did not run the accepted job')
		}
		ex.stop()
		ran_after_stop := ex.run_one()!
		if ran_after_stop {
			return error('bounded timeout benchmark ran an unexpected job after stop')
		}
		ex.wait()!
	}
	return total
}

fn bench_many_producers(producers int) !int {
	mut ex := executor.new(queue_size: producers)!
	accepted := chan bool{cap: producers}
	ran := chan int{cap: producers}
	mut threads := []thread{}

	for i in 0 .. producers {
		threads << spawn fn [mut ex, accepted, ran, i] () {
			ex.post_with_context(context.background(), fn [ran, i] () ! {
				value := i + 1
				ran <- value
				return
			}) or {
				accepted <- false
				return
			}
			accepted <- true
		}()
	}
	for _ in 0 .. producers {
		accepted_ok := <-accepted
		if !accepted_ok {
			return error('producer admission failed')
		}
	}
	drained := ex.drain_pending(producers)!
	if drained != producers {
		return error('multi-producer benchmark drained ${drained} jobs instead of ${producers}')
	}

	mut total := 0
	for _ in 0 .. producers {
		total += <-ran
	}
	for producer_thread in threads {
		producer_thread.wait()
	}
	ex.stop()
	ran_after_stop := ex.run_one()!
	if ran_after_stop {
		return error('multi-producer benchmark ran an unexpected job after stop')
	}
	ex.wait()!
	return total
}