import context
import time
import x.async as xasync

fn test_pool_rejects_invalid_config() {
	xasync.new_pool(workers: 0, queue_size: 1) or {
		assert err.msg() == 'async: pool worker count must be positive'
		return
	}
	assert false
}

fn test_pool_rejects_invalid_queue_size() {
	xasync.new_pool(workers: 1, queue_size: 0) or {
		assert err.msg() == 'async: pool queue size must be positive'
		return
	}
	assert false
}

fn test_pool_respects_worker_count() {
	mut pool := xasync.new_pool(workers: 2, queue_size: 4)!
	started := chan bool{cap: 4}
	release := chan bool{cap: 4}
	for _ in 0 .. 4 {
		pool.try_submit(fn [started, release] (mut ctx context.Context) ! {
			_ = ctx
			started <- true
			_ := <-release
		})!
	}

	wait_for_bool_signal(started, 'first pool job did not start')
	wait_for_bool_signal(started, 'second pool job did not start')
	select {
		_ := <-started {
			assert false, 'pool started more jobs than worker count before release'
		}
		100 * time.millisecond {}
	}

	for _ in 0 .. 4 {
		release <- true
	}
	pool.close()!
}

fn test_pool_queue_full_returns_backpressure_error() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
	started := chan bool{cap: 1}
	release := chan bool{cap: 2}
	blocking_job := fn [started, release] (mut ctx context.Context) ! {
		_ = ctx
		started <- true
		_ := <-release
	}

	pool.try_submit(blocking_job)!
	wait_for_bool_signal(started, 'blocking pool job did not start')
	pool.try_submit(blocking_job)!
	pool.try_submit(blocking_job) or {
		assert err.msg() == 'async: pool queue is full'
		release <- true
		release <- true
		pool.close()!
		return
	}
	assert false
}

fn test_pool_submit_with_context_waits_until_capacity_is_available() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
	started := chan bool{cap: 2}
	release := chan bool{cap: 2}
	attempting := chan bool{cap: 1}
	accepted := chan bool{cap: 1}
	ran := chan bool{cap: 1}
	blocking_job := fn [started, release] (mut ctx context.Context) ! {
		_ = ctx
		started <- true
		_ := <-release
	}

	pool.try_submit(blocking_job)!
	wait_for_bool_signal(started, 'first blocking pool job did not start')
	pool.try_submit(blocking_job)!

	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 <- true
		}) or {
			accepted <- false
			return
		}
		accepted <- true
	}()
	wait_for_bool_signal(attempting, 'bounded submit thread did not start')
	select {
		_ := <-accepted {
			assert false, 'bounded submit returned while pool was full'
		}
		50 * time.millisecond {}
	}

	release <- true
	wait_for_bool_signal(accepted, 'bounded submit did not accept after capacity opened')
	release <- true
	wait_for_bool_signal(ran, 'accepted bounded submit job did not run')
	pool.close()!
	submit_thread.wait()
}

fn test_pool_submit_with_timeout_returns_timeout_without_leaking_acceptance() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
	started := chan bool{cap: 2}
	release := chan bool{cap: 2}
	finished := chan bool{cap: 2}
	blocking_job := fn [started, release, finished] (mut ctx context.Context) ! {
		_ = ctx
		started <- true
		_ := <-release
		finished <- true
	}

	pool.try_submit(blocking_job)!
	wait_for_bool_signal(started, 'first blocking pool job did not start')
	pool.try_submit(blocking_job)!

	pool.submit_with_timeout(20 * time.millisecond, fn (mut ctx context.Context) ! {
		_ = ctx
	}) or {
		assert err.msg() == 'async: timeout'
		release <- true
		wait_for_bool_signal(finished, 'first blocking pool job did not finish')
		pool.submit_with_timeout(1 * time.second, fn (mut ctx context.Context) ! {
			_ = ctx
		})!
		release <- true
		pool.close()!
		return
	}
	assert false
}

