package main

import (
	"crypto/rand"
	"flag"
	"log"
	"math/big"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"golang.org/x/net/context"

	"code.justin.tv/common/chitin"
	_ "code.justin.tv/common/golibs/bininfo"
	"code.justin.tv/ear/kinesis/combiner"
	"code.justin.tv/release/trace/internal/pipeline"
	"code.justin.tv/release/trace/internal/sink"
	"code.justin.tv/release/trace/pbmsg"
	"code.justin.tv/release/trace/scanproto"
)

const (
	defaultTxTimeout = 65 * time.Second
	kinesisRegion    = "us-west-2"
	MAX_RECORD_SIZE  = 1000000 // 1 MB
)

func main() {
	var (
		inFile  string
		outFile string

		unused string

		inKinesis  bool
		outKinesis string

		readWorkers  int
		writeWorkers int

		txTimeout time.Duration
	)

	flag.StringVar(&inFile, "in-file", "events.pb", "input file of encoded EventSets")
	flag.StringVar(&outFile, "out-file", "transactions.pb", "output file to write encoded Transactions")

	flag.StringVar(&unused, "zk", "", "Unused. Reserved for backwards compatibility.")
	flag.StringVar(&unused, "kafka", "", "Unused. Reserved for backwards compatibility.")
	flag.StringVar(&unused, "kafka-group", "", "Unused. Reserved for backwards compatibility.")
	flag.StringVar(&unused, "in-topic", "", "Unused. Reserved for backwards compatibility.")
	flag.StringVar(&unused, "out-topic", "", "Unused. Reserved for backwards compatibility.")

	flag.BoolVar(&inKinesis, "in-kinesis", false, "set to true to run under the KCL MultiLangDaemon, accepting messages from stdin. Overrides in-file and in-topic.")
	flag.StringVar(&outKinesis, "out-kinesis-stream", "", "kinesis stream to publish transactions to")

	flag.IntVar(&readWorkers, "nread", 1, "number of read workers")
	flag.IntVar(&writeWorkers, "nwrite", 4, "number of write workers")
	flag.DurationVar(&txTimeout, "tx-timeout", defaultTxTimeout, "maximum allowed duration for a transaction")

	flag.Parse()

	var (
		source scanproto.EventSource
		dest   sink.Sink
		err    error
	)

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

	switch {
	case inKinesis:
		source = scanproto.NewKinesisEventSource()
	default:
		source = scanproto.NewFileEventSource(inFile, readWorkers)
	}

	switch {
	case outKinesis != "":
		cfg := sink.NewKinesisConfig(outKinesis)
		cfg.OnError = pipeline.LogErrors(cfg.Logger)
		dest, err = sink.Kinesis(context.Background(), cfg)
		if err != nil {
			log.Fatalf("unable to set up kinesis sink, err=%q", err)
		}
	default:
		log.Printf("writing to file location=%q", outFile)
		dest = newFileTransactionSink(outFile, writeWorkers).consume
	}

	evCh := source.Events()
	txCh := make(chan *pbmsg.Transaction, writeBuffer)
	blobCh := make(chan *combiner.Blob, writeBuffer)
	blobber := combiner.NewRecordCombiner(1, MAX_RECORD_SIZE, 1*time.Second)

	// Waitgroup to ensure clean shutdown of the pipeline, flushing
	// all records before exiting.
	wg := &sync.WaitGroup{}

	wg.Add(1)
	go func() {
		log.Println("source.run start")
		source.Run()
		go func() {
			for e := range source.Errors() {
				log.Printf("ERROR!! err=%q", e)
				source.Stop()
				return
			}
		}()
		wg.Done()
		log.Println("source.run end")
	}()

	// If we're running in kinesis, the MultiLangDaemon is in charge
	// of how we run, and it shuts us down through a special message
	// over stdin, not through signals. Signals in that case are
	// unexpected. Otherwise, though, SIGINT and SIGTERM should
	// trigger a graceful shutdown.
	if !inKinesis {
		log.Println("turning on signal handling")
		sigch := make(chan os.Signal, 1)
		signal.Notify(sigch, syscall.SIGINT, syscall.SIGTERM)
		go func() {
			for _ = range sigch {
				log.Println("source.stop start")
				wg.Add(1)
				source.Stop()
				wg.Done()
				log.Println("source.stop end")
			}
		}()
	}

	wg.Add(1)
	go func() {
		log.Println("processEventStream start")
		processEventStream(evCh, txCh, txTimeout)
		wg.Done()
		log.Println("processEventStream end")
	}()

	wg.Add(1)
	go func() {
		log.Println("combineTransactions start")
		combineTransactions(txCh, blobber)
		wg.Done()
		log.Println("combineTransactions end")
	}()

	/*
		wg.Add(1)
		go func() {
			log.Println("blobber.errors start")
			for e := range blobber.Errors() {
				log.Printf("Blobber error!! err=%q", e)
			}
			wg.Done()
			log.Println("blobber.errors end")
		}()
	*/

	wg.Add(1)
	go func() {
		log.Println("blob.Key setter start")

		var buf [16]byte
		key := new(big.Int)
		for b := range blobber.Blobs() {
			// Since we're only using a single virtual shard we need to
			// overwrite b.Key, as it'll be a constant ("1"). Generate a random
			// 128-bit integer key.
			rand.Read(buf[:])
			key.SetBytes(buf[:])
			b.Key = key.String()
			blobCh <- b
		}
		close(blobCh)
		wg.Done()
		log.Println("blob.Key setter end")
	}()

	wg.Add(1)
	go func() {
		log.Println("sink.consume start")
		dest(blobCh)
		wg.Done()
		log.Println("sink.consume end")
	}()

	wg.Wait()
}
