package twitchserver

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	_ "net/http/pprof" // importing twitchhttp implies intent to host pprof metrics, so we automatically enable it.
	"os"
	"os/signal"
	"runtime"
	"sync"
	"syscall"
	"time"

	"golang.org/x/sync/errgroup"

	"github.com/cactus/go-statsd-client/statsd"

	"code.justin.tv/common/chitin"
	"code.justin.tv/common/config"
	"code.justin.tv/common/gometrics"
	"code.justin.tv/feeds/ctxlog"
	"code.justin.tv/foundation/xray"

	goji "goji.io"
	"goji.io/middleware"
	"goji.io/pat"
)

var (
	stdSignals = []os.Signal{syscall.SIGINT}

	downMu         sync.RWMutex // Protects isShuttingDown.
	isShuttingDown bool         // Used by the healthcheck endpoint to signal server shut down to load balancers.

	downC = make(chan struct{}, 1) // Used by Shutdown and ShutdownNow.

	downTimeMu  sync.RWMutex  // Protects downTimeout.
	downTimeout time.Duration // The amount of time to wait for a server to gracefully shut down.

	healthCheckMu   sync.Mutex
	healthCheckHook func() bool
)

// Server allows simple definition of an HTTP service which meets Twitch requirements for production readiness.
func init() {
	config.Register(map[string]string{
		"bind-address":       ":8000",
		"debug-bind-address": ":6000",
	})
	rand.Seed(time.Now().UnixNano())
}

// ServerConfig configures a twitchhttp server.
// All configuration is optional and will fall back to reasonable defaults
type ServerConfig struct {
	Addr      string         // defaults to config.Resolve("bind-address")
	DebugAddr string         // defaults to config.Resolve("debug-address")
	Statter   statsd.Statter // defaults to config.Statsd()

	// Timeout configuration. These values default to 0 (ignored by graceful.Server)
	// Documentation for these values can be found here: https://golang.org/pkg/net/http/#Server
	ReadTimeout       time.Duration
	ReadHeaderTimeout time.Duration
	WriteTimeout      time.Duration
	IdleTimeout       time.Duration

	// GracefulTimeout is a duration which a server will attempt to shut down gracefully before shutting down abruptly. Defaults to 10s.
	GracefulTimeout time.Duration
	// ctxlog will append logging dimensions to the context using this key. A logger may
	// want to log those context values.
	DimensionKey interface{}
	/* ElevateKey is a key for request contexts. The value will be a boolean which dictates whether or not logs should be elevated to a
	higher log level. Logging implementations may inspect the context for this key and respond accordingly. It must be set to the same
	value as twitchclient. */
	ElevateKey interface{}
	// Stop is a channel that initiates shutdown of the server when a value is received.
	Stop <-chan struct{}

	// Hooks that are
	PreShutdownHooks []func()
	PostStartupHooks []func()
	HealthCheckHook  func() bool
}

// NewConfig creates a new configuration using the config values bind-address and debug-bind-address.
func NewConfig() *ServerConfig {
	return fillConfig(&ServerConfig{})
}

func fillConfig(conf *ServerConfig) *ServerConfig {
	if conf == nil {
		conf = NewConfig()
	}
	if conf.Addr == "" {
		conf.Addr = config.Resolve("bind-address")
	}
	if conf.DebugAddr == "" {
		conf.DebugAddr = config.Resolve("debug-bind-address")
	}
	if conf.Statter == nil {
		conf.Statter = config.Statsd()
	}
	if conf.GracefulTimeout == 0 {
		conf.GracefulTimeout = 10 * time.Second
	}
	if conf.DimensionKey == nil {
		conf.DimensionKey = new(int)
	}
	if conf.HealthCheckHook == nil {
		conf.HealthCheckHook = func() bool { return true }
	}
	return conf
}

// NewServer allocates and returns a Server. The /debug/running endpoint is registered for load balancer health checks.
func NewServer() *goji.Mux {
	debugRouter := goji.SubMux()
	debugRouter.HandleFunc(pat.Get("/running"), healthcheck)

	router := goji.NewMux()
	router.Use(RollbarMiddleware)
	router.Use(SkipMiddleware(debugRouter))
	router.Handle(pat.New("/debug/*"), debugRouter)

	return router
}

// Generate a base handler which wraps chitin, xray, and ctxlog
func baseHandler(h http.Handler, ctxLog *ctxlog.Ctxlog) (http.Handler, error) {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := chitin.Context(w, r)
		r = r.WithContext(ctx)
		h.ServeHTTP(w, r)
	})

	err := chitin.ExperimentalTraceProcessOptIn()

	if err != nil {
		return nil, err
	}

	ctxlogHandler := &ctxlog.CtxHandler{
		Next:      handler,
		Ctxlog:    ctxLog,
		Logger:    &defaultLogger{},
		IDFetcher: xray.TraceID,
	}

	xrayHandler := xray.Handler(ctxlogHandler)
	chitinHandler := chitin.Handler(xrayHandler)

	return chitinHandler, nil
}

