// +build go1.7

package chitin

import (
	"context"
	"fmt"
	"net/http"
	"sync"
	"time"

	"code.justin.tv/common/chitin/internal/trace"
)

func fetchContext(w http.ResponseWriter, r *http.Request) context.Context {
	ctx := r.Context()
	if trace.HasTraceValue(ctx) {
		return ctx
	}

	err := fmt.Errorf("chitin:" +
		" the provided http.Request doesn't seem to have come" +
		" from a chitin-enabled http.Handler. See godoc for the" +
		" chitin.Context function for advice on its usage.")
	panic(err)
}

func requestWithContext(ctx context.Context, src *http.Request) (*http.Request, context.Context) {
	child := mergeContexts(src.Context(), ctx)

	dst := new(http.Request).WithContext(child)

	return dst, child
}

func contextOfRequest(req *http.Request) context.Context {
	return req.Context()
}

// 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
}
