// Package dnsengine contains the UDP and TCP listeners and the query pipeline
// that ties the authoritative index, the policy engine, the cache and the
// recursive resolver together.
package dnsengine

import (
	"context"
	"net"
	"net/netip"
	"strconv"
	"strings"
	"time"

	"github.com/miekg/dns"

	"github.com/owen/vibedns/internal/cache"
	"github.com/owen/vibedns/internal/models"
	"github.com/owen/vibedns/internal/netutil"
	"github.com/owen/vibedns/internal/policy"
	"github.com/owen/vibedns/internal/runtimecfg"
	"github.com/owen/vibedns/internal/version"
)

// outcome records everything the query log and metrics need about one query.
type outcome struct {
	source      string
	rcode       int
	cacheHit    bool
	blocked     bool
	answerCount int
	decision    policy.Decision
	upstream    string
}

// ServeDNS implements dns.Handler. It is the single entry point for every
// query, over both UDP and TCP.
func (s *Server) ServeDNS(w dns.ResponseWriter, req *dns.Msg) {
	start := time.Now()
	protocol := "udp"
	if _, ok := w.RemoteAddr().(*net.TCPAddr); ok {
		protocol = "tcp"
	}
	client := netutil.AddrFromNetAddr(w.RemoteAddr())
	snap := s.runtime.Current()

	// Rate limiting happens before any work is done. Exceeding clients are
	// dropped without a response: replying would let an attacker use us as an
	// amplifier, which is exactly what the limiter exists to prevent.
	if !s.limiter.Allow(client) {
		s.metrics.RateLimited.Add(1)
		s.metrics.ObserveQuery(qtypeName(req), "DROPPED", models.SourceRateLimited, protocol, time.Since(start))
		s.logQuery(req, client, protocol, outcome{source: models.SourceRateLimited, rcode: -1}, snap, start)
		return
	}

	resp, out := s.respond(req, client, snap)
	if resp == nil {
		return
	}

	s.finalise(req, resp, protocol, snap)
	if err := w.WriteMsg(resp); err != nil {
		s.log.Debug("could not write DNS response", "client", client.String(), "error", err)
	}

	out.rcode = resp.Rcode
	out.answerCount = len(resp.Answer)
	elapsed := time.Since(start)
	s.metrics.ObserveQuery(qtypeName(req), dns.RcodeToString[resp.Rcode], out.source, protocol, elapsed)
	s.logQueryOutcome(req, client, protocol, out, snap, elapsed)
}

