package server

import (
	"context"
	"log/slog"
	"time"

	"github.com/miekg/dns"
)

type Server struct {
	addr string
	logger *slog.Logger
	inner *dns.Server
}

func New(addr string, logger *slog.Logger) (*Server, error) {
	mux := dns.NewServeMux()
	mux.HandleFunc(".", handleQuery)

	inner := &dns.Server{
		Addr: addr,
		Net: "udp",
		Handler: mux,
		UDPSize: 4096,
		ReadTimeout: 5 * time.Second,
		WriteTimeout: 5 * time.Second,
	}
	return &Server{
		addr: addr,
		logger: logger,
		inner: inner,
	}, nil
}

func (s *Server) Run(ctx context.Context) error {
	errCh := make(chan error, 1)
	go func() {
		s.logger.Info("udp listener active", "addr", s.addr)
		errCh <- s.inner.ListenAndServe()
	}()

	select {
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5 *time.Second)
		defer cancel()
		if err := s.inner.ShutdownContext(shutdownCtx); err != nil {
			s.logger.Error("graceful shutdown failed", "err", err)
			return err
		}
		return ctx.Err()
	case err := <-errCh:
		return err
	}
}

func (s *Server) Close() error {
	return s.inner.Shutdown()
}

func 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
	}

	q := req.Question[0]
	resp := new(dns.Msg)
	resp.SetReply(req)
	resp.Authoritative = false

	if q.Name == "example.com." && q.Qtype == dns.TypeA {
		resp.Answer = []dns.RR{
			&dns.A{
				Hdr: dns.RR_Header{
					Name: q.Name,
					Rrtype: dns.TypeA,
					Class: dns.ClassINET,
					Ttl: 60,
				},
				A: []byte{127,0,0,1},
			},
		}
	} else {
		resp.Rcode = dns.RcodeNameError
	}

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