package main

import (
	"errors"
	"flag"
	"io"
	"net"
	"net/http"
	"sync"

	"github.com/pingcap/log"
	"go.uber.org/zap"
)

type targetStatus struct {
	healthy bool
	mu      sync.RWMutex
}

var (
	allowCIDR   string
	allowSubnet *net.IPNet

	globalStatus = targetStatus{
		healthy: true,
	}

	ErrGetClientIP     = "failed to get client IP"
	ErrInvalidClientIP = "invalid client IP"
	ErrIllegalClientIP = "illegal client IP"
)

func init() {
	flag.StringVar(&allowCIDR, "allow-cidr", "", "allowed source IP CIDR")
	flag.Parse()
}

func getClientIP(r *http.Request) (net.IP, error) {
	clientIPStr, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		log.Error(ErrGetClientIP, zap.String("RemoteAddr", r.RemoteAddr), zap.Error(err))
		return nil, err
	}
	clientIP := net.ParseIP(clientIPStr)
	if clientIP == nil {
		err := errors.New(ErrInvalidClientIP)
		log.Error(ErrGetClientIP, zap.String("clientIP", clientIPStr), zap.Error(err))
		return nil, err
	}
	return clientIP, nil
}

func checkClientIP(clientIP net.IP) error {
	if !allowSubnet.Contains(clientIP) && !clientIP.IsLoopback() {
		log.Error(ErrIllegalClientIP, zap.Stringer("clientIP", clientIP))
		return errors.New(ErrIllegalClientIP)
	}
	return nil
}

func main() {
	var err error
	_, allowSubnet, err = net.ParseCIDR(allowCIDR)
	if err != nil {
		log.Fatal("invalid allowCIDR", zap.String("allowCIDR", allowCIDR), zap.Error(err))
	}

	// statusHandler Returns current target status
	statusHandler := func(w http.ResponseWriter, r *http.Request) {
		globalStatus.mu.RLock()
		defer globalStatus.mu.RUnlock()
		if globalStatus.healthy {
			w.WriteHeader(http.StatusOK)
			_, err := io.WriteString(w, "Up")
			if err != nil {
				http.Error(w, "failed to write response", http.StatusInternalServerError)
			}
		} else {
			http.Error(w, "Down", http.StatusInternalServerError)
		}
	}

	// toggleHandler toggles the state of the target
	toggleHandler := func(w http.ResponseWriter, r *http.Request) {
		globalStatus.mu.Lock()
		defer globalStatus.mu.Unlock()

		clientIP, err := getClientIP(r)
		if err != nil || clientIP == nil {
			http.Error(w, ErrGetClientIP, http.StatusBadRequest)
			return
		}
		log.Info("receive request", zap.String("Host", r.Host), zap.Stringer("clientIP", clientIP))
		err = checkClientIP(clientIP)
		if err != nil {
			http.Error(w, ErrIllegalClientIP, http.StatusForbidden)
			return
		}

		if globalStatus.healthy {
			globalStatus.healthy = false
		} else {
			globalStatus.healthy = true
		}
		w.WriteHeader(http.StatusOK)
		_, err = io.WriteString(w, "Toggled")
		if err != nil {
			http.Error(w, "failed to write response", http.StatusInternalServerError)
		}
	}

	// downHandler sets the status of the target to unhealthy
	downHandler := func(w http.ResponseWriter, r *http.Request) {
		globalStatus.mu.Lock()
		defer globalStatus.mu.Unlock()

		clientIP, err := getClientIP(r)
		if err != nil || clientIP == nil {
			http.Error(w, ErrGetClientIP, http.StatusBadRequest)
			return
		}
		log.Info("receive request", zap.String("Host", r.Host), zap.Stringer("clientIP", clientIP))
		err = checkClientIP(clientIP)
		if err != nil {
			http.Error(w, ErrIllegalClientIP, http.StatusForbidden)
			return
		}

		if globalStatus.healthy {
			globalStatus.healthy = false
		}
		w.WriteHeader(http.StatusOK)
		_, err = io.WriteString(w, "Down")
		if err != nil {
			http.Error(w, "failed to write response", http.StatusInternalServerError)
		}
	}

	http.HandleFunc("/healthz", statusHandler)
	http.HandleFunc("/toggle", toggleHandler)
	http.HandleFunc("/down", downHandler)

	err = http.ListenAndServe(":8080", nil)
	if err != nil {
		log.Fatal("failed to listen and serve", zap.Error(err))
	}
}