package pbmsg

import (
	"io"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

const apiHostPort = "trace-api.prod.us-west2.justin.tv:11143"

// connect connects to the Trace gRPC API, returning a TraceClient. The
// connection to the gRPC API server will be closed when ctx is canceled.
func connect(ctx context.Context) (TraceClient, error) {
	conn, err := grpc.Dial(apiHostPort, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}
	go func() {
		<-ctx.Done()
		conn.Close()
	}()
	return NewTraceClient(conn), nil
}

// Firehose connects to a Trace API server, requesting the stream of data at
// sampling rate, using the provided filters. The connection is closed when ctx
// is canceled.
func Firehose(ctx context.Context, sampling float64, filters ...*Comparison) (Trace_FirehoseClient, error) {
	client, err := connect(ctx)
	if err != nil {
		return nil, err
	}
	return client.Firehose(ctx, &FirehoseRequest{
		Sampling: sampling,
		Query: &Query{
			Comparisons: filters,
		},
	})
}

// ConsumeFirehose streams data from the firehose and runs handle on each
// Transaction received through the Trace_FirehoseClient. If it encounters any
// error other than gRPC's control pseudo-errors (EOF, 'Canceled', and
// 'DeadlineExceeded') then it will exit and return the error.
func ConsumeFirehose(f Trace_FirehoseClient, handle func(*Transaction)) error {
	for {
		tx, err := f.Recv()
		if err == nil {
			handle(tx)
			continue
		}
		if err == io.EOF {
			return nil
		}
		if err != nil {
			switch grpc.Code(err) {
			case codes.Canceled:
				return nil
			case codes.DeadlineExceeded:
				return nil
			default:
				return err
			}
		}
	}
}