fn test_pool_submit_with_context_parent_cancel_does_not_accept_job() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
	started := chan bool{cap: 1}
	release := chan bool{cap: 2}
	attempting := chan bool{cap: 1}
	would_run := chan bool{cap: 1}
	result := chan string{cap: 1}
	parent_ctx, cancel := xasync.with_cancel()
	blocking_job := fn [started, release] (mut ctx context.Context) ! {
		_ = ctx
		started <- true
		_ := <-release
	}

	pool.try_submit(blocking_job)!
	wait_for_bool_signal(started, 'blocking pool job did not start')
	pool.try_submit(blocking_job)!

	submit_thread := spawn fn [mut pool, parent_ctx, attempting, would_run, result] () {
		attempting <- true
		pool.submit_with_context(parent_ctx, fn [would_run] (mut ctx context.Context) ! {
			_ = ctx
			would_run <- true
		}) or {
			result <- err.msg()
			return
		}
		result <- 'accepted'
	}()
	wait_for_bool_signal(attempting, 'bounded submit thread did not start')
	select {
		msg := <-result {
			assert false, 'bounded submit returned before parent cancellation: ${msg}'
		}
		50 * time.millisecond {}
	}

	cancel()
	select {
		msg := <-result {
			assert msg == 'context canceled'
		}
		1 * time.second {
			assert false, 'bounded submit did not return after parent cancellation'
		}
	}
	release <- true
	release <- true
	pool.close()!
	assert_no_bool_signal(would_run, 'job was accepted after parent cancellation')
	submit_thread.wait()
}

fn test_pool_close_wakes_waiting_submitter_without_accepting_job() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
	started := chan bool{cap: 1}
	release := chan bool{cap: 2}
	attempting := chan bool{cap: 1}
	would_run := chan bool{cap: 1}
	result := chan string{cap: 1}
	closed := chan bool{cap: 1}
	blocking_job := fn [started, release] (mut ctx context.Context) ! {
		_ = ctx
		started <- true
		_ := <-release
	}

	pool.try_submit(blocking_job)!
	wait_for_bool_signal(started, 'blocking pool job did not start')
	pool.try_submit(blocking_job)!

	submit_thread := spawn fn [mut pool, attempting, would_run, result] () {
		attempting <- true
		pool.submit_with_context(context.background(), fn [would_run] (mut ctx context.Context) ! {
			_ = ctx
			would_run <- true
		}) or {
			result <- err.msg()
			return
		}
		result <- 'accepted'
	}()
	wait_for_bool_signal(attempting, 'bounded submit thread did not start')
	select {
		msg := <-result {
			assert false, 'bounded submit returned before close: ${msg}'
		}
		50 * time.millisecond {}
	}

	close_thread := spawn fn [mut pool, closed] () {
		pool.close() or {
			closed <- false
			return
		}
		closed <- true
	}()
	select {
		msg := <-result {
			assert msg == 'async: pool is closed'
		}
		1 * time.second {
			assert false, 'bounded submit did not return after pool close started'
		}
	}
	release <- true
	release <- true
	wait_for_bool_signal(closed, 'pool close did not finish after accepted jobs were released')
	assert_no_bool_signal(would_run, 'job was accepted after pool close started')
	submit_thread.wait()
	close_thread.wait()
}

fn test_pool_bounded_submit_rejects_nil_job() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
	nil_job := unsafe { xasync.JobFn(nil) }
	pool.submit_with_context(context.background(), nil_job) or {
		assert err.msg() == 'async: job function is nil'
		pool.submit_with_timeout(20 * time.millisecond, nil_job) or {
			assert err.msg() == 'async: job function is nil'
			pool.close()!
			return
		}
		assert false
	}
	assert false
}

fn test_pool_submit_after_close_is_refused() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
	pool.close()!
	pool.try_submit(fn (mut ctx context.Context) ! {
		_ = ctx
	}) or {
		assert err.msg() == 'async: pool is closed'
		return
	}
	assert false
}

fn test_pool_wait_is_one_shot() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
	pool.wait()!
	pool.wait() or {
		assert err.msg() == 'async: pool wait was already called'
		return
	}
	assert false
}

