package main

import (
	"bufio"
	"bytes"
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"math"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"
)

type txid uint64
type spanid string
type ns int64

type record struct {
	Kind   recordType `json:"K"`
	Time   ns         `json:"T"`
	Server string     `json:"S"`
	Host   string     `json:"M"`
	Pid    int        `json:"P"`

	Tx   txid              `json:"I"`
	Span spanid            `json:"R"`
	Data map[string]string `json:"D"`
}

type span struct {
	Tx     txid
	Span   spanid
	Parent spanid
	Out    callOut
	In     callIn
	Spans  map[spanid]span
}

type callOut struct {
	Start, End ns
	StartAttr  map[string]string
	EndAttr    map[string]string
}

type callIn struct {
	Start, End ns
	StartAttr  map[string]string
	EndAttr    map[string]string
	// ignoring extra "log" messages for now
}

type recordType int

func (rt recordType) String() string {
	switch rt {
	default:
		return "unknown"
	case rtConnect:
		return "connect"
	case rtSpanStart:
		return "span start"
	case rtSpanEnd:
		return "span end"
	case rtCallStart:
		return "call start"
	case rtCallEnd:
		return "call end"
	}
}

const (
	rtConnect   recordType = 1
	rtSpanStart recordType = 3
	rtSpanEnd   recordType = 4
	rtCallStart recordType = 5
	rtCallEnd   recordType = 6
)

func main() {
	var wg sync.WaitGroup

	jsonFile := "/dev/stdin"
	flag.StringVar(&jsonFile, "json", jsonFile, "JSON-format input file location")
	addr := ":8080"
	flag.StringVar(&addr, "http", addr, "HTTP listen address")

	flag.Parse()

	httpErr := make(chan error, 1)
	go func() {
		err := http.ListenAndServe(addr, nil)
		httpErr <- err
	}()

	file, err := os.Open(jsonFile)
	if err != nil {
		log.Fatalf("file-err=%q", err)
	}
	defer func(c io.Closer) {
		err := c.Close()
		if err != nil {
			log.Printf("close-err=%q", err)
		}
	}(file)

	ch := make(chan record, 100)
	inShards := make([]chan<- record, 0, 1<<4)
	records := make(chan map[txid][]record, cap(inShards))
	for i := 0; i < cap(inShards); i++ {
		ch := make(chan record, 100)
		inShards = append(inShards, ch)
		wg.Add(1)
		go func(ch <-chan record) {
			defer wg.Done()
			m := make(map[txid][]record)
			defer func() {
				records <- m
			}()
			for i := 0; ; i++ {
				v, ok := <-ch
				if !ok {
					break
				}
				if v.Tx%txid(cap(inShards)) == 1 {
					m[v.Tx] = append(m[v.Tx], v)
				}
			}
		}(ch)
	}

	go read(inShards, ch)

	log.Printf("reading")
	err = scan(ch, file)
	if err != nil {
		log.Fatalf("err=%q", err)
	}
	close(ch)
	wg.Wait()
	close(records)
	log.Printf("done")

	parentThumbs := make(map[string]int)
	parentThumbprints := make(map[string]int)
	byThumbprint := make(map[string][]span)

	for m := range records {
		for _, tx := range m {
			spans := make(map[spanid]span)
			for _, record := range tx {
				sp, ok := spans[record.Span]
				if !ok {
					sp.Tx = record.Tx
					sp.Span = record.Span
					sp.Spans = make(map[spanid]span)
					if last := strings.LastIndex(string(record.Span), "."); last != -1 {
						sp.Parent = record.Span[:last]
					}
				}
				switch record.Kind {
				case rtSpanStart:
					sp.In.Start = record.Time
					sp.In.StartAttr = record.Data
				case rtSpanEnd:
					sp.In.End = record.Time
					sp.In.EndAttr = record.Data
				case rtCallStart:
					sp.Out.Start = record.Time
					sp.Out.StartAttr = record.Data
				case rtCallEnd:
					sp.Out.End = record.Time
					sp.Out.EndAttr = record.Data
				}
				spans[record.Span] = sp
			}

			var root span
			for _, span := range spans {
				if span.Span == span.Parent {
					root = span
					continue
				}
				if parent, ok := spans[span.Parent]; ok {
					parent.Spans[span.Span] = span
				}
			}

			tp := thumbprint(root)
			parentThumbprints[tp]++
			byThumbprint[tp] = append(byThumbprint[tp], root)

			if root.In.StartAttr != nil {
				parentThumbs[root.In.StartAttr["thumb"]]++
			}
		}
	}

	prefix := "/thumbprint/"
	http.Handle(prefix, http.StripPrefix(strings.TrimRight(prefix, "/"), thumbHandler(byThumbprint)))

	log.Printf("ready")

	err = <-httpErr
	if err != nil {
		log.Fatalf("http-err=%q", err)
	}
}

