package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"runtime"
	"sort"
	"sync"
	"time"

	"code.justin.tv/common/chitin"
	"code.justin.tv/release/trace/analysis/routes"
	"code.justin.tv/release/trace/analysis/tx"
	"code.justin.tv/release/trace/api"
	"code.justin.tv/release/trace/scanproto"
	"code.justin.tv/rhys/nursery/trace/txstore"
	"code.justin.tv/spencer/quantiles"
	"github.com/ajstarks/svgo"
	"github.com/golang/protobuf/proto"
	"github.com/spenczar/tdigest"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/opt"
	"github.com/syndtr/goleveldb/leveldb/util"
)

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

var _ quantiles.QuantileEstimator = tdigest.New(compression)

var quants = []float64{0.00, 0.25, 0.50, 0.75, 0.90, 0.95, 0.99, 1.00}

func main() {
	filename := flag.String("input", "/dev/stdin", "Input filename")
	dbdir := flag.String("leveldb", "", "Path to leveldb database")
	hostport := flag.String("http", "127.0.0.1:8080", "Host and port for http listener")
	txCount := flag.Int("count", 100000, "Number of transactions to process")
	targetSvc := flag.String("svc", "code.justin.tv/web/web", "Name of service to inspect")
	targetRPC := flag.String("rpc", "api/twitch/viewer#info", "Name of RPC to inspect")
	tick := flag.Duration("tick", 1*time.Second, "Bucket size")

	svgOut := flag.String("svgout", "", "SVG output filename")

	outfile := flag.String("output", "/dev/null", "Output filename")
	start := flag.String("start", "2016-04-01T19:14:00Z", "Start time for writing to output file")
	recDur := flag.Duration("duration", 10*time.Second, "Duration to write to output file")

	flag.Parse()

	startT, err := time.Parse(time.RFC3339Nano, *start)
	if err != nil {
		log.Fatalf("Parse err=%q", err)
	}

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

	if *svgOut != "" {
		outf, err := os.OpenFile(*svgOut, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
		if err != nil {
			log.Fatalf("OpenFile err=%q", err)
		}
		defer func() {
			err := outf.Close()
			if err != nil {
				log.Fatalf("Close err=%q", err)
			}
		}()

		demoSVG(outf)
		return
	}

	var mu sync.Mutex
	var contents []byte

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "image/svg+xml")

		mu.Lock()
		b := contents
		mu.Unlock()

		fmt.Fprintf(w, "%s", b)
	})

	http.HandleFunc("/new", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "image/svg+xml")

		demoSVG(w)
	})

	srv := &http.Server{
		Handler: chitin.Handler(http.DefaultServeMux),
	}

	l, err := net.Listen("tcp", *hostport)
	if err != nil {
		log.Fatalf("Listen err=%q", err)
	}

	go func() {
		err := srv.Serve(l)
		if err != nil {
			log.Fatalf("Serve err=%q", err)
		}
	}()

	var src scanproto.TransactionSource

	var db *leveldb.DB
	switch {
	case *dbdir != "":
		db, err = leveldb.OpenFile(*dbdir, &opt.Options{
			ErrorIfMissing:         true,
			ReadOnly:               true,
			OpenFilesCacheCapacity: 200,
		})
		if err != nil {
			log.Fatalf("OpenFile err=%q", err)
		}
		// We'll keep the db open in order to serve dynamic http requests

		it := db.NewIterator(util.BytesPrefix([]byte(`txid="`)), nil)
		src = txstore.NewReaderLevelDB(it, runtime.NumCPU())
	default:
		f, err := os.Open(*filename)
		if err != nil {
			log.Fatalf("Open err=%q", err)
		}
		defer func() {
			err := f.Close()
			if err != nil {
				log.Fatalf("Close err=%q", err)
			}
		}()
		src = txstore.NewReaderSource(f, runtime.NumCPU())
	}

	http.Handle("/db/", &levelHandler{db: db, uriPrefix: "/db"})

	outf, err := os.OpenFile(*outfile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatalf("OpenFile err=%q", err)
	}
	defer func() {
		err := outf.Close()
		if err != nil {
			log.Fatalf("Close err=%q", err)
		}
	}()
	out := bufio.NewWriter(outf)
	defer func() {
		err := out.Flush()
		if err != nil {
			log.Fatalf("Flush err=%q", err)
		}
	}()

	defer src.Stop()
	go src.Run()

	alltime := make(map[selector]*segment)
	smalltime := make(map[selector]map[timespan]*segment)

	// sc := scanproto.NewEventSetScanner(f)
	// sc.Buffer(nil, maxTxSize)
	errs := src.Errors()
	txs := src.Transactions()
	for i, n := 0, 0; i < *txCount; i++ {
		if errs == nil && txs == nil {
			break
		}
		// var txs api.TransactionSet

		// err := proto.Unmarshal(sc.Bytes(), &txs)
		// if err != nil {
		// 	log.Fatalf("Unmarshal err=%q bytes %s", err, sc.Bytes())
		// }

		// for _, tx := range txs.Transaction {

		select {
		case err, ok := <-errs:
			if !ok {
				errs = nil
				continue
			}
			if err != nil {
				log.Fatalf("read err=%q", err)
			}
			continue
		case tx, ok := <-txs:
			if !ok {
				txs = nil
				continue
			}
			save := false
			eachCall([]*api.Call{tx.GetRoot()}, func(c *api.Call) {
				sd := api.ServiceDuration(c)

				var serverSpan timespan
				var t int64
				if stamps := c.ServiceTimestamps; len(stamps) > 0 {
					t = stamps[0].Time

					serverSpan = timespan{
						seconds:  stamps[0].Time / 1e9,
						nanos:    int32(stamps[0].Time % 1e9),
						duration: time.Duration(stamps[len(stamps)-1].Time - stamps[0].Time),
					}
				}

				if t := time.Unix(0, t); !startT.After(t) && startT.Add(*recDur).After(t) {
					save = true
				}

				for _, s := range callSelectors(c) {
					seg := alltime[s]
					if seg == nil {
						seg = &segment{
							key: s,
						}
						seg.asServer.quant = newQuantile()
						seg.asClient.quant = newQuantile()
						seg.asHTTPClient.quant = newQuantile()
						seg.asSQLClient.quant = newQuantile()
					}
					alltime[s] = seg

					seg.count++
					seg.asServer.total += sd
					if false {
						seg.asServer.quant.Add(sd.Seconds(), 1)
					}

					if t == 0 {
						continue
					}

					rounded := ((t / tick.Nanoseconds()) * tick.Nanoseconds())
					span := timespan{ // 1e9 shows up here per the definition of api.Timestamp
						seconds:  int64(rounded / 1e9),
						nanos:    int32(rounded % 1e9),
						duration: *tick,
					}
					m := smalltime[s]
					if m == nil {
						m = make(map[timespan]*segment)
						smalltime[s] = m
					}
					seg = m[span]
					if seg == nil {
						seg = &segment{
							key:  s,
							span: span,
						}
						seg.asServer.quant = newQuantile()
						seg.asClient.quant = newQuantile()
						seg.asHTTPClient.quant = newQuantile()
						seg.asSQLClient.quant = newQuantile()
					}
					m[span] = seg

					seg.count++
					seg.asServer.total += sd
					if false || seg.count == 1 || rand.Intn(10) == 0 {
						seg.asServer.quant.Add(sd.Seconds(), 1)
					}

					for _, pair := range []struct {
						keep func(*api.Call) bool
						ptr  *timeStats
					}{} {
						var calls []timespan
						for _, sub := range c.Subcalls {
							if !pair.keep(sub) {
								continue
							}

							times := sub.GetClientTimestamps()
							if len(times) < 2 {
								continue
							}
							calls = append(calls, timespan{ // 1e9 shows up here per the definition of api.Timestamp
								seconds:  times[0].Time / 1e9,
								nanos:    int32(times[0].Time % 1e9),
								duration: time.Duration(times[len(times)-1].Time - times[0].Time),
							})
						}
						covered := serverSpan.coverage(calls...)

						pair.ptr.total += covered
						if false || seg.count == 1 || rand.Intn(10) == 0 {
							pair.ptr.quant.Add(covered.Seconds(), 1)
						}
					}
				}
			})
			if save {
				b, err := proto.Marshal(&api.TransactionSet{Transaction: []*api.Transaction{tx}})
				if err == nil {
					fmt.Fprintf(out, "%s", b)
				}
			}
		}

		if i > n {
			log.Printf("i=%d", i)
			n += n + 1
		}
	}
	// err = sc.Err()
	// if err != nil {
	// 	log.Fatalf("Scan err=%q", err)
	// }

	var buf bytes.Buffer
	g := svg.New(&buf)
	g.Start(10000, 2000)

	var ss selectors
	for s := range alltime {
		ss = append(ss, s)
	}
	sort.Sort(ss)
	var i int
	for _, s := range ss {
		if s.hostname != "" {
			continue
		}
		seg := alltime[s]
		if seg.count < 100 {
			continue
		}

		log.Printf("svcname=%q rpc=%q hostname=%q count=%d total=%f p90=%f p99=%f",
			seg.key.svcname, seg.key.rpc, seg.key.hostname,
			seg.count, seg.asServer.total.Seconds(),
			seg.asServer.quant.Quantile(0.90),
			seg.asServer.quant.Quantile(0.99))

		i++
	}

	_, _ = *targetSvc, *targetRPC

	for j, sel := range []selector{
		// {svcname: *targetSvc, rpc: *targetRPC},
		{svcname: "code.justin.tv/chat/tmi/clue", rpc: "authenticate"},
		{svcname: "code.justin.tv/video/usher/api/usher", rpc: ""},
		{svcname: "code.justin.tv/web/streams-api", rpc: "followed"},
		{svcname: "code.justin.tv/web/streams-api", rpc: "stream"},
		{svcname: "code.justin.tv/web/streams-api", rpc: "streams"},
		{svcname: "code.justin.tv/web/web", rpc: "api/channels#access_token"},
		{svcname: "code.justin.tv/web/web", rpc: "api/channels#chat_properties"},
		{svcname: "code.justin.tv/web/web", rpc: "api/channels#show"},
		{svcname: "code.justin.tv/web/web", rpc: "api/channels#viewer"},
		{svcname: "code.justin.tv/web/web", rpc: "api/game_follows#live"},
		{svcname: "code.justin.tv/web/web", rpc: "api/twitch/viewer#info"},
		{svcname: "code.justin.tv/web/web", rpc: "api/twitch/viewer#token"},
		{svcname: "code.justin.tv/web/web", rpc: "kraken/follows#show"},
		{svcname: "code.justin.tv/web/web", rpc: "kraken/users#index"},
		{svcname: "code.justin.tv/web/web", rpc: "kraken/users#show"},
	} {
		target := smalltime[sel]
		var ts timespans
		for s := range target {
			ts = append(ts, s)
		}
		sort.Sort(ts)
		log.Printf("len=%d", len(ts))

		i = 0
		var qs []quartiles
		for _, s := range ts {
			seg := target[s]

			stats := seg.asServer

			log.Printf("svcname=%q rpc=%q hostname=%q start=%s count=%d total=%f p90=%f p99=%f",
				seg.key.svcname, seg.key.rpc, seg.key.hostname, seg.span.Start().Format(time.RFC3339Nano),
				seg.count, stats.total.Seconds(),
				stats.quant.Quantile(0.90),
				stats.quant.Quantile(0.99))

			qs = append(qs, quartiles{
				0: stats.quant.Quantile(0.00),
				1: stats.quant.Quantile(0.50),
				2: stats.quant.Quantile(0.90),
				3: stats.quant.Quantile(0.95),
				4: stats.quant.Quantile(0.99),
			})

			i++
		}

		c := chart{
			tick:   4,
			height: 60,

			yscale: (log10{logmin: -3, logmax: 1, perDecade: 15}).scale,
			yticks: []float64{1, 0.1, 0.01},
		}
		g.Translate(20, 20+80*j)
		g.Text(0, 0,
			fmt.Sprintf("svcname=%q rpc=%q hostname=%q", sel.svcname, sel.rpc, sel.hostname),
			"font-size:14;fill:red")
		c.plotQuartiles(g, qs)
		g.Gend()
	}
	g.End()

	mu.Lock()
	contents = buf.Bytes()
	mu.Unlock()

	// trigger deferred functions
	runtime.Goexit()
}

