package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"sort"
	"syscall"
	"time"

	"code.justin.tv/common/chitin"
	_ "code.justin.tv/common/golibs/bininfo"
	"code.justin.tv/release/trace/analysis/tx"
	"code.justin.tv/release/trace/api"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

var (
	serverHostPort = flag.String("server", "trace-api.prod.us-west2.justin.tv:11143", "The server hostport")
	host           = flag.String("host", "", "Fully-qualified name of host to inspect")
	sampling       = flag.Float64("sample", 0.1, "Sampling rate")
	duration       = flag.Duration("t", time.Second, "How long to capture data before returning a result")
)

func main() {
	flag.Parse()

	err := chitin.ExperimentalTraceProcessOptIn()
	if err != nil {
		log.Fatalf("trace enable err=%q", err)
	}

	conn, err := grpc.Dial(*serverHostPort, grpc.WithInsecure())
	if err != nil {
		log.Fatalf("dial err=%q", err)
	}
	defer conn.Close()
	client := api.NewTraceClient(conn)
	ctx, cancel := context.WithTimeout(context.Background(), *duration)

	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, syscall.SIGINT)
	go func() {
		<-sigch
		cancel()
	}()

	firehose, err := client.Firehose(ctx, &api.FirehoseRequest{
		Sampling: *sampling,
		Query: &api.Query{
			Comparisons: []*api.Comparison{
				&api.Comparison{
					ServiceHost: &api.StringComparison{Value: *host},
				},
			},
		},
	})
	if err != nil {
		log.Fatalf("firehose err=%q", err)
	}

	result := consumeFirehose(firehose)
	if _, err := result.WriteTo(os.Stdout); err != nil {
		log.Fatalf("dump err=%q", err)
	}
}

func consumeFirehose(firehose api.Trace_FirehoseClient) io.WriterTo {
	var tx *api.Transaction
	var err error

	var d = &data{}

	for tx, err = firehose.Recv(); err == nil; tx, err = firehose.Recv() {
		var match func(*api.Call) bool
		match = func(call *api.Call) bool {
			if svc := call.GetSvc(); svc != nil && svc.Host == *host {
				return true
			}
			return false
		}

		var record func(*api.Call)
		record = func(call *api.Call) {
			start, end := api.ServiceStart(call), api.ServiceEnd(call)
			if start.IsZero() || end.IsZero() {
				return
			}
			d.calls = append(d.calls, &node{tx: tx, call: call})
			// log.Fatalf("%s", proto.MarshalTextString(tx))
		}

		var search func(*api.Call)
		search = func(call *api.Call) {
			if match(call) {
				record(call)
			}
			for _, sub := range call.Subcalls {
				search(sub)
			}
		}
		search(tx.Root)
	}

	if err != nil {
		switch grpc.Code(err) {
		case codes.DeadlineExceeded:
		default:
			d.err = err
		}
	}

	sort.Sort(MakeInterface(len(d.calls), func(i, j int) { d.calls[i], d.calls[j] = d.calls[j], d.calls[i] },
		func(i, j int) bool {
			return api.ServiceStart(d.calls[i].call).Before(api.ServiceStart(d.calls[j].call))
		},
	))

	// trim off the ends, which are often outliers on the time axis
	if l := len(d.calls); l > 10 {
		d.calls = d.calls[l/10 : 9*l/10]
	}

	var earliest time.Time
	pidThreads := make(map[int]map[int]time.Time)
	for _, node := range d.calls {
		start, end := api.ServiceStart(node.call), api.ServiceEnd(node.call)
		if earliest.IsZero() {
			earliest = start
		}
		if start.Before(earliest) || end.Before(earliest) {
			continue
		}

		type event struct {
			Name      string      `json:"name"`
			Category  string      `json:"cat"`
			Phase     string      `json:"ph"`
			Timestamp float64     `json:"ts"`
			Duration  float64     `json:"dur"`
			ID        string      `json:"id"`
			Pid       int         `json:"pid"`
			Tid       int         `json:"tid"`
			Args      interface{} `json:"args"`
		}

		// find the lowest-numbered thread id to "claim" this call
		var pid int
		if node.call.GetSvc() != nil {
			pid = int(node.call.GetSvc().Pid)
		}
		var tid int
		threads, ok := pidThreads[pid]
		if !ok {
			threads = make(map[int]time.Time)
			pidThreads[pid] = threads
			var name string
			if node.call.GetSvc() != nil {
				name = node.call.GetSvc().Name
			}
			d.Events = append(d.Events, event{
				Name:  "process_name",
				Phase: "M",
				Pid:   pid,
				Args: struct {
					Name string `json:"name"`
				}{
					Name: name,
				},
			})
		}
		for tid = 0; tid < len(threads); tid++ {
			if start.After(threads[tid]) {
				threads[tid] = end
				break
			}
		}
		if tid == len(threads) {
			threads[tid] = end
		}
		// thread ids start at 1 in the catapult ui
		tid += 1

		name := callName(node.call)

		d.Events = append(d.Events, event{
			Name:      name,
			Phase:     "X",
			Timestamp: start.Sub(earliest).Seconds() * float64(time.Second/time.Microsecond),
			Duration:  end.Sub(start).Seconds() * float64(time.Second/time.Microsecond),
			Pid:       pid,
			Tid:       tid,
			Args: map[string]interface{}{
				"txid": idForTx(node.tx).String() + pathString(node.call.Path),
			},
		})
	}

	return d
}

func callName(call *api.Call) string {
	if p := call.GetParams().GetHttp(); p != nil {
		if p.Route != "" {
			return p.Route
		}
		return fmt.Sprintf("HTTP %s", p.Method)
	}
	if p := call.GetParams().GetGrpc(); p != nil {
		return p.Method
	}
	return "unknown"
}

type sorter struct {
	Length int
	SwapFn func(i, j int)
	LessFn func(i, j int) bool
}

func (s sorter) Len() int           { return s.Length }
func (s sorter) Swap(i, j int)      { s.SwapFn(i, j) }
func (s sorter) Less(i, j int) bool { return s.LessFn(i, j) }

func MakeInterface(n int, swap func(i, j int), less func(i, j int) bool) sort.Interface {
	return sorter{Length: n, SwapFn: swap, LessFn: less}
}

type data struct {
	calls []*node
	err   error

	Events []interface{} `json:"traceEvents"`
}

func (d *data) WriteTo(w io.Writer) (int64, error) {
	var nn int64

	buf, err := json.Marshal(d)
	if err != nil {
		return nn, err
	}

	n, err := fmt.Fprintf(w, "%s\n", buf)
	nn += int64(n)
	if err != nil {
		return nn, err
	}

	if d.err != nil {
		return nn, d.err
	}
	return nn, nil
}

type node struct {
	tx   *api.Transaction
	call *api.Call
}

func idForTx(t *api.Transaction) tx.TransactionID {
	if t == nil {
		return tx.TransactionID{0, 0}
	}
	txid := t.TransactionId
	switch len(txid) {
	case 0:
		return tx.TransactionID{0, 0}
	case 1:
		return tx.TransactionID{0, txid[0]}
	default:
		return tx.TransactionID{txid[0], txid[1]}
	}
}

func pathString(path []uint32) string {
	var b bytes.Buffer
	for _, p := range path {
		fmt.Fprintf(&b, ".%d", p)
	}
	return b.String()
}