fn test_pool_refuses_nil_job() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
	nil_job := unsafe { xasync.JobFn(nil) }
	pool.try_submit(nil_job) or {
		assert err.msg() == 'async: job function is nil'
		pool.close()!
		return
	}
	assert false
}

fn test_pool_close_waits_for_accepted_jobs() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
	release := chan bool{cap: 1}
	closed := chan bool{cap: 1}
	pool.try_submit(fn [release] (mut ctx context.Context) ! {
		_ = ctx
		_ := <-release
	})!

	close_thread := spawn fn [mut pool, closed] () {
		pool.close() or {
			closed <- false
			return
		}
		closed <- true
	}()
	select {
		_ := <-closed {
			assert false, 'pool close returned before accepted job completed'
		}
		100 * time.millisecond {}
	}
	release <- true
	select {
		ok := <-closed {
			assert ok
		}
		1 * time.second {
			assert false, 'pool close did not return after accepted job completed'
		}
	}
	close_thread.wait()
}

fn test_pool_first_error_is_propagated() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 2)!
	started := chan bool{cap: 1}
	release := chan bool{cap: 1}
	pool.try_submit(fn [started, release] (mut ctx context.Context) ! {
		_ = ctx
		started <- true
		_ := <-release
		return error('first pool failure')
	})!
	wait_for_bool_signal(started, 'first pool failure job did not start')
	pool.try_submit(fn (mut ctx context.Context) ! {
		_ = ctx
		return error('second pool failure')
	})!
	release <- true
	pool.close() or {
		assert err.msg() == 'first pool failure'
		return
	}
	assert false
}

fn test_pool_error_does_not_drop_accepted_jobs() {
	mut pool := xasync.new_pool(workers: 1, queue_size: 1)!
	accepted_job_ran := chan bool{cap: 1}
	pool.try_submit(fn (mut ctx context.Context) ! {
		_ = ctx
		return error('first pool failure')
	})!
	pool.try_submit(fn [accepted_job_ran] (mut ctx context.Context) ! {
		_ = ctx
		accepted_job_ran <- true
	})!
	pool.close() or {
		assert err.msg() == 'first pool failure'
		select {
			did_run := <-accepted_job_ran {
				assert did_run
			}
			1 * time.second {
				assert false, 'accepted pool job did not run after earlier error'
			}
		}
		return
	}
	assert false
}

fn test_pool_concurrent_errors_return_one_error_and_drain_accepted_jobs() {
	error_jobs := 4
	ok_jobs := 8
	mut pool := xasync.new_pool(workers: error_jobs, queue_size: ok_jobs)!
	started := chan bool{cap: error_jobs}
	release := chan bool{cap: error_jobs}
	completed_ok := chan bool{cap: ok_jobs}
	for i in 0 .. error_jobs {
		pool.try_submit(fn [started, release, i] (mut ctx context.Context) ! {
			_ = ctx
			started <- true
			_ := <-release
			return error('pool concurrent failure ${i}')
		})!
	}
	for _ in 0 .. ok_jobs {
		pool.try_submit(fn [completed_ok] (mut ctx context.Context) ! {
			_ = ctx
			completed_ok <- true
		})!
	}
	for _ in 0 .. error_jobs {
		wait_for_bool_signal(started, 'pool error job did not start')
	}
	for _ in 0 .. error_jobs {
		release <- true
	}
	pool.close() or {
		assert err.msg().starts_with('pool concurrent failure ')
		for _ in 0 .. ok_jobs {
			wait_for_bool_signal(completed_ok, 'accepted ok pool job did not complete')
		}
		return
	}
	assert false
}

