package main

import (
	"flag"
	"log"
	"math/rand"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	_ "github.com/pingcap/tidb-enterprise/minimum-edition/dashboard-ui/statik"
	"github.com/rakyll/statik/fs"
)

var (
	pdEndpoints   string
	promEndpoint  string
	tidbEndpoints string
)

func init() {
	flag.StringVar(&pdEndpoints, "pd-endpoints", "http://localhost:2379", "pd endpoints")
	flag.StringVar(&promEndpoint, "prometheus-endpoint", "http://localhost:9090", "prometheus endpoint")
	flag.StringVar(&tidbEndpoints, "tidb-endpoints", "http://localhost:10080", "tidb endpoints")
	flag.Parse()
}

type Proxy struct {
	backends []*httputil.ReverseProxy
}

func rewriteResponseHeader(resp *http.Response) error {
	resp.Header.Set("Access-Control-Allow-Headers", "accept, content-type, authorization")
	resp.Header.Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
	resp.Header.Set("Access-Control-Allow-Origin", "*")
	return nil
}

func NewProxy(targets string) *Proxy {
	backends := []*httputil.ReverseProxy{}
	for _, target := range strings.Split(targets, ",") {
		u, err := url.Parse(target)
		if err != nil {
			continue
		}
		backend := httputil.NewSingleHostReverseProxy(u)
		backend.ModifyResponse = rewriteResponseHeader
		backends = append(backends, backend)
	}
	return &Proxy{backends: backends}
}

func (p *Proxy) Handle(w http.ResponseWriter, r *http.Request) {
	// random proxy to one backend
	backend := p.backends[rand.Intn(len(p.backends))]
	backend.ServeHTTP(w, r)
}

func main() {
	statikFS, err := fs.New()
	staticHandler := http.FileServer(statikFS)
	if err != nil {
		log.Fatal(err)
	}
	pdURLs := strings.Split(pdEndpoints, ",")
	pdURL, err := url.Parse(pdURLs[0])
	if err != nil {
		log.Fatalf("failed to parse pd url %s: %v", pdURLs[0], err)
	}
	if len(pdURLs) > 1 || pdURL.Host != "localhost:2379" { // don't run proxy when PD already available at localhost:2379
		go func() {
			proxy := NewProxy(pdEndpoints)
			mux := http.NewServeMux()
			mux.HandleFunc("/", proxy.Handle)
			log.Fatal(http.ListenAndServe(":2379", mux))
		}()
	}

	promURL, err := url.Parse(promEndpoint)
	if err != nil {
		log.Fatalf("failed to parse prometheus url %s: %v", promEndpoint, err)
	}
	if promURL.Host != "localhost:9090" { // don't run proxy when Prometheus already available at localhost:9090
		go func() {
			proxy := NewProxy(promEndpoint)
			mux := http.NewServeMux()
			mux.HandleFunc("/", proxy.Handle)
			log.Fatal(http.ListenAndServe(":9090", mux))
		}()
	}

	tidbURLs := strings.Split(tidbEndpoints, ",")
	tidbURL, err := url.Parse(tidbURLs[0])
	if err != nil {
		log.Fatalf("failed to parse tidb url %s: %v", tidbURLs[0], err)
	}
	if len(tidbURLs) > 1 || tidbURL.Host != "localhost:10080" { // don't run proxy when TiDB already available at localhost:10080
		go func() {
			proxy := NewProxy(tidbEndpoints)
			mux := http.NewServeMux()
			mux.HandleFunc("/", proxy.Handle)
			log.Fatal(http.ListenAndServe(":10080", mux))
		}()
	}

	// Serves up the index.html file regardless of the path.
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		r.URL.Path = "/"
		staticHandler.ServeHTTP(w, r)
	})
	http.Handle("/static/", staticHandler)
	log.Fatal(http.ListenAndServe(":8080", nil))
}