package server

import (
	"compress/gzip"
	"context"
	"io"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/alecthomas/kingpin/v2"
	gokitlog "github.com/go-kit/log"
	"github.com/golang/snappy"
	"github.com/pingcap/log"
	"github.com/prometheus/common/promlog"
	"github.com/prometheus/common/route"
	"github.com/prometheus/common/version"
	"github.com/prometheus/exporter-toolkit/web"
	"github.com/prometheus/exporter-toolkit/web/kingpinflag"
	"github.com/tidbcloud/serverless-metering/pkg/config"
	"github.com/tidbcloud/serverless-metering/pkg/store"
)

type MetricServer struct {
	server      *http.Server
	gokitLogger gokitlog.Logger
	quitCh      chan struct{}
	flags       *web.FlagConfig
}

func NewMetricServer(storage store.Storage, cfg *config.ServerConfig) *MetricServer {
	app := kingpin.New(filepath.Base(os.Args[0]), "The Pushgateway").UsageWriter(os.Stdout)
	flags := kingpinflag.AddFlags(app, cfg.MetricAddress)

	app.Version(version.Print("pushgateway"))
	app.HelpFlag.Short('h')
	kingpin.MustParse(app.Parse([]string{}))
	quitCh := make(chan struct{})

	forwarder := newMetricForwarder()
	r := route.New()
	enableMetricForwarder := config.Global().EnableRcu
	pushAPIPath := "/metrics"
	r.Put(pushAPIPath+"/job"+"/:job/*labels", Push(storage, forwarder, true, true, false, enableMetricForwarder))
	r.Post(pushAPIPath+"/job"+"/:job/*labels", Push(storage, forwarder, false, true, false, enableMetricForwarder))
	r.Put(pushAPIPath+"/job"+"/:job", Push(storage, forwarder, true, true, false, enableMetricForwarder))
	r.Post(pushAPIPath+"/job"+"/:job", Push(storage, forwarder, false, true, false, enableMetricForwarder))

	mux := http.NewServeMux()
	mux.Handle("/", decodeRequest(r))
	server := &http.Server{Handler: mux}
	return &MetricServer{
		server:      server,
		quitCh:      quitCh,
		flags:       flags,
		gokitLogger: promlog.New(&promlog.Config{}),
	}
}

func (s *MetricServer) Run() error {
	go func() {
		_ = shutdownServerOnQuit(s.server, s.quitCh)
	}()
	return web.ListenAndServe(s.server, s.flags, s.gokitLogger)
}

func decodeRequest(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close() // Make sure the underlying io.Reader is closed.
		switch contentEncoding := r.Header.Get("Content-Encoding"); strings.ToLower(contentEncoding) {
		case "gzip":
			gr, err := gzip.NewReader(r.Body)
			if err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
			defer gr.Close()
			r.Body = gr
		case "snappy":
			r.Body = io.NopCloser(snappy.NewReader(r.Body))
		default:
			// Do nothing.
		}

		h.ServeHTTP(w, r)
	})
}

// shutdownServerOnQuit shutdowns the provided server upon closing the provided
// quitCh or upon receiving a SIGINT or SIGTERM.
func shutdownServerOnQuit(server *http.Server, quitCh <-chan struct{}) error {
	notifier := make(chan os.Signal, 1)
	signal.Notify(notifier, os.Interrupt, syscall.SIGTERM)

	select {
	case <-notifier:
		log.Info("received SIGINT/SIGTERM; exiting gracefully...")
		break
	case <-quitCh:
		log.Warn("received termination request via web service, exiting gracefully...")
		break
	}
	return server.Shutdown(context.Background())
}