package server

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

	"github.com/miekg/dns"
	"sdns/internal/blocklist"
	"sdns/internal/cache"
)

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, blocked := 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(),
		"blocked", blocked,
	)
}

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

	if s.cache != nil {
		key := cache.Key{Name: q.Name, Qtype: q.Qtype, Class: q.Qclass}
		if cached, ok := s.cache.Get(key); ok {
			return cached, false
		}
	}

	if s.blocklist != nil && s.blocklist.IsBlocked(q.Name) {
		return s.blockedResponse(req), true
	}

	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],
		)
		m := new(dns.Msg)
		m.SetRcode(req, dns.RcodeServerFailure)
		return m, false
	}

	reply.Id = req.Id

	if s.cache != nil && reply.Rcode != dns.RcodeServerFailure {
		key := cache.Key{Name: q.Name, Qtype: q.Qtype, Class: q.Qclass}
		s.cache.Set(key, reply, 0)
	}

	return reply, false
}

func (s *Server) blockedResponse(req *dns.Msg) *dns.Msg {
	m := new(dns.Msg)
	m.SetReply(req)
	m.Authoritative = true

	if s.blocklist.Response() == blocklist.ResponseNXDOMAIN {
		m.Rcode = dns.RcodeNameError
		return m
	}

	q := req.Question[0]
	switch q.Qtype {
	case dns.TypeA:
		m.Answer = append(m.Answer, &dns.A{
			Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60},
			A:   net.IPv4(0, 0, 0, 0),
		})
	case dns.TypeAAAA:
		m.Answer = append(m.Answer, &dns.AAAA{
			Hdr:  dns.RR_Header{Name: q.Name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 60},
			AAAA: net.IPv6zero,
		})
	}
	return m
}