// Serve accepts connections to the listener and dispatches HTTP requests, as well as starting a debug server for pprof debugging.
// Go runtime metrics begin emitting to statsd.
func Serve(ln net.Listener, h http.Handler, serverConfig *ServerConfig) error {
	serverConfig = fillConfig(serverConfig)

	downMu.Lock()
	// Allows for shutdown and subsequent startup for servers.
	// This is NOT the intended usage and this may cause issues for you.
	// TODO: refactor out global variables.
	isShuttingDown = false
	downMu.Unlock()

	healthCheckMu.Lock()
	healthCheckHook = serverConfig.HealthCheckHook
	healthCheckMu.Unlock()

	if serverConfig.DebugAddr != "" {
		go func() {
			log.Printf("debug listening on %v", serverConfig.DebugAddr)
			log.Print(http.ListenAndServe(serverConfig.DebugAddr, nil))
		}()
	}

	gometrics.Monitor(serverConfig.Statter, time.Second*5)

	ctxLog := &ctxlog.Ctxlog{
		CtxDims:       &ctxDimensions{serverConfig.DimensionKey},
		ElevateKey:    serverConfig.ElevateKey,
		StartingIndex: rand.Int63(),
	}
	handler, err := baseHandler(h, ctxLog)
	if err != nil {
		return err
	}

	log.Printf("application listening on %v", serverConfig.Addr)
	server := &http.Server{
		Addr:              serverConfig.Addr,
		Handler:           handler,
		ReadTimeout:       serverConfig.ReadTimeout,
		ReadHeaderTimeout: serverConfig.ReadHeaderTimeout,
		WriteTimeout:      serverConfig.WriteTimeout,
		IdleTimeout:       serverConfig.IdleTimeout,
	}

	runPostStartupHooks(serverConfig)

	var g errgroup.Group
	g.Go(waitForShutdown(server, serverConfig, downC))

	err = server.Serve(ln)
	if err != http.ErrServerClosed {
		return err
	}

	err = g.Wait()
	log.Println("Server is shut down...")

	return err
}

// ListenAndServe starts the application server, as well as a debug server for pprof debugging.
// Go runtime metrics begin emitting to statsd.
func ListenAndServe(h http.Handler, serverConfig *ServerConfig) error {
	serverConfig = fillConfig(serverConfig)
	ln, err := net.Listen("tcp", serverConfig.Addr)
	if err != nil {
		return err
	}
	return Serve(tcpKeepAliveListener{ln.(*net.TCPListener)}, h, serverConfig)
}

// AddDefaultSignalHandlers will handle OS signals in an opinionated but sane manner. At the moment, that means
//
// SIGUSR2 will gracefully shut down
// SIGINT will gracefully shut down
// SIGTERM will gracefully shut down
func AddDefaultSignalHandlers() {
	stdSignals = append(stdSignals, syscall.SIGTERM, syscall.SIGUSR2)
}

// AddShutdownSignal will add listeners to OS signals.
// When any of those signals are received, the http server will shut down gracefully.
func AddShutdownSignal(sig ...os.Signal) {
	stdSignals = append(stdSignals, sig...)
}

// ClearSignalHandlers removes any existing signal handling, including SIGINT, preventing shutdown
// of the server.
func ClearSignalHandlers() {
	stdSignals = []os.Signal{}
}

// Shutdown will gracefully shut down the HTTP server.
// No future requests will be accepted, but in-progress requests will complete.
func Shutdown(t time.Duration) {
	setDownTimeout(t)
	goDown()
}

// ShutdownNow initiates server shutdown with the configured graceful shutdown timeout.
// The timeout is set very short so that we immediately close all connections forcefully.
func ShutdownNow() {
	goDown()
}

// RegisterHealthCheckHook sets a global function that is executed once per
// health check request.
func RegisterHealthCheckHook(check func() bool) {
	healthCheckMu.Lock()
	healthCheckHook = check
	healthCheckMu.Unlock()
}

func setDownTimeout(t time.Duration) {
	downTimeMu.Lock()
	downTimeout = t
	downTimeMu.Unlock()
}

func goDown() {
	select {
	case downC <- struct{}{}:
	default:
	}
}

// RegisterPreShutdownHook allows consumers to register a hook to run at PreShutdown
//
// The PreShutdown event fires after the server receives a SIGINT, SIGUSR2 or call to Shutdown()
// but before the server has started severing any existing connections.
// (This event may be used to notify load balancers that the server should be taken out of rotation)
func (c *ServerConfig) RegisterPreShutdownHook(hook func()) {
	c.PreShutdownHooks = append(c.PreShutdownHooks, hook)
}