type thumbHandler map[string][]span

var thumbIndex = template.Must(template.New("").Parse(`
<html>
<head>
</head>
<body>

{{range .}}
<p>{{.Count}} <a href={{.Thumb}}>{{.Thumb}}</a></p>
{{end}}

</body>
</html>
`[1:]))

func (th thumbHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if _, ok := th[strings.TrimLeft(r.URL.String(), "/")]; ok {
		th.serveThumb(w, r)
		return
	}

	list := make(thumbCountSort, 0, len(th))
	for thumb, elts := range th {
		list = append(list, thumbCount{Thumb: thumb, Count: len(elts)})
	}
	sort.Sort(sort.Reverse(list))

	w.Header().Set("Content-Type", "text/html")
	w.WriteHeader(http.StatusOK)

	err := thumbIndex.Execute(w, list)
	if err != nil {
		log.Printf("err: %v", err)
	}
}

var thumbSingle = template.Must(template.New("").Parse(`
<html>
<head>
  <title></title>
  <script src="http://trifacta.github.io/vega/lib/d3.v3.min.js"></script>
  <script src="http://trifacta.github.io/vega/vega.js"></script>
</head>
<body>

<div id="vis"></div>

{{range .Spans}}
<p>{{.}}</a></p>
{{end}}

</body>
<script type="text/javascript">
function parse(spec) {
  vg.parse.spec(spec, function(chart) {
    chart({el:"#vis"}).update();
  });
}
parse({{.Spec}});
</script>
</html>
`[1:]))

