package chitin

import (
	"time"

	"code.justin.tv/common/chitin/internal/trace"
	rt "code.justin.tv/release/trace"
	"golang.org/x/net/context"
)

type key int

var (
	keyArrivalTime = new(key)
	keyRecorder    = new(key)
)

func markArrivalTime(ctx context.Context, val time.Time) context.Context {
	return context.WithValue(ctx, keyArrivalTime, val)
}

func getArrivalTime(ctx context.Context) time.Time {
	return ctx.Value(keyArrivalTime).(time.Time)
}

func setRecorder(parent context.Context, val recorder) context.Context {
	return context.WithValue(parent, keyRecorder, val)
}

// GetTraceId returns the Trace ID of the current inter-process request tree
// and request, or nil if the provided context has no attached Trace
// annotations.
//
// The methods of a nil return value are safe to call, though any strings they
// return won't correspond to a real transaction.
func GetTraceID(ctx context.Context) *rt.ID {
	return trace.GetTraceID(ctx)
}

// ExperimentalNewSpan issues a new Trace span id, to be used for tracing
// calls to subsystems.
//
// Users making calls to external systems on behalf of an incoming request
// must pass a sub-context to the relevant RPC libraries. These child contexts
// should have their own span ids (from this function), and may also have
// separate cancellations or timeouts (via the x/net/context package).
func ExperimentalNewSpan(parent context.Context) context.Context {
	return trace.NewSpan(parent)
}
