package trace

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

	"code.justin.tv/release/trace"

	"crypto/rand"

	"golang.org/x/net/context"
	"google.golang.org/grpc/metadata"
)

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

type key int

var (
	keyTransaction = new(key)
)

type transaction struct {
	TxID [2]uint64
	Path callPath
	// accessed atomically
	spanCount uint32
}

// 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 := ctx.Value(keyTransaction).(*transaction)
	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 := ctx.Value(keyTransaction).(*transaction)
	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)
}

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 context.WithValue(parent, keyTransaction, value)
}

// ContextFromGRPCMetadata issues a new Context based on parent, but including
// the Trace values in the MD. If the Trace values are missing, a new
// transaction id (with high entropy) is issued.
func ContextFromGRPCMetadata(parent context.Context, md metadata.MD) context.Context {
	var id [2]uint64
	var path []uint32

	// TODO(rhys): pull data out of MD
	slc := md[metadataKey]
	if len(slc) != 1 {
		return newTransaction(parent)
	}
	str := slc[0]
	rest := strings.TrimLeft(str, "0123456789abcdefABCDEF")
	hexPrefix := str[:len(str)-len(rest)]
	if len(hexPrefix) != len(id)*16 {
		return newTransaction(parent)
	}

	for i := range id {
		b, err := hex.DecodeString(hexPrefix[i*16 : (i+1)*16])
		if err != nil {
			return newTransaction(parent)
		}
		id[i] = binary.LittleEndian.Uint64(b)
	}

	for i, seg := range strings.Split(rest, ".") {
		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 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 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 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("%d", value.TxID[0]))
	h.Add(headerID, fmt.Sprintf("%d", value.TxID[1]))
	h.Set(headerSubID, value.Path.String())
}

// AugmentGRPCMetadata adds Trace transaction information to a gRPC metadata
// value. It does not modify the internal Trace record state.
func AugmentGRPCMetadata(ctx context.Context, md metadata.MD) {
	value, ok := ctx.Value(keyTransaction).(*transaction)
	if !ok {
		delete(md, metadataKey)
		return
	}

	txid := txidString(value.TxID)
	path := value.Path.String()

	md[metadataKey] = []string{txid + path}
}

func newTransaction(parent context.Context) context.Context {
	var id [2]uint64
	for i := range id {
		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{}{}:
			}
			atomic.AddInt64(&exp.EntropyUnderflow, 1)
			// Be sure to get what we came for. Yes, it will probably involve
			// waiting this time.
			id[i] = <-ids
		case id[i] = <-ids:
		}
	}

	value := &transaction{
		TxID: id,
		Path: []uint32(nil),
	}
	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")
		}
	}
}