// respond produces the reply message and describes how it was produced.
func (s *Server) respond(req *dns.Msg, client netip.Addr, snap *runtimecfg.Snapshot) (*dns.Msg, outcome) {

	if req.Opcode != dns.OpcodeQuery {
		return errorReply(req, dns.RcodeNotImplemented), outcome{source: models.SourceError}
	}
	if len(req.Question) != 1 {
		// Multiple questions in one message are not defined by any RFC and no
		// real client sends them.
		return errorReply(req, dns.RcodeFormatError), outcome{source: models.SourceError}
	}

	q := req.Question[0]
	qname := strings.ToLower(dns.Fqdn(q.Name))
	do := requestDO(req)

	if q.Qclass == dns.ClassCHAOS {
		return s.chaosReply(req, snap), outcome{source: models.SourceLocal}
	}
	if q.Qclass != dns.ClassINET {
		return errorReply(req, dns.RcodeRefused), outcome{source: models.SourceRefused}
	}

	// 1. Client policy. Blocking runs first so a policy applies even to names
	//    that a local zone would otherwise answer.
	decision := snap.Policy.Evaluate(client, qname)
	if decision.Blocked {
		s.metrics.Blocked.Add(1)
		return s.blockReply(req, decision), outcome{
			source: models.SourceBlocked, blocked: true, decision: decision,
		}
	}

	// 2. Authoritative zones always win over recursion.
	if resp := snap.Zones.Answer(req, do); resp != nil {
		s.metrics.Authoritative.Add(1)
		return resp, outcome{source: models.SourceAuthoritative, decision: decision}
	}

	// 3. Recursion, subject to the ACL. Authoritative answers above remain
	//    available to clients that are not allowed to recurse.
	if !snap.Settings.DNS.Recursion {
		s.metrics.Refused.Add(1)
		return errorReply(req, dns.RcodeRefused), outcome{source: models.SourceRefused, decision: decision}
	}
	if !snap.ACL.Allowed(client) {
		s.metrics.Refused.Add(1)
		s.log.Debug("recursion denied by ACL", "client", client.String(), "name", qname)
		return errorReply(req, dns.RcodeRefused), outcome{source: models.SourceRefused, decision: decision}
	}

	// 4. Cache.
	key := cache.KeyFor(dns.Question{Name: qname, Qtype: q.Qtype, Qclass: q.Qclass}, do)
	if res := s.cache.Get(key, req); res.Hit {
		s.metrics.CacheHits.Add(1)
		src := models.SourceCache
		if res.Stale {
			s.metrics.StaleServed.Add(1)
			src = models.SourceStale
		}
		res.Msg.RecursionAvailable = true
		return res.Msg, outcome{source: src, cacheHit: true, decision: decision}
	}
	s.metrics.CacheMisses.Add(1)

	// 5. Forward upstream.
	ctx, cancel := context.WithTimeout(s.ctx, s.forwardBudget(snap))
	defer cancel()

	fstart := time.Now()
	result, err := s.resolver.Resolve(ctx, req)
	s.metrics.ObserveResolver(time.Since(fstart), err != nil)
	if err != nil {
		s.log.Debug("recursive resolution failed", "name", qname, "type", qtypeName(req), "error", err)
		s.metrics.Errors.Add(1)
		return errorReply(req, dns.RcodeServerFailure), outcome{source: models.SourceError, decision: decision}
	}

	s.metrics.Recursive.Add(1)
	// 6. Cache the answer.
	s.cache.Put(key, result.Msg)

	resp := result.Msg
	resp.Id = req.Id
	resp.Question = req.Question
	resp.RecursionAvailable = true
	return resp, outcome{source: models.SourceRecursive, decision: decision, upstream: result.Upstream}
}

// forwardBudget bounds the total time spent forwarding one query, leaving the
// client's own timeout some headroom.
func (s *Server) forwardBudget(snap *runtimecfg.Snapshot) time.Duration {
	per := time.Duration(snap.Settings.Resolver.TimeoutMS) * time.Millisecond
	attempts := snap.Settings.Resolver.Retries + 1
	total := per * time.Duration(attempts)
	if total > 15*time.Second {
		total = 15 * time.Second
	}
	if total < per {
		total = per
	}
	return total
}

// blockReply builds the response for a policy-blocked query.
func (s *Server) blockReply(req *dns.Msg, d policy.Decision) *dns.Msg {
	q := req.Question[0]
	m := new(dns.Msg)
	m.SetReply(req)
	m.RecursionAvailable = true
	m.Authoritative = true

	ttl := uint32(60)
	if d.Policy != nil && d.Policy.TTL > 0 {
		ttl = d.Policy.TTL
	}

	switch d.Action() {
	case models.BlockRefused:
		m.Rcode = dns.RcodeRefused
		return m

	case models.BlockSinkhole:
		switch q.Qtype {
		case dns.TypeA:
			if d.Policy != nil && d.Policy.SinkholeV4.IsValid() {
				m.Answer = append(m.Answer, &dns.A{
					Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: ttl},
					A:   d.Policy.SinkholeV4.AsSlice(),
				})
			}
		case dns.TypeAAAA:
			if d.Policy != nil && d.Policy.SinkholeV6.IsValid() {
				m.Answer = append(m.Answer, &dns.AAAA{
					Hdr:  dns.RR_Header{Name: q.Name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: ttl},
					AAAA: d.Policy.SinkholeV6.AsSlice(),
				})
			}
		}
		if len(m.Answer) == 0 {
			// Sinkholing only makes sense for address queries; everything else
			// gets an empty NOERROR so clients do not retry in a loop.
			m.Ns = append(m.Ns, syntheticSOA(q.Name, ttl))
		}
		return m

	default: // NXDOMAIN
		m.Rcode = dns.RcodeNameError
		m.Ns = append(m.Ns, syntheticSOA(q.Name, ttl))
		return m
	}
}

