package twirpserviceslohook

import (
	"context"
	"strconv"
	"time"

	"github.com/twitchtv/twirp"
)

type contextKey int

const (
	requestTimeKey contextKey = iota
)

// MethodName is a helper type to better document the map TwirpSLOTracker.SLO
type MethodName string

// ServiceName is a helper type to better document the map TwirpSLOTracker.SLO
type ServiceName string

func defaultFailedStatusCodes(code int) bool {
	return code < 200 || code >= 500
}

var defaultStatTracker = noopStatTracker{}

type noopStatTracker struct{}

func (noopStatTracker) Passed(serviceName string, methodName string) {}
func (noopStatTracker) Failed(serviceName string, methodName string) {}

var _ StatTracker = defaultStatTracker

// StatTracker is a sink that can track passing and failing SLO stats
type StatTracker interface {
	Passed(serviceName string, methodName string)
	Failed(serviceName string, methodName string)
}

// TwirpSLOTracker can create twirp service hooks that track SLO metrics.  A method fails SLO if it is a bad status
// code or takes too long.
type TwirpSLOTracker struct {
	// Map of service name -> method name -> maximum time a method can take while still meeting SLO
	SLO map[ServiceName]map[MethodName]time.Duration
	// SLO to use if none is present in the SLO map
	DefaultSLO time.Duration
	// Where to send stats.  By default, stats are dropped
	StatTracker StatTracker
	// Which status codes count as failures.  By default, all non 2xx, 3xx, 4xx are failures
	FailingStatusCodes func(code int) bool `json:"-"`
}

// MergeSLO merges multiple SLO definitions
func MergeSLO(slos ...map[ServiceName]map[MethodName]time.Duration) map[ServiceName]map[MethodName]time.Duration {
	ret := map[ServiceName]map[MethodName]time.Duration{}
	for _, slo := range slos {
		for k, v := range slo {
			if ret[k] == nil {
				ret[k] = map[MethodName]time.Duration{}
			}
			for mn, t := range v {
				ret[k][mn] = t
			}
		}
	}
	return ret
}

func (t *TwirpSLOTracker) statTracker() StatTracker {
	if t.StatTracker == nil {
		return defaultStatTracker
	}
	return t.StatTracker
}

func (t *TwirpSLOTracker) defaultSLO() time.Duration {
	if t.DefaultSLO == 0 {
		return time.Millisecond * 150
	}
	return t.DefaultSLO
}

func (t *TwirpSLOTracker) failingStatusCodes() func(code int) bool {
	if t.FailingStatusCodes == nil {
		return defaultFailedStatusCodes
	}
	return t.FailingStatusCodes
}

func (t *TwirpSLOTracker) passes(serviceName ServiceName, methodName MethodName, reqTime time.Duration) bool {
	timing := t.SLO[serviceName][methodName]
	if timing == 0 {
		return reqTime <= t.defaultSLO()
	}
	return reqTime <= timing
}

// ServerHook creates a twirp server hook for this stat tracker
func (t *TwirpSLOTracker) ServerHook() *twirp.ServerHooks {
	return &twirp.ServerHooks{
		RequestReceived: func(ctx context.Context) (context.Context, error) {
			ctx = context.WithValue(ctx, requestTimeKey, time.Now())
			return ctx, nil
		},
		// ResponseSent is called when all bytes of a response (including an error
		// response) have been written. Because the ResponseSent hook is terminal, it
		// does not return a context.
		ResponseSent: func(ctx context.Context) {
			serviceName, _ := twirp.ServiceName(ctx)
			methodName, _ := twirp.MethodName(ctx)

			if codeFromCtx, ok := twirp.StatusCode(ctx); ok {
				if codeAsInt, okAsInt := strconv.Atoi(codeFromCtx); okAsInt == nil {
					if t.failingStatusCodes()(codeAsInt) {
						t.statTracker().Failed(serviceName, methodName)
						return
					}
				}
			}
			if recvTime, ok := ctx.Value(requestTimeKey).(time.Time); ok {
				reqDuration := time.Since(recvTime)
				if t.passes(ServiceName(serviceName), MethodName(methodName), reqDuration) {
					t.statTracker().Passed(serviceName, methodName)
				} else {
					t.statTracker().Failed(serviceName, methodName)
				}
			}
		},
	}
}
