package main

import (
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"runtime/pprof"
	"time"

	"code.justin.tv/common/chitin"
	_ "code.justin.tv/common/golibs/bininfo"
	"code.justin.tv/release/trace/analysis/callgraph"
	"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/persistent"
	"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 (
	defaultTopic     = "trace-transactions"
	defaultFlush     = time.Second * 30
	defaultOutPath   = ""
	defaultURLPrefix = "/"

	txConvertBuffer = 1024

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

func main() {
	var (
		topic      string
		outPath    string
		outDB      string
		urlPrefix  string
		flush      time.Duration
		cpuprofile string
		useAPI     bool
		services   string
		hosts      string
		server     bool

		useKinesis bool

		unused string
	)

	flag.StringVar(&topic, "topic", defaultTopic, "kinesis topic providing transactions")
	flag.DurationVar(&flush, "flush", defaultFlush, "time between flushing transactions into a report")
	flag.StringVar(&outPath, "outdir", defaultOutPath, "Output directory name")
	flag.StringVar(&outDB, "db", "", "Output LevelDB directory name")
	flag.StringVar(&urlPrefix, "urlprefix", defaultURLPrefix, "Prefix for all internal URLs in reports")
	flag.StringVar(&cpuprofile, "cpuprofile", "", "Write a CPU profile to the specified file before exiting")
	flag.BoolVar(&useAPI, "api", false, "Use the Trace API instead of kinesis (allows you to run txreport outside the VPC)")
	flag.StringVar(&services, "services", "", "Comma delimited list of services to filter to")
	flag.StringVar(&hosts, "hosts", "", "Comma delimited list of hosts to filter to")
	flag.BoolVar(&server, "server", false, "Runs an HTTP server at the end of the run for easy local use")
	flag.BoolVar(&useKinesis, "kinesis", true, "Use Kinesis as the transaction source")

	flag.StringVar(&unused, "kafka", "", "Unused. Reserved for backwards compatibility.")

	flag.Parse()

	var err error

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

	if cpuprofile != "" {
		cpu, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatalf("cpuprofile err=%q", err)
		}
		defer func(c io.Closer) {
			err := c.Close()
			if err != nil {
				log.Printf("cpuprofile close err=%q", err)
			}
		}(cpu)
		log.Println("starting cpuprofile")
		err = pprof.StartCPUProfile(cpu)
		if err != nil {
			log.Fatalf("cpuprofile err=%q", err)
		}
		defer pprof.StopCPUProfile()
	}

	var source scanproto.TransactionSource
	switch {
	case useAPI:
		source = scanproto.NewSimpleApiTransactionSource("trace-api.prod.us-west2.justin.tv:11143", services, hosts)
	case useKinesis:
		conf := aws.NewConfig().WithRegion(kinesisRegion)
		sess := session.New(conf)
		client := kinesis.New(sess)
		source = scanproto.NewKinesisTransactionSource(context.Background(), client, topic)
	}
	limitedSource := scanproto.LimitedTransactionSource(source, flush)

	go func() {
		log.Println("start reading input")
		limitedSource.Run()
		log.Println("done reading input")
	}()
	go func() {
		for e := range limitedSource.Errors() {
			log.Fatalf("source err=%v", e)
		}
	}()
	input := limitedSource.Transactions()

	// hold all transactions in memory
	txs := tx.NewInMemoryTransactionSet()

	var db *persistent.DB
	if outDB != "" {
		db, err = persistent.OpenDB(outDB)
		if err != nil {
			log.Fatalf("db err=%q", err)
		}
		defer func() {
			log.Printf("db compacting")
			err := db.CompactAll()
			if err != nil {
				log.Printf("db compact err=%q", err)
			}
			err = db.Close()
			if err != nil {
				log.Printf("db close err=%q", err)
			}
			log.Printf("db closed")
		}()
	}

	// Convert 'raw' pbmsg.Transactions into tx.Transactions because
	// our existing analysis code expects the latter.
	log.Println("start converting input")
	i := 0
	for t := range input {
		i += 1
		txs.Add(&tx.Transaction{
			Root: callgraph.TransactionToCallgraph(t),
			ID:   tx.IDForTx(t),

			Tx: cleanedAPITransaction(t),
		})
	}
	log.Printf("done converting input - read %d txs", i)

	log.Println("start creating report")
	report := NewReport(outPath, urlPrefix)
	report.DB = db
	report.Populate(txs)
	log.Println("done creating report")
	log.Println("start dumping html")
	err = report.DumpHTML()
	if err != nil {
		log.Fatalf("html write err=%q", err)
	}

	log.Println("finished writing")

	if server {
		fs := http.FileServer(http.Dir(outPath))
		http.Handle("/", fs)
		log.Println("serving report on port 8080")
		http.ListenAndServe(":8080", nil)
	}
}

// cleanAPITransaction copies the provided Transaction and removes any
// potentially-sensitive information.
func cleanedAPITransaction(tx *api.Transaction) *api.Transaction {
	ourTx := proto.Clone(tx).(*api.Transaction)
	cleanAPITransaction(ourTx)
	return ourTx
}

// cleanAPITransaction removes any potentially-sensitive information from the
// provided Transaction. It modifies the Transaction in place.
//
// Callers can use proto.Clone to deep-copy the Transaction if they have a
// read-only value.
func cleanAPITransaction(tx *api.Transaction) {
	cleanAPICall(tx.GetRoot())
}

// cleanAPICall removes any potentially-sensitive information from the
// provided Call. Specifically, this means removing the HTTP URI path and
// attaching the name of the route.
func cleanAPICall(call *api.Call) {
	for _, c := range call.GetSubcalls() {
		cleanAPICall(c)
	}

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

		svc := call.GetSvc()
		if svc != nil {
			rm := routes.Routes[svc.Name]
			httpParams.Route = rm.LookupTrimSlash(path, httpParams.Method.String())
		}

	}
}
