package contextlogger

import (
	"time"

	log "github.com/Sirupsen/logrus"
	xnetcontext "golang.org/x/net/context"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/stats"
)

// GrpcStatsLogger implements stats.Handler to log gRPC requests.
type GrpcStatsLogger struct {
	Logger *log.Logger
}

// TagRPC adds a key into context with a log.Fields value for accumulating things to log.
func (l *GrpcStatsLogger) TagRPC(ctx xnetcontext.Context, tagInfo *stats.RPCTagInfo) xnetcontext.Context {
	return ContextWithLogFields(ctx)
}

// HandleRPC is called several times for each phase of the RPC.
// We accumulate fields to log from the different phases into the context and
// then we emit the log entry at the end.
func (l *GrpcStatsLogger) HandleRPC(ctx xnetcontext.Context, s stats.RPCStats) {
	logFields, _ := getLogFields(ctx)

	logFields.Lock()
	defer logFields.Unlock()

	switch t := s.(type) {
	case *stats.Begin:
		logFields.Fields["begin_time"] = t.BeginTime
	case *stats.InHeader:
		logFields.Fields["method"] = t.FullMethod
		logFields.Fields["remote_addr"] = t.RemoteAddr
	case *stats.OutPayload:
		logFields.Fields["out_payload_bytes"] = t.WireLength
	case *stats.End:
		logFields.Fields["end_time"] = t.EndTime
		logFields.Fields["error"] = t.Error
		logFields.Fields["duration"] = logFields.Fields["end_time"].(time.Time).Sub(
			logFields.Fields["begin_time"].(time.Time))

		if md, ok := metadata.FromContext(ctx); ok {
			for key, values := range md {
				for _, value := range values {
					if dk, dv, err := metadata.DecodeKeyValue(key, value); err == nil {
						logFields.Fields[dk] = dv
					}
				}
			}
		}

		l.Logger.WithFields(logFields.Fields).Info()
	}
}

// TagConn isn't needed, we only care about requests.
func (l *GrpcStatsLogger) TagConn(ctx xnetcontext.Context, tagInfo *stats.ConnTagInfo) xnetcontext.Context {
	return ctx
}

// HandleConn isn't needed, we only care about requests.
func (l *GrpcStatsLogger) HandleConn(ctx xnetcontext.Context, s stats.ConnStats) {
}
