/*

Package trace assists with instrumenting chitin processes.  It helps with
transfering trace information between processes, and it emits RPC lifecycle
events to a trace collector.

To simplify the process of sending trace events, the events may be sent in UDP
packets.  There may be a local daemon accepting packets over the loopback
interface (with the advantage of its larger MTU).  Not much is decided yet.

For now, we'll assume that we can send UDP packets to localhost:8943.  The
port comes from "TXID" converted to phone digits.  IANA shows the port as
unassigned, and it's low enough that the OS won't give the port away.  At the
very least, nothing bad is likely to come from sending the packets.

*/
package trace

import (
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"code.justin.tv/common/chitin/internal/github.com/showmax/go-fqdn"
	"code.justin.tv/common/chitin/internal/perp"
	"code.justin.tv/common/golibs/pkgpath"
	"code.justin.tv/release/trace/common"
	"code.justin.tv/release/trace/events"
	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"
)

const (
	flushInterval = 100 * time.Millisecond
	// The largest possible IP packet is just under 64kB.  Our production
	// linux loopback MTU is 16436, OSX appears to be 16384.  Our production
	// ethernet MTU is set to 1500, so we'll play it safe and use that as a
	// baseline (but save plenty of bytes for the various headers on layer 4
	// and below).
	maxPacket = 1400

	unknownHostname = "unknown.justin.tv"
	unknownSvcname  = "code.justin.tv/common/chitin/unknown"
)

var (
	keyProcInfo = new(key)
)

var (
	traceAddr = &net.UDPAddr{
		IP:   net.IPv4(127, 0, 0, 1),
		Port: 8943,
	}
)

var infoHolder = struct {
	sync.RWMutex
	ProcInfo
}{}

type ProcInfo struct {
	pid      int32
	hostname string
	svcname  string

	dest io.Writer
}

// An InfoOpt changes configuration on a ProcInfo value.  The design comes
// from Rob Pike's blog post "Self-referential functions and the design of
// options".
//
// http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
type InfoOpt func(*ProcInfo)

func SetDestination(w io.Writer) InfoOpt {
	return func(info *ProcInfo) {
		info.dest = &lockWriter{w: w}
	}
}

func ProcessOptIn(opts ...InfoOpt) error {
	infoHolder.Lock()
	defer infoHolder.Unlock()

	info := &infoHolder.ProcInfo

	for _, opt := range opts {
		opt(info)
	}

	return populateInfo(info)
}

func WithInfo(parent context.Context, opts ...InfoOpt) (context.Context, error) {
	// TODO(rhys): less terrible function name

	info := new(ProcInfo)

	for _, opt := range opts {
		opt(info)
	}

	err := populateInfo(info)
	if err != nil {
		return parent, err
	}
	return context.WithValue(parent, keyProcInfo, info), nil
}

func populateInfo(info *ProcInfo) error {
	if info.pid == 0 {
		info.pid = int32(os.Getpid())
		if info.pid == 0 {
			info.pid = -1
		}
	}

	if info.hostname == "" {
		info.hostname = fqdn.Get()

		if info.hostname == "" {
			// prevent retriggering this case
			info.hostname = unknownHostname
		}
	}

	if info.svcname == "" {
		name, ok := pkgpath.Main()
		if !ok || name == "" {
			name = unknownSvcname
		}
		info.svcname = name
	}

	if info.dest == nil {
		sbc, err := dialShardedBatchedConn()
		if err != nil {
			return err
		}
		info.dest = sbc
	}

	return nil
}

// Client events

// SendRequestHeadPrepared is called immediately before a call to
// (net/http.RoundTripper).RoundTrip
func SendRequestHeadPrepared(ctx context.Context, req *http.Request) {
	var method = common.Method(common.Method_value[req.Method])
	if req.Method == "" {
		method = common.Method_GET
	}

	peer := req.URL.Host
	if strings.LastIndex(peer, ":") > strings.LastIndex(peer, "]") {
		// has explicit port
	} else {
		port := map[string]string{"http": "80", "https": "443"}[req.URL.Scheme]
		// Blind concatenation with a colon is correct here - if the host is
		// an ipv6 address, it will already be surrounded with square brackets
		// (e.g. "[::1]").
		peer = req.URL.Host + ":" + port
	}

	SendEvent(ctx, &events.Event{
		Kind: common.Kind_REQUEST_HEAD_PREPARED,
		Extra: &events.Extra{
			Peer: peer,
			Http: &events.ExtraHTTP{
				Method:  method,
				UriPath: req.URL.Path,
			},
		},
	})
}

// SendRequestBodySent is called when req.Body is fully consumed by net/http.
// This could be via req.Body.Close, or via req.Body.Read returning an io.EOF.
func SendRequestBodySent(ctx context.Context) {
	SendEvent(ctx, &events.Event{
		Kind: common.Kind_REQUEST_BODY_SENT,
	})
}

// SendResponseHeadReceived is called immediately after a call to
// (net/http.RoundTripper).RoundTrip
func SendResponseHeadReceived(ctx context.Context, resp *http.Response, err error) {
	statusCode := -1
	if err == nil {
		statusCode = resp.StatusCode
	}

	SendEvent(ctx, &events.Event{
		Kind: common.Kind_RESPONSE_HEAD_RECEIVED,
		Extra: &events.Extra{
			Http: &events.ExtraHTTP{
				StatusCode: int32(statusCode),
			},
		},
	})
}

// SendResponseBodyReceived is called when resp.Body is fully consumed. This
// could be via resp.Body.Close, or via resp.Body.Read returning an io.EOF.
func SendResponseBodyReceived(ctx context.Context) {
	SendEvent(ctx, &events.Event{
		Kind: common.Kind_RESPONSE_BODY_RECEIVED,
	})
}

