import context
import sync
import time
import x.async as xasync

@[heap]
struct ControlledDeadlineParent {
	deadline_at time.Time
mut:
	done_ch   chan int
	mutex     &sync.Mutex = sync.new_mutex()
	err_value IError      = none
}

fn new_controlled_deadline_parent(deadline_at time.Time) &ControlledDeadlineParent {
	return &ControlledDeadlineParent{
		deadline_at: deadline_at
		done_ch:     chan int{}
		mutex:       sync.new_mutex()
	}
}

fn (ctx &ControlledDeadlineParent) deadline() ?time.Time {
	return ctx.deadline_at
}

fn (ctx &ControlledDeadlineParent) value(key context.Key) ?context.Any {
	_ = key
	return none
}

fn (mut ctx ControlledDeadlineParent) done() chan int {
	ctx.mutex.lock()
	done := ctx.done_ch
	ctx.mutex.unlock()
	return done
}

fn (mut ctx ControlledDeadlineParent) err() IError {
	ctx.mutex.lock()
	err := ctx.err_value
	ctx.mutex.unlock()
	return err
}

fn (mut ctx ControlledDeadlineParent) close_with_error(err IError) {
	ctx.mutex.lock()
	ctx.err_value = err
	if !ctx.done_ch.closed {
		ctx.done_ch.close()
	}
	ctx.mutex.unlock()
}

fn context_error_or_closed_without_error(mut ctx context.Context) ! {
	err := ctx.err()
	if err !is none {
		return err
	}
	return error('context closed without error')
}

fn test_with_timeout_returns_timeout_error() {
	xasync.with_timeout(5 * time.millisecond, fn (mut ctx context.Context) ! {
		_ = ctx
		time.sleep(50 * time.millisecond)
	}) or {
		assert err.msg() == 'async: timeout'
		return
	}
	assert false
}

fn test_with_timeout_returns_without_waiting_for_ignored_cancellation() {
	stopwatch := time.new_stopwatch()
	xasync.with_timeout(10 * time.millisecond, fn (mut ctx context.Context) ! {
		_ = ctx
		time.sleep(250 * time.millisecond)
	}) or {
		assert err.msg() == 'async: timeout'
		assert stopwatch.elapsed() < 150 * time.millisecond
		return
	}
	assert false
}

fn test_with_timeout_zero_duration_expires_immediately() {
	xasync.with_timeout(0 * time.millisecond, fn (mut ctx context.Context) ! {
		done := ctx.done()
		select {
			_ := <-done {
				return context_error_or_closed_without_error(mut ctx)
			}
			1 * time.second {
				return
			}
		}
	}) or {
		assert err.msg() == 'async: timeout'
		return
	}
	assert false
}

fn test_with_timeout_negative_duration_is_already_expired() {
	xasync.with_timeout(-1 * time.millisecond, fn (mut ctx context.Context) ! {
		done := ctx.done()
		select {
			_ := <-done {
				return context_error_or_closed_without_error(mut ctx)
			}
			1 * time.second {
				return
			}
		}
	}) or {
		assert err.msg() == 'async: timeout'
		return
	}
	assert false
}

fn test_with_timeout_success_before_timeout() {
	seen := chan int{cap: 1}
	xasync.with_timeout(1 * time.second, fn [seen] (mut ctx context.Context) ! {
		_ = ctx
		seen <- 7
	})!
	assert <-seen == 7
}

fn test_with_timeout_returns_job_error() {
	xasync.with_timeout(1 * time.second, fn (mut ctx context.Context) ! {
		_ = ctx
		return error('job failed')
	}) or {
		assert err.msg() == 'job failed'
		return
	}
	assert false
}

fn test_with_timeout_preserves_job_error_matching_context_message_before_timeout() {
	xasync.with_timeout(1 * time.second, fn (mut ctx context.Context) ! {
		_ = ctx
		return error('context deadline exceeded')
	}) or {
		assert err.msg() == 'context deadline exceeded'
		return
	}
	assert false
}

