package pancheri

import (
	"github.com/miekg/dns"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/sirupsen/logrus"
	"strings"
)

type Handler struct {
	C *Config
	R *Resolver
	A *Authority
	B *Blackholer
}

var (
	numberOfRequests = promauto.NewCounter(prometheus.CounterOpts{
		Name: "pancheri_handler_requests_total",
		Help: "Number of DNS requests that have been handled by pancheri",
	})
	numberOfARequests = promauto.NewCounter(prometheus.CounterOpts{
		Name: "pancheri_handler_a_requests_total",
		Help: "Number of type A DNS requests that have been handled by pancheri",
	})
	numberOfARequestsResultingInCNAME = promauto.NewCounter(prometheus.CounterOpts{
		Name: "pancheri_handler_a_to_cname_requests_total",
		Help: "Number of type A DNS requests that resolved as a CNAME that have been handled by pancheri",
	})
	numberOfAAAARequests = promauto.NewCounter(prometheus.CounterOpts{
		Name: "pancheri_handler_aaaa_requests_total",
		Help: "Number of type A DNS requests that have been handled by pancheri",
	})
	numberOfAAAARequestsResultingInCNAME = promauto.NewCounter(prometheus.CounterOpts{
		Name: "pancheri_handler_aaaa_to_cname_requests_total",
		Help: "Number of type AAAA DNS requests that resolved as a CNAME that have been handled by pancheri",
	})
	numberOfCNAMERequests = promauto.NewCounter(prometheus.CounterOpts{
		Name: "pancheri_handler_cname_requests_total",
		Help: "Number of type CNAME DNS requests that have been handled by pancheri",
	})
	numberOfTXTRequests = promauto.NewCounter(prometheus.CounterOpts{
		Name: "pancheri_handler_txt_requests_total",
		Help: "Number of type CNAME DNS requests that have been handled by pancheri",
	})
	numberOfUnsupportedRequests = promauto.NewCounter(prometheus.CounterOpts{
		Name: "pancheri_handler_unsupported_requests_total",
		Help: "Number of unsupported type requests that have been handled by pancheri",
	})
	numberOfMultiQuestionRequests = promauto.NewCounter(prometheus.CounterOpts{
		Name: "pancheri_handler_multiq_requests_total",
		Help: "Number of unsupported multi-question requests that have been handled by pancheri",
	})
	numberOfUpstreamedRequests = promauto.NewCounter(prometheus.CounterOpts{
		Name: "pancheri_handler_upstreamed_requests_total",
		Help: "Number of requests that were sent upstream",
	})
)

