package server

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

	"codeberg.org/miekg/dns"
	"codeberg.org/miekg/dns/rdata"
	"linum/internal/blocklist"
	"linum/internal/cache"
)

func (s *Server) handleQuery(ctx context.Context, w dns.ResponseWriter, req *dns.Msg) {
	if len(req.Question) == 0 {
		m := new(dns.Msg)
		m.Rcode = dns.RcodeFormatError
		m.Response = true
		m.ID = req.ID
		m.Question = req.Question
		io.Copy(w, m)
		return
	}

	resp, blocked := s.buildResponse(req)

	resp.ID = req.ID
	resp.Data = nil
	if _, err := io.Copy(w, resp); err != nil {
		slog.Error("write response failed",
			"err",   err,
			"qname", req.Question[0].Header().Name,
			"qtype", dns.TypeToString[dns.RRToType(req.Question[0])],
		)
		return
	}
	slog.Info("query served",
		"qname",   req.Question[0].Header().Name,
		"qtype",   dns.TypeToString[dns.RRToType(req.Question[0])],
		"rcode",   dns.RcodeToString[resp.Rcode],
		"blocked", blocked,
	)
}

func (s *Server) buildResponse(req *dns.Msg) (*dns.Msg, bool) {
	if len(req.Question) == 0 {
		m := new(dns.Msg)
		m.Rcode = dns.RcodeFormatError
		m.Response = true
		m.ID = req.ID
		m.Question = req.Question
		return m, false
	}
	q := req.Question[0]
	qname := q.Header().Name
	qtype := dns.RRToType(q)
	qclass := q.Header().Class

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

	if s.cache != nil {
		key := cache.Key{Name: qname, Qtype: qtype, Class: qclass}
		if cached, ok := s.cache.Get(key); ok {
			cached.ID = req.ID
			return cached, false
		}
	}

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

	reply, err := s.resolver.Lookup(ctx, qname, qtype)
	if err != nil {
		slog.Error("resolution failed",
			"err",   err,
			"qname", qname,
			"qtype", dns.TypeToString[qtype],
		)
		m := new(dns.Msg)
		m.Rcode = dns.RcodeServerFailure
		m.Response = true
		m.ID = req.ID
		m.Question = req.Question
		return m, false
	}

	reply.ID = req.ID

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

	return reply, false
}

func (s *Server) blockedResponse(req *dns.Msg) *dns.Msg {
	m := new(dns.Msg)
	m.Response = true
	m.ID = req.ID
	m.Opcode = req.Opcode
	m.RecursionDesired = req.RecursionDesired
	m.Question = req.Question
	m.Authoritative = true

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

	q := req.Question[0]
	qname := q.Header().Name
	qtype := dns.RRToType(q)

	switch qtype {
	case dns.TypeA:
		m.Answer = append(m.Answer, &dns.A{
			Hdr: dns.Header{Name: qname, Class: dns.ClassINET, TTL: 60},
			A:   rdata.A{Addr: netip.AddrFrom4([4]byte{})},
		})
	case dns.TypeAAAA:
		m.Answer = append(m.Answer, &dns.AAAA{
			Hdr:   dns.Header{Name: qname, Class: dns.ClassINET, TTL: 60},
			AAAA: rdata.AAAA{Addr: netip.IPv6Unspecified()},
		})
	}
	return m
}