package trace

import (
	"encoding/binary"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"sync/atomic"
	"time"

	"crypto/rand"

	"golang.org/x/net/context"
)

var (
	headerID    = http.CanonicalHeaderKey("Trace-ID")
	headerSubID = http.CanonicalHeaderKey("Trace-Span")
)

type key int

var (
	keyTransaction = new(key)
)

type transaction struct {
	ID    uint64
	SubID string
	// accessed atomically
	spanCount int64
}

// Logfmt returns a text version of the Trace values stored in the provided
// Context. It can be used as a prefix for log lines (if the user adds
// intermediate whitespace).
func Logfmt(ctx context.Context) string {
	value, ok := ctx.Value(keyTransaction).(*transaction)
	if !ok {
		return fmt.Sprintf("trace-id=%18s trace-span=%q", "", "")
	}
	return fmt.Sprintf("trace-id=%#016x trace-span=%q", value.ID, value.SubID)
}

// ContextFromHeader issues a new Context based on parent, but including the
// Trace values in the Header. If the Trace values are missing, a new
// transaction id (with high entropy) is issued.
func ContextFromHeader(parent context.Context, h http.Header) context.Context {
	id, err := strconv.ParseUint(h.Get(headerID), 0, 64)
	if err != nil {
		return newTransaction(parent)
	}
	subid := h.Get(headerSubID)
validateSpan:
	for _, c := range subid {
		switch {
		default:
			subid = "invalid"
			break validateSpan
		case c == '.':
		case '0' <= c && c <= '9':
		case 'a' <= c && c <= 'z':
		case 'A' <= c && c <= 'Z':
		}
	}

	value := &transaction{
		ID:    id,
		SubID: subid,
	}
	return context.WithValue(parent, keyTransaction, value)
}

// NewSpan 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 NewSpan(parent context.Context) context.Context {
	parentTx, ok := parent.Value(keyTransaction).(*transaction)
	if !ok {
		// We've been passed a Context without an attached Trace identity.
		// This may be a bug (so we could log the event). Now is not an
		// appropriate time to attach a new Trace identity, so we'll just
		// return the original Context.
		return parent
	}

	spanCount := atomic.AddInt64(&parentTx.spanCount, 1)
	subTx := &transaction{
		ID:    parentTx.ID,
		SubID: fmt.Sprintf("%s.%d", parentTx.SubID, spanCount-1),
	}

	return context.WithValue(parent, keyTransaction, subTx)
}

// AugmentHeader adds Trace transaction information to a Header. It does not
// modify the internal Trace record state.
//
// When making subcalls on behalf of an incoming request, a user must generate
// a new Trace span (via the NewSpan function) and then
func AugmentHeader(ctx context.Context, h http.Header) {
	value, ok := ctx.Value(keyTransaction).(*transaction)
	if !ok {
		h.Del(headerID)
		h.Del(headerSubID)
		return
	}
	h.Set(headerID, fmt.Sprintf("%#016x", value.ID))
	h.Set(headerSubID, value.SubID)
}

func newTransaction(parent context.Context) context.Context {
	var id uint64
	select {
	default:
		// Oops, we've got to wait for more entropy. Record the underflow ...
		// but of course there are probably a bunch of other starving
		// consumers.
		select {
		default:
		case underflow <- struct{}{}:
		}
		// Be sure to get what we came for. Yes, it will probably involve
		// waiting this time.
		id = <-ids
	case id = <-ids:
	}

	value := &transaction{
		ID:    id,
		SubID: "",
	}
	return context.WithValue(parent, keyTransaction, value)
}

const (
	idQueueSize = 1 << 10
)

var (
	ids       <-chan uint64
	underflow chan<- struct{}
)

func init() {
	c := make(chan uint64, idQueueSize)
	ids = c
	go genIds(c)

	u := make(chan struct{}, 1)
	underflow = u
	go logUnderflow(u, 1*time.Second)
}

func genIds(ids chan<- uint64) {
	var i uint64
	for {
		err := binary.Read(rand.Reader, binary.BigEndian, &i)
		if err != nil {
			log.Fatalf("entropy-error=%q", err)
		}

		ids <- i
	}
}

func logUnderflow(underflow <-chan struct{}, period time.Duration) {
	t := time.NewTicker(period)
	defer t.Stop()

	for _ = range t.C {
		select {
		default:
		case <-underflow:
			log.Printf("entropy-underflow")
		}
	}
}