func (th thumbHandler) serveThumb(w http.ResponseWriter, r *http.Request) {
	list := th[strings.TrimLeft(r.URL.String(), "/")]
	if len(list) == 0 {
		http.NotFound(w, r)
		return
	}

	bins := make(map[int]int)
	for _, elt := range list {
		if start, end := elt.In.Start, elt.In.End; start != 0 && end != 0 {
			delta := end - start
			if delta > 0 {
				bins[int(math.Log10(float64(delta))*10)]++
			}
		}
	}

	type vegadata struct {
		Name   string        `json:"name"`
		Values []interface{} `json:"values"`
	}
	type contents struct {
		Spans []span
		Spec  struct {
			Width   int `json:"width"`
			Height  int `json:"height"`
			Padding struct {
				Top    int `json:"top"`
				Bottom int `json:"bottom"`
				Left   int `json:"left"`
				Right  int `json:"right"`
			} `json:"padding"`
			Data   []vegadata               `json:"data"`
			Scales []map[string]interface{} `json:"scales"`
			Axes   []map[string]interface{} `json:"axes"`
			Marks  []map[string]interface{} `json:"marks"`
		}
	}

	var body contents
	body.Spans = list
	if max := 10; len(list) > max {
		body.Spans = list[:max]
	}
	body.Spec.Width = 400
	body.Spec.Height = 200
	body.Spec.Padding.Top = 10
	body.Spec.Padding.Bottom = 20
	body.Spec.Padding.Left = 60
	body.Spec.Padding.Right = 10
	body.Spec.Scales = []map[string]interface{}{
		{
			"name": "x", "type": "log", "range": "width", "nice": true,
			"domain": map[string]string{"data": "table", "field": "data.x"},
		},
		{
			"name": "y", "type": "log", "range": "height", "nice": true,
			"domain":    map[string]string{"data": "table", "field": "data.y"},
			"domainMin": 0.1,
		},
	}
	body.Spec.Axes = []map[string]interface{}{
		{"type": "x", "scale": "x"},
		{"type": "y", "scale": "y"},
	}
	body.Spec.Marks = []map[string]interface{}{
		{
			"type": "rect",
			"from": map[string]string{"data": "table"},
			"properties": map[string]interface{}{
				"enter": map[string]interface{}{
					"x":  map[string]interface{}{"scale": "x", "field": "data.x"},
					"x2": map[string]interface{}{"scale": "x", "field": "data.x2"},
					"y":  map[string]interface{}{"scale": "y", "field": "data.y"},
					"y2": map[string]interface{}{"scale": "y", "value": 0.1},
				},
				"update": map[string]interface{}{
					"fill": map[string]interface{}{"value": "steelblue"},
				},
				"hover": map[string]interface{}{
					"fill": map[string]interface{}{"value": "red"},
				},
			},
		},
	}

	data := []interface{}{}
	ordered := make([]int, 0, len(bins))
	for bin := range bins {
		ordered = append(ordered, bin)
	}
	sort.Ints(ordered)
	for _, bin := range ordered {
		count := bins[bin]
		type xy struct {
			X  int `json:"x"`
			X2 int `json:"x2"`
			Y  int `json:"y"`
		}
		x := int(math.Exp((float64(bin)+0.05)/10*math.Log(10))) / int(time.Microsecond)
		x2 := int(math.Exp((float64(bin)+0.95)/10*math.Log(10))) / int(time.Microsecond)
		data = append(data, xy{X: x, X2: x2, Y: count})
	}
	body.Spec.Data = []vegadata{
		{
			Name:   "table",
			Values: data,
		},
	}

	err := thumbSingle.Execute(w, body)
	if err != nil {
		log.Printf("err: %v", err)
	}
}

type thumbCount struct {
	Thumb string
	Count int
}

type thumbCountSort []thumbCount

func (t thumbCountSort) Len() int      { return len(t) }
func (t thumbCountSort) Swap(i, j int) { t[i], t[j] = t[j], t[i] }
func (t thumbCountSort) Less(i, j int) bool {
	if ti, tj := t[i].Count, t[j].Count; ti != tj {
		return ti < tj
	}
	if ti, tj := t[i].Thumb, t[j].Thumb; ti != tj {
		return ti < tj
	}
	return false
}

func thumbprint(sp span) string {
	childThumbs := make([]string, 0, len(sp.Spans))
	for _, child := range sp.Spans {
		childThumbs = append(childThumbs, thumbprint(child))
	}
	sort.Strings(childThumbs)

	var w bytes.Buffer
	fmt.Fprintf(&w, "%q\n", outThumb(sp))
	fmt.Fprintf(&w, "%q\n", inThumb(sp))
	for _, t := range childThumbs {
		fmt.Fprintf(&w, "%s\n", t)
	}
	return fmt.Sprintf("%x", sha256.Sum256(w.Bytes()))
}

func inThumb(sp span) string {
	if sp.In.StartAttr == nil {
		return ""
	}
	return sp.In.StartAttr["thumb"]
}

func outThumb(sp span) string {
	if sp.Out.StartAttr == nil {
		return ""
	}
	return sp.Out.StartAttr["thumb"]
}

//

func scan(ch chan<- record, r io.Reader) error {
	sc := bufio.NewScanner(r)
	for i := 0; sc.Scan(); i++ {
		var v record

		err := json.Unmarshal(sc.Bytes(), &v)
		if err != nil {
			return err
		}

		ch <- v
	}

	err := sc.Err()
	if err != nil {
		return err
	}

	return nil
}

func read(out []chan<- record, in <-chan record) {
	defer func() {
		for _, ch := range out {
			close(ch)
		}
	}()
	for v := range in {
		out[int(v.Tx%txid(len(out)))] <- v
	}
}