fn test_pool_close_drains_many_accepted_jobs_while_finishing() {
	jobs := 12
	workers := 3
	mut pool := xasync.new_pool(workers: workers, queue_size: jobs - workers)!
	started := chan bool{cap: jobs}
	release := chan bool{cap: jobs}
	finished := chan bool{cap: jobs}
	closed := chan bool{cap: 1}
	for _ in 0 .. jobs {
		pool.try_submit(fn [started, release, finished] (mut ctx context.Context) ! {
			_ = ctx
			started <- true
			_ := <-release
			finished <- true
		})!
	}
	for _ in 0 .. workers {
		wait_for_bool_signal(started, 'initial pool job did not start')
	}
	close_thread := spawn fn [mut pool, closed] () {
		pool.close() or {
			closed <- false
			return
		}
		closed <- true
	}()
	wait_until_pool_rejects_as_closed(mut pool)
	assert_no_bool_signal(closed, 'pool close returned while accepted jobs were still blocked')

	for _ in 0 .. jobs - 1 {
		release <- true
	}
	for _ in 0 .. jobs - 1 {
		wait_for_bool_signal(finished, 'accepted pool job did not finish')
	}
	assert_no_bool_signal(closed, 'pool close returned before the last accepted job finished')

	release <- true
	wait_for_bool_signal(finished, 'last accepted pool job did not finish')
	wait_for_bool_signal(closed, 'pool close did not drain accepted jobs')
	close_thread.wait()
}

fn test_pool_parent_cancellation_is_observed_by_cooperative_job() {
	parent_ctx, cancel := xasync.with_cancel()
	mut pool := xasync.new_pool_with_context(parent_ctx, workers: 1, queue_size: 1)!
	started := chan bool{cap: 1}
	pool.try_submit(fn [started] (mut ctx context.Context) ! {
		started <- true
		done := ctx.done()
		select {
			_ := <-done {
				return ctx.err()
			}
			1 * time.second {
				return error('pool job did not observe parent cancellation')
			}
		}
		return error('unreachable')
	})!
	wait_for_bool_signal(started, 'pool job did not start before parent cancellation')
	cancel()
	pool.close() or {
		assert err.msg() == 'context canceled'
		return
	}
	assert false
}

fn test_pool_non_cooperative_job_finishes_naturally_after_parent_cancel() {
	parent_ctx, cancel := xasync.with_cancel()
	mut pool := xasync.new_pool_with_context(parent_ctx, workers: 1, queue_size: 1)!
	finished := chan bool{cap: 1}
	pool.try_submit(fn [finished] (mut ctx context.Context) ! {
		_ = ctx
		time.sleep(20 * time.millisecond)
		finished <- true
	})!
	cancel()
	pool.close()!
	assert <-finished
}

fn test_pool_short_stress_many_jobs() {
	jobs := 100
	mut pool := xasync.new_pool(workers: 4, queue_size: jobs)!
	done := chan int{cap: jobs}
	for i in 0 .. jobs {
		pool.try_submit(fn [done, i] (mut ctx context.Context) ! {
			_ = ctx
			done <- i
		})!
	}
	pool.close()!

	mut seen := []bool{len: jobs}
	for _ in 0 .. jobs {
		i := <-done
		assert i >= 0
		assert i < jobs
		assert !seen[i]
		seen[i] = true
	}
	for was_seen in seen {
		assert was_seen
	}
}

fn test_pool_bounded_submit_stress_waiting_submitters_are_released() {
	workers := 2
	queue_size := 2
	initial_jobs := workers + queue_size
	submitters := 8
	mut pool := xasync.new_pool(workers: workers, queue_size: queue_size)!
	started := chan bool{cap: initial_jobs}
	release := chan bool{cap: initial_jobs}
	attempting := chan bool{cap: submitters}
	accepted := chan int{cap: submitters}
	ran := chan int{cap: submitters}
	blocking_job := fn [started, release] (mut ctx context.Context) ! {
		_ = ctx
		started <- true
		_ := <-release
	}

	for _ in 0 .. initial_jobs {
		pool.try_submit(blocking_job)!
	}
	for _ in 0 .. workers {
		wait_for_bool_signal(started, 'initial pool worker did not start')
	}

	mut submit_threads := []thread{}
	for i in 0 .. submitters {
		submit_threads << spawn fn [mut pool, attempting, accepted, ran, i] () {
			attempting <- true
			pool.submit_with_context(context.background(), fn [ran, i] (mut ctx context.Context) ! {
				_ = ctx
				ran <- i
			}) or {
				failed_index := -1 - i
				accepted <- failed_index
				return
			}
			accepted <- i
		}()
	}
	for _ in 0 .. submitters {
		wait_for_bool_signal(attempting, 'bounded submitter did not start')
	}

	for _ in 0 .. initial_jobs {
		release <- true
	}

	mut seen_accepted := []bool{len: submitters}
	for _ in 0 .. submitters {
		i := wait_for_int_signal(accepted, 'bounded submitter was not accepted')
		assert i >= 0
		assert i < submitters
		assert !seen_accepted[i]
		seen_accepted[i] = true
	}
	mut seen_ran := []bool{len: submitters}
	for _ in 0 .. submitters {
		i := wait_for_int_signal(ran, 'accepted bounded submitter job did not run')
		assert i >= 0
		assert i < submitters
		assert !seen_ran[i]
		seen_ran[i] = true
	}

	pool.close()!
	for t in submit_threads {
		t.wait()
	}
}