type timeStats struct {
	quant quantiles.QuantileEstimator
	total time.Duration
}

type segment struct {
	count        int64
	asServer     timeStats
	asClient     timeStats
	asHTTPClient timeStats
	asSQLClient  timeStats

	key  selector
	span timespan
}

type selector struct {
	svcname  string
	rpc      string
	hostname string
}

type selectors []selector

func (ss selectors) Len() int      { return len(ss) }
func (ss selectors) Swap(i, j int) { ss[i], ss[j] = ss[j], ss[i] }
func (ss selectors) Less(i, j int) bool {
	si, sj := ss[i], ss[j]
	if si.svcname != sj.svcname {
		return si.svcname < sj.svcname
	}
	if si.rpc != sj.rpc {
		return si.rpc < sj.rpc
	}
	if si.hostname != sj.hostname {
		return si.hostname < sj.hostname
	}
	return false
}

type timespan struct {
	seconds  int64
	nanos    int32
	duration time.Duration
}

func (t timespan) Start() time.Time {
	return time.Unix(t.seconds, int64(t.nanos)).UTC()
}

func (t timespan) End() time.Time {
	return t.Start().Add(t.duration)
}

type timespans []timespan

func (ts timespans) Len() int      { return len(ts) }
func (ts timespans) Swap(i, j int) { ts[i], ts[j] = ts[j], ts[i] }
func (ts timespans) Less(i, j int) bool {
	si, sj := ts[i], ts[j]
	if sis, sjs := si.Start(), sj.Start(); !sis.Equal(sjs) {
		return sis.Before(sjs)
	}
	if si.duration != sj.duration {
		return si.duration < sj.duration
	}
	return false
}