fn test_with_timeout_refuses_nil_job() {
	nil_job := unsafe { xasync.JobFn(nil) }
	xasync.with_timeout(1 * time.second, nil_job) or {
		assert err.msg() == 'async: job function is nil'
		return
	}
	assert false
}

fn test_with_timeout_returns_timeout_for_cooperative_job() {
	mut timed_out := false
	xasync.with_timeout(50 * time.millisecond, fn (mut ctx context.Context) ! {
		done := ctx.done()
		select {
			_ := <-done {
				return context_error_or_closed_without_error(mut ctx)
			}
			1 * time.second {
				return error('cooperative timeout job was not canceled')
			}
		}
	}) or {
		assert err.msg() == 'async: timeout'
		timed_out = true
	}
	assert timed_out
}

fn test_with_timeout_context_uses_parent_context() {
	mut parent_ctx, cancel := xasync.with_cancel()
	cancel()
	xasync.with_timeout_context(parent_ctx, 1 * time.second, fn (mut ctx context.Context) ! {
		done := ctx.done()
		select {
			_ := <-done {
				return context_error_or_closed_without_error(mut ctx)
			}
			1 * time.second {}
		}
	}) or {
		assert err.msg() == 'context canceled'
		return
	}
	assert false
}

fn test_with_timeout_context_preserves_already_expired_parent_deadline() {
	mut background := context.background()
	parent_ctx, parent_cancel := context.with_timeout(mut background, -1 * time.millisecond)
	defer {
		parent_cancel()
	}

	xasync.with_timeout_context(parent_ctx, 1 * time.second, fn (mut ctx context.Context) ! {
		_ = ctx
		return error('job should not run with an already expired parent')
	}) or {
		assert err.msg() == 'context deadline exceeded'
		return
	}
	assert false
}

fn test_with_timeout_context_preserves_parent_deadline_that_expires_first() {
	mut background := context.background()
	parent_ctx, parent_cancel := context.with_timeout(mut background, 10 * time.millisecond)
	defer {
		parent_cancel()
	}

	xasync.with_timeout_context(parent_ctx, 1 * time.second, fn (mut ctx context.Context) ! {
		_ = ctx
		time.sleep(100 * time.millisecond)
	}) or {
		assert err.msg() == 'context deadline exceeded'
		return
	}
	assert false
}

fn test_with_timeout_context_waits_for_controlled_parent_done_and_returns_exact_parent_error() {
	mut parent_ctx := new_controlled_deadline_parent(time.now().add(-1 * time.second))
	started := chan bool{cap: 1}
	release := chan bool{cap: 1}
	result := chan string{cap: 1}

	caller := spawn fn [mut parent_ctx, started, release, result] () {
		xasync.with_timeout_context(context.Context(parent_ctx), 1 * time.second, fn [started, release] (mut ctx context.Context) ! {
			_ = ctx
			started <- true
			_ := <-release
		}) or {
			result <- err.msg()
			return
		}
		result <- 'ok'
	}()

	select {
		did_start := <-started {
			assert did_start
		}
		1 * time.second {
			assert false, 'with_timeout_context returned before starting the job for a pending parent'
		}
	}
	select {
		msg := <-result {
			assert false, 'with_timeout_context returned before parent done closed: ${msg}'
		}
		else {}
	}

	parent_ctx.close_with_error(error('controlled parent deadline'))
	select {
		msg := <-result {
			assert msg == 'controlled parent deadline'
		}
		1 * time.second {
			assert false, 'with_timeout_context did not return after controlled parent closed'
		}
	}
	release <- true
	caller.wait()
}

fn test_with_timeout_context_uses_own_timeout_before_parent_deadline() {
	mut background := context.background()
	parent_ctx, parent_cancel := context.with_timeout(mut background, 1 * time.second)
	defer {
		parent_cancel()
	}

	xasync.with_timeout_context(parent_ctx, 10 * time.millisecond, fn (mut ctx context.Context) ! {
		_ = ctx
		time.sleep(100 * time.millisecond)
	}) or {
		assert err.msg() == 'async: timeout'
		return
	}
	assert false
}