module onecontext

import context
import sync
import time

// canceled is the error returned when the cancel function is called on a merged context
pub const canceled = error('canceled context')

@[heap]
struct OneContext {
mut:
	ctx        context.Context
	ctxs       []context.Context
	done       chan int
	err        IError           = none
	err_mutex  &sync.Mutex      = sync.new_mutex()
	cancel_fn  context.CancelFn = unsafe { nil }
	cancel_ctx context.Context
}

// merge allows to merge multiple contexts
// it returns the merged context
pub fn merge(ctx context.Context, ctxs ...context.Context) (context.Context, context.CancelFn) {
	mut background := context.background()
	cancel_ctx, cancel := context.with_cancel(mut &background)
	mut octx := &OneContext{
		done:       chan int{cap: 3}
		ctx:        ctx
		ctxs:       ctxs
		cancel_fn:  cancel
		cancel_ctx: cancel_ctx
	}
	spawn octx.run()
	return context.Context(octx), context.CancelFn(cancel)
}

// deadline returns the earliest deadline among all merged contexts,
// or none if no context has a deadline set.
pub fn (octx OneContext) deadline() ?time.Time {
	mut min := time.Time{}
	mut found := false

	if deadline := octx.ctx.deadline() {
		min = deadline
		found = true
	}

	for ctx in octx.ctxs {
		if deadline := ctx.deadline() {
			if !found || deadline < min {
				min = deadline
			}
			found = true
		}
	}

	if !found {
		return none
	}

	return min
}

// done returns the done channel, which is closed when the merged context is canceled.
pub fn (octx OneContext) done() chan int {
	return octx.done
}

// err returns the error from the merged context, or `none` if not yet canceled.
pub fn (mut octx OneContext) err() IError {
	octx.err_mutex.lock()
	defer {
		octx.err_mutex.unlock()
	}
	return octx.err
}

// value looks up a value by key across all merged contexts, returning
// the first match found or none if no context holds the key.
pub fn (octx OneContext) value(key context.Key) ?context.Any {
	if value := octx.ctx.value(key) {
		return value
	}

	for ctx in octx.ctxs {
		if value := ctx.value(key) {
			return value
		}
	}

	return none
}

// run starts listening for cancellation signals from all merged contexts.
pub fn (mut octx OneContext) run() {
	mut wrapped_ctx := &octx.ctx
	if octx.ctxs.len == 1 {
		mut first_ctx := &octx.ctxs[0]
		octx.run_two_contexts(mut wrapped_ctx, mut first_ctx)
		return
	}

	octx.run_multiple_contexts(mut wrapped_ctx)
	for mut ctx in octx.ctxs {
		octx.run_multiple_contexts(mut ctx)
	}
}

// str returns a string representation of the OneContext.
pub fn (octx OneContext) str() string {
	return ''
}

// cancel cancels the merged context with the given error, closing
// the done channel and propagating cancellation to the underlying context.
pub fn (mut octx OneContext) cancel(err IError) {
	octx.cancel_fn()
	octx.err_mutex.lock()
	octx.err = err
	octx.err_mutex.unlock()
	if !octx.done.closed {
		octx.done <- 0
		octx.done.close()
	}
}

// run_two_contexts spawns a listener that cancels the merged context
// when either of the two given contexts is done.
pub fn (mut octx OneContext) run_two_contexts(mut ctx1 context.Context, mut ctx2 context.Context) {
	octx_cancel_done := octx.cancel_ctx.done()
	c1done := ctx1.done()
	c2done := ctx2.done()
	spawn fn (mut octx OneContext, octx_cancel_done chan int, c1done chan int, c2done chan int, mut ctx1 context.Context, mut ctx2 context.Context) {
		select {
			_ := <-octx_cancel_done {
				octx.cancel(canceled)
			}
			_ := <-c1done {
				octx.cancel(ctx1.err())
			}
			_ := <-c2done {
				octx.cancel(ctx2.err())
			}
		}
	}(mut &octx, octx_cancel_done, c1done, c2done, mut ctx1, mut ctx2)
}

// run_multiple_contexts spawns a listener that cancels the merged context
// when the given context is done.
pub fn (mut octx OneContext) run_multiple_contexts(mut ctx context.Context) {
	octx_cancel_done := octx.cancel_ctx.done()
	cdone := ctx.done()
	spawn fn (mut octx OneContext, octx_cancel_done chan int, cdone chan int, mut ctx context.Context) {
		select {
			_ := <-octx_cancel_done {
				octx.cancel(canceled)
			}
			_ := <-cdone {
				octx.cancel(ctx.err())
			}
		}
	}(mut &octx, octx_cancel_done, cdone, mut ctx)
}