// syntheticSOA gives a blocked or synthesised negative answer something for the
// client to derive a negative cache TTL from.
func syntheticSOA(name string, ttl uint32) *dns.SOA {
	return &dns.SOA{
		Hdr: dns.RR_Header{
			Name: dns.Fqdn(name), Rrtype: dns.TypeSOA, Class: dns.ClassINET, Ttl: ttl,
		},
		Ns:      "localhost.",
		Mbox:    "hostmaster." + dns.Fqdn(name),
		Serial:  1,
		Refresh: 3600,
		Retry:   600,
		Expire:  86400,
		Minttl:  ttl,
	}
}

// chaosReply answers version.bind and hostname.bind in the CHAOS class.
func (s *Server) chaosReply(req *dns.Msg, snap *runtimecfg.Snapshot) *dns.Msg {
	q := req.Question[0]
	m := new(dns.Msg)
	m.SetReply(req)
	m.Authoritative = true

	if q.Qtype != dns.TypeTXT {
		m.Rcode = dns.RcodeRefused
		return m
	}
	name := strings.ToLower(q.Name)
	if !snap.Settings.DNS.ExposeVersion {
		// Revealing the software version by default only helps an attacker.
		m.Rcode = dns.RcodeRefused
		return m
	}
	var value string
	switch name {
	case "version.bind.", "version.server.":
		value = version.Name + " " + version.Version
	case "hostname.bind.", "id.server.":
		value = s.hostname
	default:
		m.Rcode = dns.RcodeRefused
		return m
	}
	m.Answer = append(m.Answer, &dns.TXT{
		Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeTXT, Class: dns.ClassCHAOS, Ttl: 0},
		Txt: []string{value},
	})
	return m
}

// finalise applies EDNS to the response and truncates it if it will not fit in
// the client's UDP buffer.
func (s *Server) finalise(req, resp *dns.Msg, protocol string, snap *runtimecfg.Snapshot) {
	resp.Compress = true

	advertised := uint16(dns.MinMsgSize) // 512, the pre-EDNS limit
	if opt := req.IsEdns0(); opt != nil && snap.Settings.DNS.EDNSEnabled {
		clientSize := opt.UDPSize()
		if clientSize < dns.MinMsgSize {
			clientSize = dns.MinMsgSize
		}
		ourSize := uint16(snap.Settings.DNS.EDNSUDPSize)
		if clientSize < ourSize {
			advertised = clientSize
		} else {
			advertised = ourSize
		}
		// Echo an OPT record so the client knows we speak EDNS, mirroring the
		// DO bit it asked for.
		resp.SetEdns0(ourSize, opt.Do())
	}

	if protocol == "tcp" {
		return // TCP carries up to 64 KiB; no truncation needed
	}

	maxUDP := uint16(snap.Settings.DNS.MaxUDPResponse)
	if advertised < maxUDP {
		maxUDP = advertised
	}
	if resp.Len() > int(maxUDP) {
		resp.Truncate(int(maxUDP))
		if resp.Truncated {
			s.metrics.TruncatedResp.Add(1)
		}
	}
}

func errorReply(req *dns.Msg, rcode int) *dns.Msg {
	m := new(dns.Msg)
	m.SetRcode(req, rcode)
	m.RecursionAvailable = true
	return m
}

func requestDO(req *dns.Msg) bool {
	opt := req.IsEdns0()
	return opt != nil && opt.Do()
}

func qtypeName(req *dns.Msg) string {
	if len(req.Question) == 0 {
		return "NONE"
	}
	if s, ok := dns.TypeToString[req.Question[0].Qtype]; ok {
		return s
	}
	return "TYPE" + strconv.Itoa(int(req.Question[0].Qtype))
}