package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"code.justin.tv/common/chitin"
	_ "code.justin.tv/common/golibs/bininfo"
	"code.justin.tv/release/trace/api"
	"github.com/spenczar/tdigest"
	"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")
	caller = flag.String("caller", "code.justin.tv/edge/visage/cmd/visage",
		"The caller service to analyze")
	callee = flag.String("callee", "code.justin.tv/web/cohesion/cmd/cohesion",
		"The callee service to analyze")
	sampling      = flag.Float64("sample", 0.1, "Sampling rate")
	duration      = flag.Duration("t", time.Second, "How long to capture data before returning a result")
	callThreshold = flag.Int("fanout", 30, "usage")
)

const (
	compression = 100 // sane default for tdigest
)

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{
					ServiceName: &api.StringComparison{Value: *caller},
				},
			},
			Subqueries: []*api.Query{
				&api.Query{
					Comparisons: []*api.Comparison{
						&api.Comparison{
							ServiceName: &api.StringComparison{Value: *callee},
						},
					},
				},
			},
		},
	})
	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{
		caller: newDatum(),
		callee: newDatum(),
	}

	for tx, err = firehose.Recv(); err == nil; tx, err = firehose.Recv() {
		var match func(*api.Call) bool
		match = func(call *api.Call) bool {
			count := 0
			if svc := call.GetSvc(); svc != nil && svc.Name == *caller {
				for _, sub := range call.Subcalls {
					if svc := sub.GetSvc(); svc != nil && svc.Name == *callee {
						count++
						if count >= *callThreshold {
							return true
						}
					}
				}
			}
			return false
		}

		var record func(*api.Call)
		record = func(call *api.Call) {
			d.caller.count++
			d.caller.clientDurations.Add(api.ServiceDuration(call).Seconds(), 1)
			d.caller.serviceDurations.Add(api.ServiceDuration(call).Seconds(), 1)

			for _, sub := range call.Subcalls {
				if svc := sub.GetSvc(); svc != nil && svc.Name == *callee {
					d.callee.count++
					d.callee.clientDurations.Add(api.ClientDuration(sub).Seconds(), 1)
					d.callee.serviceDurations.Add(api.ServiceDuration(sub).Seconds(), 1)
				}
			}
		}

		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
		}
	}
	return d
}

type data struct {
	caller *datum
	callee *datum
	err    error
}

type datum struct {
	clientDurations  tdigest.TDigest
	serviceDurations tdigest.TDigest
	count            int
}

func newDatum() *datum {
	return &datum{
		clientDurations:  tdigest.New(compression),
		serviceDurations: tdigest.New(compression),
	}
}

func p(q float64) string {
	p := fmt.Sprintf("%0.5f", q)
	if !strings.HasPrefix(p, "0.") {
		p = "100"
	} else {
		p = strings.SplitN(p, ".", 2)[1]
		if l := 2; len(p) > l {
			p = p[:l] + strings.TrimRight(p[l:], "0")
		}
	}
	return p
}

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

	quantiles := []float64{0.00, 0.50, 0.90, 0.95, 0.99, 0.999}

	fmt.Fprintf(w, "caller: %6d calls    %s\n", d.caller.count, *caller)
	fmt.Fprintf(w, "--------------------\n")
	for _, q := range quantiles {
		dur := time.Duration(float64(time.Second) * d.caller.serviceDurations.Quantile(q))
		if d.caller.count == 0 {
			dur = 0
		}
		fmt.Fprintf(w, " % 5s: %10.3fms\n", "p"+p(q),
			dur.Seconds()*float64(time.Second/time.Millisecond))
	}
	fmt.Fprintf(w, "\n")
	fmt.Fprintf(w, "callee: %6d calls    %s\n", d.callee.count, *callee)
	fmt.Fprintf(w, "--------------------\n")
	for _, q := range quantiles {
		cdur := time.Duration(float64(time.Second) * d.callee.clientDurations.Quantile(q))
		sdur := time.Duration(float64(time.Second) * d.callee.serviceDurations.Quantile(q))
		if d.callee.count == 0 {
			cdur, sdur = 0, 0
		}
		fmt.Fprintf(w, " % 5s: %10.3fms / %10.3fms\n", "p"+p(q),
			cdur.Seconds()*float64(time.Second/time.Millisecond),
			sdur.Seconds()*float64(time.Second/time.Millisecond))
	}

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