package server

import (
	"context"
	"github.com/miekg/dns"
	"log/slog"
	"time"
)

func (s *Server) handleQuery(w dns.ResponseWriter, req *dns.Msg) {
	if len(req.Question) == 0 {
		m := new(dns.Msg)
		m.SetRcode(req, dns.RcodeFormatError)
		_ = w.WriteMsg(m)
		return
	}

	resp := s.buildResponse(req)

	if err := w.WriteMsg(resp); err != nil {
		slog.Error("write response failed",
			"err", err,
			"qname", req.Question[0].Name,
			"qtype", dns.TypeToString[req.Question[0].Qtype],
		)
		return
	}
	slog.Info("query served",
		"qname", req.Question[0].Name,
		"qtype", dns.TypeToString[req.Question[0].Qtype],
		"rcode", dns.RcodeToString[resp.Rcode],
		"client", w.RemoteAddr().String(),
	)
}

func (s *Server) buildResponse(req *dns.Msg) *dns.Msg {
	if len(req.Question) == 0 {
		m := new(dns.Msg)
		m.SetRcode(req, dns.RcodeFormatError)
		return m
	}
	q := req.Question[0]
	resp := new(dns.Msg)
	resp.SetReply(req)
	resp.Authoritative = false

	if opt := req.IsEdns0(); opt != nil {
		resp.SetEdns0(4096, false)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	reply, err := s.resolver.Lookup(ctx, q.Name, q.Qtype)
	if err != nil {
		slog.Error("resolution failed",
			"err", err,
			"qname", q.Name,
			"qtype", dns.TypeToString[q.Qtype],
		)
		resp.Rcode = dns.RcodeServerFailure
		return resp
	}

	return reply
}