// RegisterPostStartupHook allows consumers to register a hook to run at PostStartup
//
// The PostStartup event fires once the server has started up and is ready to receive connections.
// Once the event fires it is safe to start sending requests to the server.
// (This event may be used to notify load balancers that the server can be placed into rotation)
func (c *ServerConfig) RegisterPostStartupHook(hook func()) {
	c.PostStartupHooks = append(c.PostStartupHooks, hook)
}

// ClearHooks resets all the hooks that have been registered
func (c *ServerConfig) ClearHooks() {
	c.PreShutdownHooks = nil
	c.PostStartupHooks = nil
}

func runPreShutdownHooks(c *ServerConfig) {
	log.Printf("About to run the PreShutdown hooks")
	for _, hook := range c.PreShutdownHooks {
		hook()
	}
}

func runPostStartupHooks(c *ServerConfig) {
	log.Printf("About to run the PostStartup hooks")
	for _, hook := range c.PostStartupHooks {
		hook()
	}
}

// RollbarMiddleware intercepts panics and sends them to rollbar before re-panicing.
func RollbarMiddleware(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if p := recover(); p != nil {
				if config.RollbarErrorLogger() != nil {
					config.RollbarErrorLogger().RequestPanic(r, p)
				} else {
					// This error handling is copied from golang's net/http package
					// https://github.com/golang/go/blob/0ec62565f911575beaf7d047dfe1eae2ae02bf67/src/net/http/server.go#L1468
					// It allows us to log the error while still returning a response to the client
					const size = 64 << 10
					buf := make([]byte, size)
					buf = buf[:runtime.Stack(buf, false)]
					log.Printf("http: panic serving %v: %v\n%s", r.RemoteAddr, p, buf)
				}

				w.WriteHeader(http.StatusInternalServerError)
			}
		}()

		h.ServeHTTP(w, r)
	})
}

// SkipMiddleware takes a handler, and if the request is being routed to that handler
// it will skip all remaining middleware and immediately execute the handler.
//
// The order in which middleware is registered determines which middleware this will skip.
// For more information refer to https://godoc.org/goji.io#Mux.Use
func SkipMiddleware(desiredHandler http.Handler) func(http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if handler := middleware.Handler(r.Context()); handler == desiredHandler {
				handler.ServeHTTP(w, r)
				return
			}

			h.ServeHTTP(w, r)
		})
	}
}

func healthcheck(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Cache-Control", "no-cache, no-store, must-revalidate")

	downMu.RLock()
	down := isShuttingDown
	downMu.RUnlock()

	if down {
		w.WriteHeader(http.StatusServiceUnavailable)
		fmt.Fprint(w, "SHUTTING DOWN")
		return
	}

	healthCheckMu.Lock()
	available := healthCheckHook()
	healthCheckMu.Unlock()

	if !available {
		w.WriteHeader(http.StatusServiceUnavailable)
		fmt.Fprint(w, "unavailable")
		return
	}

	_, err := w.Write([]byte("OK"))
	if err != nil {
		log.Printf("health check failed to respond: %v", err)
	}
}

// tcpKeepAliveListener sets TCP keep-alive timeouts on accepted
// connections. It's used by ListenAndServe and ListenAndServeTLS so
// dead TCP connections (e.g. closing laptop mid-download) eventually
// go away.
// Lifted from goji and net/http.
type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return nil, err
	}
	if err := tc.SetKeepAlive(true); err != nil {
		return nil, err
	}
	if err := tc.SetKeepAlivePeriod(3 * time.Minute); err != nil {
		return nil, err
	}
	return tc, nil
}

func waitForShutdown(s *http.Server, c *ServerConfig, downC chan struct{}) func() error {
	return func() error {
		sig := make(chan os.Signal, 1)

		if len(stdSignals) > 0 {
			signal.Notify(sig, stdSignals...)
		}

		timeout := c.GracefulTimeout

		// Wait for a shutdown signal.
		select {
		case <-sig:
		case <-c.Stop:
		case <-downC:
			downTimeMu.RLock()
			timeout = downTimeout
			downTimeMu.RUnlock()
		}

		log.Println("Shutting down server...")

		downMu.Lock()
		isShuttingDown = true
		downMu.Unlock()

		runPreShutdownHooks(c)

		if timeout == 0 {
			return s.Close()
		}

		ctx, cancel := context.WithTimeout(context.Background(), timeout)
		defer cancel()

		s.SetKeepAlivesEnabled(false)
		if err := s.Shutdown(ctx); err != nil {
			return s.Close()
		}
		return nil
	}
}
