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_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)
    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_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_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
}