package main

import (
	"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/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")
	service        = flag.String("svc", "code.justin.tv/web/web", "The 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")
	quantile       = flag.Float64("q", 0.90, "Timing quantile to report, e.g. 0.90 for p90")
)

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: *service},
				},
			},
		},
	})
	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{
		hosts: make(map[string]*datum),
		datum: datum{
			digest: tdigest.New(compression),
		},
	}

	for tx, err = firehose.Recv(); err == nil; tx, err = firehose.Recv() {
		var search func(*api.Call)
		search = func(call *api.Call) {
			if svc := call.GetSvc(); svc != nil && svc.Name == *service {
				host := svc.Host
				dat, ok := d.hosts[host]
				if !ok {
					dat = &datum{
						digest: tdigest.New(compression),
					}
					d.hosts[host] = dat
				}
				dur := api.ServiceDuration(call)
				dat.digest.Add(dur.Seconds(), 1)
				dat.count++
				d.datum.digest.Add(dur.Seconds(), 1)
				d.datum.count++
			}
			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 {
	hosts map[string]*datum
	datum datum
	err   error
}

type datum struct {
	digest tdigest.TDigest
	count  int
}

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

	const total = "total"
	var longest = len(total)

	var hosts []string
	for host := range d.hosts {
		hosts = append(hosts, host)
		if len(host) > longest {
			longest = len(host)
		}
	}
	sort.Strings(hosts)

	dur := time.Duration(float64(time.Second) * d.datum.digest.Quantile(*quantile))
	n, err := fmt.Fprintf(w, "% *s: %10.3fms, %5d calls\n", longest, total,
		dur.Seconds()*float64(time.Second/time.Millisecond), d.datum.count)
	nn += int64(n)
	if err != nil {
		return nn, err
	}

	for _, host := range hosts {
		datum := d.hosts[host]
		dur := time.Duration(float64(time.Second) * datum.digest.Quantile(*quantile))
		n, err := fmt.Fprintf(w, "% *s: %10.3fms, %5d calls\n",
			longest, host,
			dur.Seconds()*float64(time.Second/time.Millisecond), datum.count)
		nn += int64(n)
		if err != nil {
			return nn, err
		}
	}

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