package main

import (
	"encoding/json"
	"flag"
	"github.com/asaskevich/govalidator"
	"github.com/hoisie/redis"
	"github.com/julienschmidt/httprouter"
	"github.com/weppos/publicsuffix-go/publicsuffix"
	"log"
	"net/http"
	"os"
	"strings"
)

var c *redis.Client

type infoResponse struct {
	RedisServer string `json:"redis"`
}

var (
	flagKey = flag.String("key", "godns:hosts", "Redis key for hash set")
	flagListen = flag.String("listen", ":8080", "listen address")
	flagServer = flag.String("server", "localhost:6379", "redis host")
	filePath = flag.String("filepath", "/var/lib/joker/dist", "file path")
)

var (
	key string
	server string
)

func main() {
	flag.Parse()

	if key = os.Getenv("REDIS_KEY"); key == "" {
		key = *flagKey
	}

	if server = os.Getenv("REDIS_ADDR"); server == "" {
		server = *flagServer
	}

	c = &redis.Client{
		Addr: server,
	}

	router := httprouter.New()

	fs := http.FileServer(http.Dir(*filePath))

	serveFiles := func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
		fs.ServeHTTP(w, r)
	}

	router.GET("/", serveFiles)
	router.GET("/js/*filepath", serveFiles)
	router.GET("/css/*filepath", serveFiles)

	router.GET("/info", getInfo)

	router.GET("/records", getRecords)
	router.GET("/records/:zone", getRecords)

	router.POST("/update", updateRecord)
	router.POST("/remove", removeRecord)

	log.Println("Starting server on " + *flagListen)

	http.ListenAndServe(*flagListen, router)
}

func getInfo(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	json.NewEncoder(w).Encode(&infoResponse{
		RedisServer: server,
	})
}

type tableRow struct {
	Name string `json:"name"`
	IP string `json:"ip"`
	ItemCount int `json:"itemCount"`
}

type response struct {
	Success bool `json:"success"`
	Message string `json:"message"`
}

func getRecords(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	zone := p.ByName("zone")

	hosts := make(map[string]string)

	if err := c.Hgetall(key, hosts); err != nil {
		errorResponse(w, http.StatusInternalServerError, "Unable to retrieve records: " + err.Error())
		return
	}

	groups := make(map[string]map[string]string)

	for host, addr := range hosts {
		domain, err := publicsuffix.Domain(host)

		if err != nil {
			domain = host
		}

		if zone != "" && !strings.HasSuffix(host, zone) {
			continue
		}

		if group, ok := groups[domain]; ok {
			group[host] = addr
		} else {
			groups[domain] = map[string]string{host:addr}
		}
	}

	out := make([]*tableRow, 0)

	if zone == "" {
		for domain, children := range groups {
			ip := children[domain]

			if len(children) > 1 {
				ip = "-"
			} else {
				domain = mapFirst(children)
				ip = children[domain]
			}

			out = append(out, &tableRow{
				Name: domain,
				IP: ip,
				ItemCount: len(children),
			})
		}
	} else {
		for domain, ip := range groups[zone] {
			out = append(out, &tableRow{
				Name: domain,
				IP: ip,
				ItemCount: 1,
			})
		}
	}

	json.NewEncoder(w).Encode(out)
}

func mapFirst(m map[string]string) string {
	for k := range m {
		return k
	}

	return ""
}

type domainRequest struct {
	Name string `json:"name"`
	IP string `json:"ip"`
	Group bool `json:"group"`
}

func updateRecord(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	if !strings.HasPrefix(r.Header.Get("Content-Type"), "application/json") {
		errorResponse(w, http.StatusBadRequest, "Invalid body type.")
		return
	}

	var req domainRequest

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		errorResponse(w, http.StatusBadRequest, "Unable to decode body: " + err.Error())
		return
	}

	req.Name = strings.ToLower(req.Name)

	// Replace wildcards with an 'a' to test the domain
	testDomain := req.Name
	testDomain = strings.Replace(testDomain, "*", "a", -1)

	if !govalidator.IsDNSName(testDomain) || !govalidator.IsIP(req.IP) {
		errorResponse(w, http.StatusBadRequest, "Invalid domain or IP")
		return
	}

	if _, err := c.Hset(key, req.Name, []byte(req.IP)); err != nil {
		errorResponse(w, http.StatusInternalServerError, "Unable to save record: " + err.Error())
		return
	}

	c.Publish("godns:update_record", []byte(req.Domain))

	json.NewEncoder(w).Encode(&response{Success: true})
}

func removeRecord(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	var req domainRequest

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		return
	}

	req.Name = strings.ToLower(req.Name)

	if !govalidator.IsDNSName(strings.Replace(req.Name, "*", "a", -1)) {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	if req.Group {
		hosts := make(map[string]string)

		if err := c.Hgetall(key, hosts); err != nil {
			errorResponse(w, http.StatusInternalServerError, "Unable to fetch records: " + err.Error())
			return
		}

		for host, _ := range hosts {
			host = strings.ToLower(host)

			if !strings.HasSuffix(host, req.Name) {
				continue
			}

			c.Hdel(key, host)

			c.Publish("godns:remove_record", []byte(host))
		}
	} else {
		if _, err := c.Hdel(key, req.Name); err != nil {
			errorResponse(w, http.StatusInternalServerError, "Unable to delete record: " + err.Error())
			return
		}

		c.Publish("godns:remove_record", []byte(req.Domain))
	}
}

func errorResponse(w http.ResponseWriter, status int, message string) {
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(&response{Success: false, Message: message})
}