package main

import (
	"bufio"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"code.justin.tv/common/chitin"
	"code.justin.tv/release/trace/api"
	"code.justin.tv/release/trace/scanproto"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"
)

const (
	// Region of the Kinesis stream to connect to.
	kinesisRegion = "us-west-2"
)

var (
	topic = flag.String("topic", "trace-transactions", "The Kinesis topic")
	t     = flag.Duration("t", 1*time.Second, "Duration of dump")
	debug = flag.String("http", ":11143", "HTTP debug port")
	out   = flag.String("output", "", "Output filename")
)

func main() {
	flag.Parse()

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

	ctx := context.Background()
	ctx, cancel := context.WithTimeout(ctx, *t)
	defer cancel()

	conf := aws.NewConfig().WithRegion(kinesisRegion)
	sess := session.New(conf)
	client := kinesis.New(sess)
	source := scanproto.NewKinesisTransactionSource(ctx, client, *topic)

	var c io.Closer = ioutil.NopCloser(nil)
	var w io.Writer = ioutil.Discard
	if *out != "" {
		f, err := os.OpenFile(*out, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
		if err != nil {
			log.Fatalf("file open err=%q", err)
		}
		c = f
		w = bufio.NewWriter(f)
	}
	defer func() {
		err := c.Close()
		if err != nil {
			log.Printf("close err=%q", err)
		}
	}()

	go func() {
		err := (&http.Server{
			Addr:    *debug,
			Handler: chitin.Handler(http.DefaultServeMux),
		}).ListenAndServe()
		if err != nil {
			log.Printf("http err=%q", err)
		}
	}()

	go func() {
		for err := range source.Errors() {
			log.Fatalf("err=%q", err)
		}
	}()

	go func() {
		source.Run()
	}()

	go func() {
		<-ctx.Done()
		source.Stop()
	}()

	var wg sync.WaitGroup
	var mu sync.Mutex
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			var i, nn, e int64
			for tx := range source.Transactions() {
				i++
				cleanTransaction(tx)
				buf, err := proto.Marshal(&api.TransactionSet{
					Transaction: []*api.Transaction{tx},
				})
				if err != nil {
					e++
					continue
				}

				mu.Lock()
				n, err := w.Write(buf)
				mu.Unlock()

				nn += int64(n)
				if err != nil {
					e++
					continue
				}
			}
			log.Printf("transaction errors=%d count=%d size=%d", e, i, nn)
		}()
	}

	wg.Wait()
}

func cleanTransaction(tx *api.Transaction) {
	cleanCall(tx.GetRoot())
}

func cleanCall(call *api.Call) {
	for _, c := range call.GetSubcalls() {
		cleanCall(c)
	}

	httpParams := call.GetParams().GetHttp()
	if httpParams != nil {
		httpParams.UriPath = ""
	}
}
