package chitin

import (
	"context"
	"sync"
	"time"
)

// mergeContexts returns a Context value that combines properties of the two
// provided Contexts. Its deadline is the earliest deadline of its arguments.
// When its Value method is called, it will return values from the first
// Context in preference to values from the second context. When either
// context is canceled, its Done channel will be closed and calls to its Err
// method will yield a non-nil result corresponding to the Err value of the
// first context to be canceled or timed out.
//
// Callers must be careful of leaking goroutines. There is no danger of a leak
// here if at least one of the provided Contexts is the Background or TODO
// value. Otherwise, the caller must ensure that at least one of the provided
// contexts will be canceled in a timely manner.
func mergeContexts(a, b context.Context) context.Context {
	if a == b {
		return a
	} else if a == context.Background() || a == context.TODO() {
		return b
	} else if b == context.Background() || b == context.TODO() {
		return a
	}

	mc := &mergedContext{
		ctxA: a,
		ctxB: b,
		done: make(chan struct{}),
	}

	for _, ctx := range []context.Context{a, b} {
		if mc.err == nil {
			mc.err = ctx.Err()
		}

		t, ok := ctx.Deadline()
		if ok && (mc.deadline.IsZero() || t.Before(mc.deadline)) {
			mc.deadline = t
		}
	}

	go func() {
		select {
		case <-a.Done():
			mc.errMu.Lock()
			mc.err = a.Err()
		case <-b.Done():
			mc.errMu.Lock()
			mc.err = b.Err()
		}
		mc.errMu.Unlock()
		close(mc.done)
	}()

	return mc
}

type mergedContext struct {
	ctxA     context.Context
	ctxB     context.Context
	deadline time.Time
	done     chan struct{}

	errMu sync.Mutex
	err   error
}

var _ context.Context = (*mergedContext)(nil)

func (mc *mergedContext) Deadline() (time.Time, bool) {
	return mc.deadline, !mc.deadline.IsZero()

}

func (mc *mergedContext) Done() <-chan struct{} { return mc.done }

func (mc *mergedContext) Err() error {
	mc.errMu.Lock()
	err := mc.err
	mc.errMu.Unlock()
	return err
}

func (mc *mergedContext) Value(key interface{}) interface{} {
	for _, ctx := range []context.Context{mc.ctxA, mc.ctxB} {
		v := ctx.Value(key)
		if v != nil {
			return v
		}
	}
	return nil
}
