package pancheri

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

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

func (h *Handler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {

	// figure out how we should resolve this
	msg := new(dns.Msg)
	msg.SetReply(r)
	msg.Authoritative = true

	if len(r.Question) != 1 {
		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 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 {
				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 {
						// 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]
				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 {
						// 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 {
				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 {
				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 {
				// 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 {
		// 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
		}
	}
}