fn test_pool_bounded_submit_stress_waiting_submitters_rejected_on_close() {
	workers := 2
	queue_size := 2
	initial_jobs := workers + queue_size
	submitters := 8
	mut pool := xasync.new_pool(workers: workers, queue_size: queue_size)!
	started := chan bool{cap: initial_jobs}
	release := chan bool{cap: initial_jobs}
	attempting := chan bool{cap: submitters}
	result := chan string{cap: submitters}
	ran := chan bool{cap: submitters}
	closed := chan bool{cap: 1}
	blocking_job := fn [started, release] (mut ctx context.Context) ! {
		_ = ctx
		started <- true
		_ := <-release
	}

	for _ in 0 .. initial_jobs {
		pool.try_submit(blocking_job)!
	}
	for _ in 0 .. workers {
		wait_for_bool_signal(started, 'initial pool worker did not start')
	}

	mut submit_threads := []thread{}
	for i in 0 .. submitters {
		submit_threads << spawn fn [mut pool, attempting, result, ran, i] () {
			attempting <- true
			pool.submit_with_context(context.background(), fn [ran] (mut ctx context.Context) ! {
				_ = ctx
				ran <- true
			}) or {
				result <- err.msg()
				return
			}
			result <- 'accepted ${i}'
		}()
	}
	for _ in 0 .. submitters {
		wait_for_bool_signal(attempting, 'bounded submitter did not start')
	}

	close_thread := spawn fn [mut pool, closed] () {
		pool.close() or {
			closed <- false
			return
		}
		closed <- true
	}()
	for _ in 0 .. submitters {
		msg := wait_for_string_signal(result, 'bounded submitter did not return after close')
		assert msg == 'async: pool is closed'
	}
	assert_no_bool_signal(ran, 'bounded submitter job ran after close started')

	for _ in 0 .. initial_jobs {
		release <- true
	}
	wait_for_bool_signal(closed, 'pool close did not finish after releasing accepted jobs')
	for t in submit_threads {
		t.wait()
	}
	close_thread.wait()
}

fn wait_for_bool_signal(signal chan bool, message string) {
	select {
		ok := <-signal {
			assert ok
		}
		1 * time.second {
			assert false, message
		}
	}
}

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

fn wait_for_string_signal(signal chan string, message string) string {
	select {
		value := <-signal {
			return value
		}
		1 * time.second {
			assert false, message
		}
	}
	return ''
}

fn assert_no_bool_signal(signal chan bool, message string) {
	select {
		_ := <-signal {
			assert false, message
		}
		else {}
	}
}

fn wait_until_pool_rejects_as_closed(mut pool xasync.Pool) {
	probe := fn (mut ctx context.Context) ! {
		_ = ctx
	}
	for _ in 0 .. 100 {
		pool.try_submit(probe) or {
			if err.msg() == 'async: pool is closed' {
				return
			}
			assert err.msg() == 'async: pool queue is full'
			time.sleep(1 * time.millisecond)
			continue
		}
		assert false, 'pool accepted probe while backlog should be full'
	}
	assert false, 'pool close did not start'
}