package trace

import (
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"strings"
	"sync/atomic"

	"code.justin.tv/release/trace"
	"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 {
	TxID [2]uint64
	Path CallPath
	// accessed atomically
	spanCount uint32
}

// FromContext returns the Transaction bound to the context, if any.
func FromContext(ctx context.Context) (*Transaction, bool) {
	value, ok := ctx.Value(keyTransaction).(*Transaction)
	return value, ok
}

// NewContext returns a copy of the parent context and associates it with a
// Transaction.
func NewContext(ctx context.Context, tx *Transaction) context.Context {
	return context.WithValue(ctx, keyTransaction, tx)
}

// 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 {
	var id [2]uint64
	var path string

	value, ok := FromContext(ctx)
	if ok {
		id = value.TxID
		path = value.Path.String()
	}

	return fmt.Sprintf("trace-id=%s trace-span=%q", TXIDString(id), path)
}

// 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.
func GetTraceID(ctx context.Context) *trace.ID {
	value, ok := FromContext(ctx)
	if !ok {
		return nil
	}
	return trace.ParseID(TXIDString(value.TxID) + value.Path.String())
}

type CallPath []uint32

func (cp CallPath) String() string {
	var dst []byte
	for _, seg := range cp {
		dst = append(dst, '.')
		dst = strconv.AppendUint(dst, uint64(seg), 10)
	}
	return string(dst)
}

// HasTraceValue returns whether the provided context includes the Trace
// transaction ID of the current RPC tree.
func HasTraceValue(ctx context.Context) bool {
	_, ok := ctx.Value(keyTransaction).(*Transaction)
	return ok
}

func TXIDString(id [2]uint64) string {
	var txid [16]byte
	binary.LittleEndian.PutUint64(txid[:8], id[0])
	binary.LittleEndian.PutUint64(txid[8:], id[1])
	return fmt.Sprintf("%02x", txid)
}

// 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 {
	var id [2]uint64
	isSet := false
	for i, val := range h[headerID] {
		isSet = true
		if i >= len(id) {
			break
		}
		var err error
		id[i], err = strconv.ParseUint(val, 10, 64)
		if err != nil {
			return NewTransaction(parent)
		}
	}
	if !isSet {
		return NewTransaction(parent)
	}

	var path []uint32
	for i, seg := range strings.Split(h.Get(headerSubID), ".") {
		if i == 0 {
			if seg != "" {
				return NewTransaction(parent)
			}
			continue
		}
		id, err := strconv.ParseUint(seg, 10, 32)
		if err != nil {
			return NewTransaction(parent)
		}
		path = append(path, uint32(id))
	}

	value := &Transaction{
		TxID: id,
		Path: path,
	}
	return NewContext(parent, 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 := FromContext(parent)
	if !ok {
		// We've been passed a Context without an attached Trace identity.
		// This is probably the process making an outbound http request on its
		// own behalf - we'll assign it a new transaction id.
		return NewTransaction(parent)
	}

	spanCount := atomic.AddUint32(&parentTx.spanCount, 1)
	subTx := &Transaction{
		TxID: parentTx.TxID,
		Path: append(append(make([]uint32, 0, len(parentTx.Path)+1), parentTx.Path...), spanCount-1),
	}

	return NewContext(parent, 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 := FromContext(ctx)
	if !ok {
		h.Del(headerID)
		h.Del(headerSubID)
		return
	}
	h.Set(headerID, fmt.Sprintf("%d", value.TxID[0]))
	h.Add(headerID, fmt.Sprintf("%d", value.TxID[1]))
	h.Set(headerSubID, value.Path.String())
}

func NewTransaction(parent context.Context) context.Context {
	var id [2]uint64
	for i := range id {
		err := binary.Read(rand.Reader, binary.BigEndian, &id[i])
		if err != nil {
			log.Fatalf("entropy-error=%q", err)
		}
	}

	value := &Transaction{
		TxID: id,
		Path: []uint32(nil),
	}
	return NewContext(parent, value)
}
