import benchmark
import context
import os
import time
import x.async as xasync

const default_group_rounds = 8
const default_group_jobs = 16
const default_task_rounds = 32
const default_pool_jobs = 64
const default_pool_workers = 4
const default_pool_bounded_rounds = 8
const default_pool_bounded_timeout_ms = 2
const default_timeout_rounds = 32
const default_every_iterations = 5
const default_every_interval_ms = 1

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

fn run_benchmarks() ! {
	group_rounds := env_int('XASYNC_BENCH_GROUP_ROUNDS', default_group_rounds, 1, 200)
	group_jobs := env_int('XASYNC_BENCH_GROUP_JOBS', default_group_jobs, 1, 512)
	task_rounds := env_int('XASYNC_BENCH_TASK_ROUNDS', default_task_rounds, 1, 500)
	pool_jobs := env_int('XASYNC_BENCH_POOL_JOBS', default_pool_jobs, 1, 1000)
	pool_workers := env_int('XASYNC_BENCH_POOL_WORKERS', default_pool_workers, 1, 64)
	pool_bounded_rounds := env_int('XASYNC_BENCH_POOL_BOUNDED_ROUNDS', default_pool_bounded_rounds,
		1, 100)
	pool_bounded_timeout_ms := env_int('XASYNC_BENCH_POOL_BOUNDED_TIMEOUT_MS',
		default_pool_bounded_timeout_ms, 1, 100)
	timeout_rounds := env_int('XASYNC_BENCH_TIMEOUT_ROUNDS', default_timeout_rounds, 1, 500)
	every_iterations := env_int('XASYNC_BENCH_EVERY_ITERATIONS', default_every_iterations, 1, 100)
	every_interval_ms :=
		env_int('XASYNC_BENCH_EVERY_INTERVAL_MS', default_every_interval_ms, 1, 100)

	println('x.async cautious benchmark')
	println('Override sizes with XASYNC_BENCH_* environment variables.')

	mut checksum := 0
	mut b := benchmark.start()
	checksum += bench_group(group_rounds, group_jobs)!
	b.measure('Group: rounds=${group_rounds}, jobs_per_round=${group_jobs}, checksum=${checksum}')

	checksum += bench_task(task_rounds)!
	b.measure('Task: rounds=${task_rounds}, checksum=${checksum}')

	checksum += bench_pool(pool_jobs, pool_workers)!
	b.measure('Pool: jobs=${pool_jobs}, workers=${pool_workers}, checksum=${checksum}')

	checksum += bench_pool_bounded_admission(pool_bounded_rounds, pool_bounded_timeout_ms)!
	b.measure('Pool bounded admission: rounds=${pool_bounded_rounds}, timeout_ms=${pool_bounded_timeout_ms}, checksum=${checksum}')

	checksum += bench_timeout(timeout_rounds)!
	b.measure('with_timeout: rounds=${timeout_rounds}, checksum=${checksum}')

	checksum += bench_every(every_iterations, every_interval_ms)!
	b.measure('every: iterations=${every_iterations}, interval_ms=${every_interval_ms}, 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_group(rounds int, jobs_per_round int) !int {
	mut total := 0
	for _ in 0 .. rounds {
		done := chan int{cap: jobs_per_round}
		mut group := xasync.new_group(context.background())
		for _ in 0 .. jobs_per_round {
			group.go(fn [done] (mut ctx context.Context) ! {
				_ = ctx
				done <- 1
			})!
		}
		group.wait()!
		for _ in 0 .. jobs_per_round {
			total += <-done
		}
	}
	return total
}

fn bench_task(rounds int) !int {
	mut total := 0
	for i in 0 .. rounds {
		mut task := xasync.run[int](fn [i] (mut ctx context.Context) !int {
			_ = ctx
			return i + 1
		})!
		total += task.wait()!
	}
	return total
}

fn bench_pool(jobs int, workers int) !int {
	done := chan int{cap: jobs}
	mut pool := xasync.new_pool(workers: workers, queue_size: jobs)!
	for _ in 0 .. jobs {
		pool.try_submit(fn [done] (mut ctx context.Context) ! {
			_ = ctx
			done <- 1
		})!
	}
	pool.close()!
	mut total := 0
	for _ in 0 .. jobs {
		total += <-done
	}
	return total
}

fn bench_pool_bounded_admission(rounds int, timeout_ms int) !int {
	mut total := 0
	admission_timeout := time.Duration(timeout_ms) * time.millisecond
	for _ in 0 .. rounds {
		mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
		started := chan bool{cap: 2}
		release := chan bool{cap: 3}
		finished := chan bool{cap: 2}
		attempting := chan bool{cap: 1}
		accepted := chan bool{cap: 1}
		ran := chan int{cap: 1}
		blocking_job := fn [started, release, finished] (mut ctx context.Context) ! {
			_ = ctx
			started <- true
			_ := <-release
			finished <- true
		}

		pool.try_submit(blocking_job)!
		wait_for_benchmark_bool(started, 'bounded admission pool job did not start')!
		pool.try_submit(blocking_job)!

		mut timed_out := false
		pool.submit_with_timeout(admission_timeout, fn (mut ctx context.Context) ! {
			_ = ctx
		}) or {
			if err.msg() != 'async: timeout' {
				return err
			}
			timed_out = true
		}
		if !timed_out {
			release <- true
			release <- true
			pool.close()!
			return error('submit_with_timeout accepted while the pool was full')
		}

		submit_thread := spawn fn [mut pool, attempting, accepted, ran] () {
			attempting <- true
			pool.submit_with_context(context.background(), fn [ran] (mut ctx context.Context) ! {
				_ = ctx
				ran <- 1
			}) or {
				accepted <- false
				return
			}
			accepted <- true
		}()
		wait_for_benchmark_bool(attempting, 'bounded admission submitter did not start')!

		release <- true
		wait_for_benchmark_bool(finished, 'bounded admission first job did not finish')!
		wait_for_benchmark_bool(accepted,
			'submit_with_context did not accept after capacity opened')!
		release <- true
		total += read_benchmark_int(ran, 'bounded admission accepted job did not run')!
		pool.close()!
		submit_thread.wait()
	}
	return total
}

fn bench_timeout(rounds int) !int {
	mut total := 0
	for _ in 0 .. rounds {
		xasync.with_timeout(1 * time.second, fn (mut ctx context.Context) ! {
			_ = ctx
		})!
		total++
	}
	return total
}

fn bench_every(iterations int, interval_ms int) !int {
	ctx, cancel := xasync.with_cancel()
	ticks := chan int{cap: iterations + 4}
	result := chan string{cap: 1}
	interval := time.Duration(interval_ms) * time.millisecond
	worker := spawn fn [ctx, ticks, result, interval] () {
		xasync.every(ctx, interval, fn [ticks] (mut ctx context.Context) ! {
			_ = ctx
			ticks <- 1
		}) or {
			result <- err.msg()
			return
		}
		result <- 'ok'
	}()

	mut total := 0
	for _ in 0 .. iterations {
		select {
			value := <-ticks {
				total += value
			}
			1 * time.second {
				cancel()
				worker.wait()
				return error('every benchmark did not receive a tick')
			}
		}
	}
	cancel()
	select {
		msg := <-result {
			if msg != 'context canceled' {
				worker.wait()
				return error('unexpected every benchmark result: ${msg}')
			}
		}
		1 * time.second {
			worker.wait()
			return error('every benchmark did not stop after cancellation')
		}
	}
	worker.wait()
	return total
}

fn wait_for_benchmark_bool(signal chan bool, message string) ! {
	select {
		ok := <-signal {
			if !ok {
				return error(message)
			}
		}
		1 * time.second {
			return error(message)
		}
	}
}

fn read_benchmark_int(signal chan int, message string) !int {
	select {
		value := <-signal {
			return value
		}
		1 * time.second {
			return error(message)
		}
	}
	return 0
}