package main

import (
	"golang.org/x/net/context"

	"code.justin.tv/release/trace/api"
)

const siphonJobCount = 10

type Siphon interface {
	PullTransactions(ctx context.Context) error
	CloseFireHose()
}

/// NewFirehoseSiphon creates a new firehoseSiphon object and returns it
func NewFirehoseSiphon(firehoseClient api.TraceClient, stats StatsWriter, serviceName string, dbList map[string]string) Siphon {
	siphon := new(firehoseSiphon)
	siphon.client = firehoseClient
	siphon.stats = stats
	siphon.serviceName = serviceName
	siphon.databaseList = dbList
	siphon.funnel = make(chan *api.Transaction)

	for i := 0; i < siphonJobCount; i++ {
		go siphon.funnelJob()
	}

	return siphon
}

func (siphon *firehoseSiphon) CloseFireHose() {
	close(siphon.funnel)
}

type firehoseSiphon struct {
	client       api.TraceClient
	stats        StatsWriter
	serviceName  string
	databaseList map[string]string
	funnel       chan *api.Transaction
}

/// PullTransactions starts up a firehose query that will run as long as possible.
/// All results from that query will be sent to the funnel channel to be worked by
/// the funnel jobs.
func (siphon *firehoseSiphon) PullTransactions(ctx context.Context) error {
	var err error

	//Start a firehose query to pull any and all transactions that include the specified watch-service
	firehose, err := siphon.client.Firehose(ctx, &api.FirehoseRequest{
		Sampling: 1.0,
		Query: &api.Query{
			Comparisons: []*api.Comparison{
				&api.Comparison{
					ServiceName: &api.StringComparison{Value: siphon.serviceName},
				},
			},
		},
	})
	if err != nil {
		return err
	}

	//Retrieve & pipe into the funnel as quickly as possible
	for {
		tx, err := firehose.Recv()
		if err != nil {
			return err
		}
		siphon.funnel <- tx
	}
}

/// funnelJob is the driver code for a set of goroutines that work the funnel,
/// taking transactions from the firehose client and sending them to stats
func (siphon *firehoseSiphon) funnelJob() {
	//Read from the funnel channel & pipe into stats
	for tx := range siphon.funnel {
		siphon.stats.AddTransaction(tx)

		siphon.seekQueries(tx.GetRoot())
	}
}

/// seekQueries accepts a trace call & works the entire call tree,
/// piping our desired sql queries into stats.  What we're looking
/// for is SQL subcalls who are children of a call from our watch-service.
/// So we are going to work the three looking for calls with the correct
/// service name, and then send all SQL subcalls to stats.
func (siphon *firehoseSiphon) seekQueries(call *api.Call) {
	if call == nil {
		return
	}

	// This call belongs to our watch-service, send all SQL subcalls to stats
	checkSQL := (call.GetSvc().Name == siphon.serviceName)

	for _, subCall := range call.GetSubcalls() {
		siphon.seekQueries(subCall)

		if checkSQL {
			sql := subCall.GetParams().GetSql()

			if sql == nil {
				continue
			}

			//Only actually send to stats if the query is against one of our desired db's
			_, shouldWatch := siphon.databaseList[sql.Dbname]
			if !shouldWatch {
				continue
			}

			siphon.stats.AddQuery(subCall)
		}
	}
}