// Server events

// SendRequestHeadReceived is called immediately before a call to
// (net/http.Handler).ServeHTTP
func SendRequestHeadReceived(ctx context.Context, req *http.Request) {
	SendEvent(ctx, &events.Event{
		Kind: common.Kind_REQUEST_HEAD_RECEIVED,
		Extra: &events.Extra{
			Peer: req.RemoteAddr,
			Http: &events.ExtraHTTP{
				Method:  common.Method(common.Method_value[req.Method]),
				UriPath: req.URL.Path,
			},
		},
	})
}

// SendResponseHeadPrepared is called immediately before a call to
// (net/http.ResponseWriter).WriteHeader
func SendResponseHeadPrepared(ctx context.Context, statusCode int) {
	SendEvent(ctx, &events.Event{
		Kind: common.Kind_RESPONSE_HEAD_PREPARED,
		Extra: &events.Extra{
			Http: &events.ExtraHTTP{
				StatusCode: int32(statusCode),
			},
		},
	})
}

func SendEvent(ctx context.Context, msg *events.Event) {
	info, ok := ctx.Value(keyProcInfo).(*ProcInfo)
	if !ok {
		info = new(ProcInfo)

		infoHolder.RLock()
		*info = infoHolder.ProcInfo
		infoHolder.RUnlock()

		if info.dest == nil {
			atomic.AddInt64(&exp.NoProcInfo, 1)
			return
		}
	}

	tx, ok := FromContext(ctx)
	if !ok {
		atomic.AddInt64(&exp.NoTransaction, 1)
		return
	}

	msg.Time = time.Now().UnixNano()

	msg.Pid = info.pid
	msg.Hostname = info.hostname
	msg.Svcname = info.svcname

	msg.TransactionId = tx.TxID[:]
	msg.Path = tx.Path

	set := &events.EventSet{Event: []*events.Event{msg}}
	buf, err := proto.Marshal(set)
	if err != nil {
		atomic.AddInt64(&exp.MarshalError, 1)
		return
	}

	n, err := info.dest.Write(buf)
	if n != len(buf) {
		atomic.AddInt64(&exp.ShortWrite, 1)
	}
	if err != nil {
		atomic.AddInt64(&exp.WriteError, 1)
		return
	}
}

func dialShardedBatchedConn() (io.Writer, error) {
	// count the number of runtime P structures, without changing the value
	pcount := runtime.GOMAXPROCS(0)

	ws := make([]io.Writer, pcount*2)
	for i := range ws {
		conn, err := net.DialUDP("udp", nil, traceAddr)
		if err != nil {
			return nil, fmt.Errorf("can't dial UDP: %s", err)
		}
		ws[i] = newBatchWriter(conn, maxPacket)
	}
	return newPerPWriter(ws), nil
}

type perPWriter struct {
	pool    *perp.Pool
	writers []io.Writer
}

func newPerPWriter(ws []io.Writer) *perPWriter {
	return &perPWriter{
		pool:    perp.NewPoolSize(len(ws)),
		writers: append([]io.Writer(nil), ws...),
	}
}

func (pw *perPWriter) Write(p []byte) (int, error) {
	ctx := context.Background()
	var (
		ok  bool
		n   int
		err error
	)

	pw.pool.Do(ctx, func(i int) {
		ok = true
		w := pw.writers[i]
		n, err = w.Write(p)
	})
	if !ok {
		err = ctx.Err()
	}

	return n, err
}

type lockWriter struct {
	mu sync.Mutex
	w  io.Writer
}

func (l *lockWriter) Write(p []byte) (int, error) {
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.w.Write(p)
}

var (
	errTooBig = errors.New("trace: message is too big for a single packet")
)

type batchWriter struct {
	mu           sync.Mutex
	w            io.Writer
	buf          []byte
	flushPending bool

	timer *time.Timer
}

func newBatchWriter(w io.Writer, size int) io.Writer {
	b := &batchWriter{
		w:   w,
		buf: make([]byte, 0, size),
	}

	b.timer = time.AfterFunc(flushInterval, func() {
		err := b.Flush()
		if err != nil {
			atomic.AddInt64(&exp.FlushError, 1)
		}
	})
	b.timer.Stop()

	return b
}

func (b *batchWriter) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()

	if !b.flushPending {
		b.timer.Reset(flushInterval)
		b.flushPending = true
	}

	if len(b.buf)+len(p) <= cap(b.buf) {
		b.buf = append(b.buf, p...)
		return len(p), nil
	}

	if len(p) <= cap(b.buf) {
		_, err := b.w.Write(b.buf)
		// In the case of a short write, we've got to dump the remaining
		// buffer.  We've lost track of the datagram boundaries, and we
		// absolutely cannot send out datagrams that begin mid-message.  So in
		// no case do we keep any part of the existing buffer around - we'll
		// just return the error we got from the real Write to our caller,
		// along with the number of bytes our caller sent.
		b.buf = append(b.buf[:0], p...)
		return len(p), err
	}

	return 0, errTooBig
}

func (b *batchWriter) Flush() error {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.flushPending {
		b.timer.Stop()
		b.flushPending = false
	}

	if len(b.buf) == 0 {
		return nil
	}

	_, err := b.w.Write(b.buf)
	// Just like in Write, we have to dump the remainder of the buffer in the
	// case of a short write.  The message boundaries are lost, so we'll just
	// bubble up the error that our io.Writer gave us and carry on.
	b.buf = b.buf[:0]
	return err
}
