package server

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

	"github.com/miekg/dns"
	"linum/internal/resolver"
	"linum/internal/blocklist"
	"linum/internal/cache"
)

type Server struct {
	logger   *slog.Logger
	resolver *resolver.Resolver
	cache *cache.Cache
	blocklist *blocklist.Blocklist
	udp      *dns.Server
	tcp      *dns.Server
	doh      *http.Server
}

func New(udpAddr, tcpAddr, dohAddr string, logger *slog.Logger,
r *resolver.Resolver, c *cache.Cache, b *blocklist.Blocklist) (*Server, error) {
	s := &Server{logger: logger, resolver: r, cache: c, blocklist: b}
	mux := dns.NewServeMux()
	mux.HandleFunc(".", s.handleQuery)

	if udpAddr != "" {
		s.udp = &dns.Server{
			Addr:         udpAddr,
			Net:          "udp",
			Handler:      mux,
			UDPSize:      4096,
			ReadTimeout:  5 * time.Second,
			WriteTimeout: 5 * time.Second,
			ReusePort:    true,
		}
	}

	if tcpAddr != "" {
		s.tcp = &dns.Server{
			Addr:         tcpAddr,
			Net:          "tcp",
			Handler:      mux,
			ReadTimeout:  5 * time.Second,
			WriteTimeout: 5 * time.Second,
		}
	}

	if dohAddr != "" {
		dohMux := http.NewServeMux()
		dohMux.HandleFunc("/dns-query", s.dohHandler)
		s.doh = &http.Server{
			Addr:         dohAddr,
			Handler:      dohMux,
			ReadTimeout:  5 * time.Second,
			WriteTimeout: 5 * time.Second,
		}
	}
	return s, nil
}
func (s *Server) Run(ctx context.Context) error {
	errCh := make(chan error, 3)
	if s.udp != nil {
		go func() {
			s.logger.Info("udp listener active", "addr", s.udp.Addr)
			errCh <- s.udp.ListenAndServe()
		}()
	}

	if s.tcp != nil {
		go func() {
			s.logger.Info("tcp listener active", "addr", s.tcp.Addr)
			errCh <- s.tcp.ListenAndServe()
		}()
	}

	if s.doh != nil {
		go func() {
			s.logger.Info("doh listener active", "addr", s.doh.Addr)
			errCh <- s.doh.ListenAndServe()
		}()
	}

	select {
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		if s.udp != nil {
			s.udp.ShutdownContext(shutdownCtx)
		}
		if s.tcp != nil {
			s.tcp.ShutdownContext(shutdownCtx)
		}
		if s.doh != nil {
			s.doh.Shutdown(shutdownCtx)
		}
		return ctx.Err()
	case err := <-errCh:
		return err
	}
}

func (s *Server) Close() error {
	if s.udp != nil {
		s.udp.Shutdown()
	}
	if s.tcp != nil {
		s.tcp.Shutdown()
	}
	if s.doh != nil {
		s.doh.Close()
	}
	return nil
}