func (h *Handler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	numberOfRequests.Inc()
	// figure out how we should resolve this
	msg := new(dns.Msg)
	msg.SetReply(r)
	msg.Authoritative = true

	if len(r.Question) != 1 {
		numberOfMultiQuestionRequests.Inc()
		msg.Rcode = dns.RcodeFormatError
		err := w.WriteMsg(msg)
		if err != nil {
			logrus.Errorf("error responding: %s", err)
			return
		}
		return
	}

	q := r.Question[0]

	// is it blackholed?
	if h.B.ShouldBlock(strings.TrimRight(q.Name, ".")) {
		// return nxdomain
		// alright, send an nxdomain
		if r.RecursionDesired {
			msg.RecursionAvailable = true
		}
		msg.Rcode = dns.RcodeNameError

		err := w.WriteMsg(msg)
		if err != nil {
			logrus.Errorf("error responding: %s", err)
			return
		}
		return
	}

	// is it ours?
	for authority, zone := range h.A.Zones {
		if strings.HasSuffix(q.Name, authority) {
			// ok, its an owned domain - respond to it
			logrus.WithFields(logrus.Fields{
				"name":  q.Name,
				"qtype": q.Qtype,
				"zone":  zone.Zonefile,
				"raw":   q.String(),
			}).Trace("responding to query for authoritative zone")

			if q.Qtype == dns.TypeA {
				numberOfARequests.Inc()
				record, ok := zone.ARecords[q.Name]
				if !ok {
					// SPECIAL CASE: for A and AAAA records, return with a CNAME if and only if that cname exists
					cname, cok := zone.CNAMERecords[q.Name]
					if cok {
						numberOfARequestsResultingInCNAME.Inc()
						// return with the CNAME record instead and resolve the CNAME
						rendered := cname.Render()
						msg.Rcode = dns.RcodeSuccess
						msg.Answer = []dns.RR{
							rendered,
						}
						msg.RecursionAvailable = true
						msg.Authoritative = true

						for authority, zone := range h.A.Zones {
							if strings.HasSuffix(cname.Target, authority) {
								// its authoritative
								if rec, ok := zone.CNAMERecords[cname.Target]; ok {
									// double cname isn't allowed right now
									logrus.WithFields(logrus.Fields{
										"in":      q.Name,
										"target":  cname.Target,
										"target2": rec.Target,
									}).Error("double cname")
									msg.Answer = []dns.RR{}
									msg.Rcode = dns.RcodeNameError
									break
								}
								if rec, ok := zone.ARecords[cname.Target]; ok {
									msg.Answer = append(msg.Answer, rec.Render())
								}
							}
						}
						if len(msg.Answer) == 1 && msg.Rcode == dns.RcodeSuccess {
							// ok, it's not ours
							// this is also not allowed right now (TODO)
							logrus.WithFields(logrus.Fields{
								"in":     q.Name,
								"target": cname.Target,
							}).Error("external cname (TODO)")
							msg.Answer = []dns.RR{}
							msg.Rcode = dns.RcodeNameError
							break
						}

						err := w.WriteMsg(msg)
						if err != nil {
							logrus.Errorf("error responding: %s", err)
							return
						}
						return
					} else {
						// send an nxdomain
						if r.RecursionDesired {
							msg.RecursionAvailable = true
						}
						msg.Rcode = dns.RcodeNameError
						err := w.WriteMsg(msg)
						if err != nil {
							logrus.Errorf("error responding: %s", err)
							return
						}
						return
					}
				}
				// return with the record
				rendered := record.Render()
				msg.Rcode = dns.RcodeSuccess
				msg.Answer = []dns.RR{
					rendered,
				}
				msg.RecursionAvailable = true
				msg.Authoritative = true
				err := w.WriteMsg(msg)
				if err != nil {
					logrus.Errorf("error responding: %s", err)
					return
				}
				return
			} else if q.Qtype == dns.TypeAAAA {
				record, ok := zone.AAAARecords[q.Name]
				numberOfAAAARequests.Inc()
				if !ok {
					// SPECIAL CASE: for A and AAAA records, return with a CNAME if and only if that cname exists
					cname, cok := zone.CNAMERecords[q.Name]
					if cok {
						numberOfAAAARequestsResultingInCNAME.Inc()
						// return with the CNAME record instead and resolve the CNAME
						rendered := cname.Render()
						msg.Rcode = dns.RcodeSuccess
						msg.Answer = []dns.RR{
							rendered,
						}
						msg.RecursionAvailable = true
						msg.Authoritative = true

						for authority, zone := range h.A.Zones {
							if strings.HasSuffix(cname.Target, authority) {
								// its authoritative
								if rec, ok := zone.CNAMERecords[cname.Target]; ok {
									// double cname isn't allowed right now
									logrus.WithFields(logrus.Fields{
										"in":      q.Name,
										"target":  cname.Target,
										"target2": rec.Target,
									}).Error("double cname")
									msg.Answer = []dns.RR{}
									msg.Rcode = dns.RcodeNameError
									break
								}
								if rec, ok := zone.AAAARecords[cname.Target]; ok {
									msg.Answer = append(msg.Answer, rec.Render())
								}
							}
						}
						if len(msg.Answer) == 1 && msg.Rcode == dns.RcodeSuccess {
							// ok, it's not ours
							// this is also not allowed right now (TODO)
							logrus.WithFields(logrus.Fields{
								"in":     q.Name,
								"target": cname.Target,
							}).Error("external cname (TODO)")
							msg.Answer = []dns.RR{}
							msg.Rcode = dns.RcodeNameError
							break
						}

						err := w.WriteMsg(msg)
						if err != nil {
							logrus.Errorf("error responding: %s", err)
							return
						}
						return
					} else {
						// send an nxdomain
						if r.RecursionDesired {
							msg.RecursionAvailable = true
						}
						msg.Rcode = dns.RcodeNameError
						err := w.WriteMsg(msg)
						if err != nil {
							logrus.Errorf("error responding: %s", err)
							return
						}
						return
					}
				}
				// return with the record
				rendered := record.Render()
				msg.Rcode = dns.RcodeSuccess
				msg.Answer = []dns.RR{
					rendered,
				}
				msg.RecursionAvailable = true
				msg.Authoritative = true
				err := w.WriteMsg(msg)
				logrus.Trace(msg.String())
				if err != nil {
					logrus.Errorf("error responding: %s", err)
					return
				}
				return
			} else if q.Qtype == dns.TypeCNAME {
				numberOfCNAMERequests.Inc()
				record, ok := zone.CNAMERecords[q.Name]
				if !ok {
					// send an nxdomain
					if r.RecursionDesired {
						msg.RecursionAvailable = true
					}
					msg.Rcode = dns.RcodeNameError
					err := w.WriteMsg(msg)
					if err != nil {
						logrus.Errorf("error responding: %s", err)
						return
					}
					return
				}
				// return with the record
				rendered := record.Render()
				msg.Rcode = dns.RcodeSuccess
				msg.Answer = []dns.RR{
					rendered,
				}
				msg.RecursionAvailable = true
				msg.Authoritative = true
				err := w.WriteMsg(msg)
				logrus.Trace(msg.String())
				if err != nil {
					logrus.Errorf("error responding: %s", err)
					return
				}
				return
			} else if q.Qtype == dns.TypeTXT {
				numberOfTXTRequests.Inc()
				record, ok := zone.TXTRecords[q.Name]
				if !ok {
					// send an nxdomain
					if r.RecursionDesired {
						msg.RecursionAvailable = true
					}
					msg.Rcode = dns.RcodeNameError
					err := w.WriteMsg(msg)
					if err != nil {
						logrus.Errorf("error responding: %s", err)
						return
					}
					return
				}
				// return with the record
				rendered := record.Render()
				msg.Rcode = dns.RcodeSuccess
				msg.Answer = []dns.RR{
					rendered,
				}
				msg.RecursionAvailable = true
				msg.Authoritative = true
				err := w.WriteMsg(msg)
				logrus.Trace(msg.String())
				if err != nil {
					logrus.Errorf("error responding: %s", err)
					return
				}
				return
			} else {
				numberOfUnsupportedRequests.Inc()
				// not supported
				logrus.WithFields(logrus.Fields{
					"name":  q.Name,
					"qtype": q.Qtype,
				}).Error("received unsupported question type")
				// send an nxdomain
				if r.RecursionDesired {
					msg.RecursionAvailable = true
				}
				msg.Rcode = dns.RcodeNameError
				err := w.WriteMsg(msg)
				if err != nil {
					logrus.Errorf("error responding: %s", err)
					return
				}
				return
			}

			return
		}
	}
	// no. do we have upstream resolution enabled?
	if h.C.Resolver.Enable {
		numberOfUpstreamedRequests.Inc()
		// alright, resolve it with the resolver
		resp, err := h.R.Resolve(q.Name, q.Qtype)
		resp.SetReply(r)

		if err != nil {
			logrus.Errorf("error resolving: %s", err)
			return
		}

		err = w.WriteMsg(resp)
		if err != nil {
			logrus.Errorf("error responding: %s", err)
			return
		}
	} else {
		// alright, send an nxdomain
		if r.RecursionDesired {
			msg.RecursionAvailable = true
		}
		msg.Rcode = dns.RcodeNameError

		err := w.WriteMsg(msg)
		if err != nil {
			logrus.Errorf("error responding: %s", err)
			return
		}
	}
}