package internal

import (
	"context"
	"time"

	"github.com/miekg/dns"
)

type DNSServer struct {
	srv *dns.Server
	dns.Handler
}

func (d *DNSServer) ListenAndServe(ctx context.Context, addr string) error {
	errs := make(chan error)

	c, canceller := context.WithCancel(ctx)
	go d.runServer(c, "udp", addr, errs)
	go d.runServer(c, "tcp", addr, errs)

	defer canceller()
	select {
	case <-ctx.Done():
		return nil
	case err := <-errs:
		return err
	}
}

func (d *DNSServer) runServer(ctx context.Context, net, addr string, in chan<- error) {
	d.srv = &dns.Server{
		Addr:        addr,
		Net:         net,
		IdleTimeout: func() time.Duration { return time.Second * 5 },
		Handler:     d,
	}

	go func() {
		<-ctx.Done()
		d.srv.Listener.Close()
	}()

	if err := d.srv.ListenAndServe(); err != nil {
		in <- err
	}
}