func eachCall(cs []*api.Call, fn func(c *api.Call)) {
	for _, c := range cs {
		if c == nil {
			continue
		}
		fn(c)
		eachCall(c.Subcalls, fn)
	}
}

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

	}
}

func callSelectors(c *api.Call) selectors {
	ss := make(selectors, 0, 4)

	ss = append(ss, selector{})

	svc := c.Svc
	if svc == nil {
		return ss
	}
	if svc.Name == "" {
		return ss
	}
	ss = append(ss, selector{svcname: svc.Name})

	var rpc string
	switch p := c.GetParams(); {
	case p.GetHttp() != nil:
		if r := p.GetHttp().Route; r != "" {
			rpc = r
		} else {
			rm := routes.Routes[svc.Name]
			rpc = rm.LookupTrimSlash(p.GetHttp().UriPath, p.GetHttp().Method.String())
		}
	case p.GetGrpc() != nil:
		rpc = p.GetGrpc().Method
	}

	if rpc == "" {
		return ss
	}
	ss = append(ss, selector{svcname: svc.Name, rpc: rpc})

	if svc.Host == "" {
		return ss
	}
	ss = append(ss, selector{svcname: svc.Name, rpc: rpc, hostname: svc.Host})

	return ss
}

func newQuantile() quantiles.QuantileEstimator {
	// return quantScale{est: newFrugal(quants...)}
	return quantiles.NewExactQuantileEstimator()
}

type quantScale struct {
	est quantiles.QuantileEstimator
}

func (qs quantScale) Add(val float64, weight int)      { qs.est.Add(val*1e3, weight) }
func (qs quantScale) Quantile(q float64) (val float64) { return qs.est.Quantile(q) / 1e3 }

type frugal struct {
	quantiles  []float64
	estimators []quantiles.QuantileEstimator
}

func newFrugal(q ...float64) frugal {
	f := frugal{
		quantiles: append([]float64(nil), q...),
	}
	sort.Float64s(f.quantiles)
	for _, q := range f.quantiles {
		f.estimators = append(f.estimators, quantiles.NewFrugalQuantileEstimator(q))
	}
	return f
}

func (f frugal) Add(val float64, weight int) {
	for _, e := range f.estimators {
		e.Add(val, weight)
	}
}

func (f frugal) Quantile(q float64) (val float64) {
	for i := range f.quantiles {
		if f.quantiles[i] == q {
			return f.estimators[i].Quantile(q)
		}
	}
	panic(fmt.Errorf("unexpected quantile %f", q))
}
