// Package dnsd is the LAN's authoritative DNS server. It answers the configured
// zone (e.g. penguins.lan) from user-owned records plus dynamically-synthesized
// per-device records, and forwards anything outside the zone to upstream
// resolvers. Modeled on the sibling hatecomputers.club hcdns server.
package dnsd

import (
	"database/sql"
	"fmt"
	"log"
	"strings"

	"github.com/miekg/dns"

	"penguins.lan/args"
	"penguins.lan/database"
	"penguins.lan/utils"
)

const maxRecursion = 10

type Handler struct {
	DB        *sql.DB
	Zone      string // fqdn, e.g. "penguins.lan."
	ServerIP  string // apex A target; "" leaves the apex unanswered
	Resolvers []string
}

// resolveExternal forwards a query to the first upstream resolver that answers.
// With no resolvers configured (the LAN has no outside internet) it returns
// nothing.
func (h *Handler) resolveExternal(name string, qtype uint16) ([]dns.RR, error) {
	if len(h.Resolvers) == 0 {
		return nil, nil
	}
	client := &dns.Client{}
	msg := &dns.Msg{}
	msg.SetQuestion(name, qtype)
	msg.RecursionDesired = true

	var lastErr error
	for _, resolver := range h.Resolvers {
		in, _, err := client.Exchange(msg, resolver)
		if err != nil {
			lastErr = err
			continue
		}
		return in.Answer, nil
	}
	return nil, lastErr
}

// synthesizeDevice answers <compact-mac>.<username>.<zone> with the device's
// current LAN IP, if the label is a MAC owned by that user and the device is
// live. These records aren't stored — they track DHCP in real time.
func (h *Handler) synthesizeDevice(name string, qtype uint16) ([]dns.RR, error) {
	if qtype != dns.TypeA {
		return nil, nil
	}
	rest, ok := strings.CutSuffix(name, "."+h.Zone)
	if !ok {
		return nil, nil
	}
	labels := strings.Split(rest, ".")
	if len(labels) != 2 {
		return nil, nil // only <mac>.<username> lives directly under the zone
	}
	macLabel, username := labels[0], labels[1]

	mac, err := utils.ParseMacCompact(macLabel)
	if err != nil {
		return nil, nil // not a MAC label
	}
	owned, err := database.MACOwnedByUsername(h.DB, mac.Format(), username)
	if err != nil {
		return nil, err
	}
	if !owned {
		return nil, nil
	}
	ip, err := database.IPForMac(h.DB, mac.Format())
	if err != nil || ip == "" {
		return nil, err
	}
	rr, err := dns.NewRR(fmt.Sprintf("%s 30 IN A %s", name, ip))
	if err != nil {
		return nil, err
	}
	return []dns.RR{rr}, nil
}

// resolveLocal answers an in-zone name: stored CNAMEs (followed recursively),
// then stored records of the requested type, then the apex and dynamic device
// records as fallbacks.
func (h *Handler) resolveLocal(name string, qtype uint16, depth int) ([]dns.RR, error) {
	if depth == 0 {
		return nil, fmt.Errorf("dns: too much recursion resolving %s", name)
	}

	var answers []dns.RR

	cnames, err := database.FindDNSRecords(h.DB, name, "CNAME")
	if err != nil {
		return nil, err
	}
	for _, rec := range cnames {
		rr, err := dns.NewRR(fmt.Sprintf("%s %d IN CNAME %s", rec.Name, rec.TTL, dns.Fqdn(rec.Content)))
		if err != nil {
			return nil, err
		}
		answers = append(answers, rr)

		if qtype != dns.TypeCNAME {
			followed, err := h.resolveLocal(dns.Fqdn(rec.Content), qtype, depth-1)
			if err != nil {
				return nil, err
			}
			answers = append(answers, followed...)
		}
	}

	qtypeName := dns.TypeToString[qtype]
	records, err := database.FindDNSRecords(h.DB, name, qtypeName)
	if err != nil {
		return nil, err
	}
	for _, rec := range records {
		rr, err := dns.NewRR(fmt.Sprintf("%s %d IN %s %s", rec.Name, rec.TTL, qtypeName, rec.Content))
		if err != nil {
			return nil, err
		}
		answers = append(answers, rr)
	}

	if len(answers) > 0 {
		return answers, nil
	}

	if name == h.Zone && qtype == dns.TypeA && h.ServerIP != "" {
		rr, err := dns.NewRR(fmt.Sprintf("%s 300 IN A %s", name, h.ServerIP))
		if err != nil {
			return nil, err
		}
		return []dns.RR{rr}, nil
	}

	return h.synthesizeDevice(name, qtype)
}

func (h *Handler) inZone(name string) bool {
	return name == h.Zone || strings.HasSuffix(name, "."+h.Zone)
}

func (h *Handler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	msg := &dns.Msg{}
	msg.SetReply(r)

	for _, q := range r.Question {
		name := dns.Fqdn(strings.ToLower(q.Name))
		if h.inZone(name) {
			msg.Authoritative = true
			answers, err := h.resolveLocal(name, q.Qtype, maxRecursion)
			if err != nil {
				log.Println("dnsd: resolve failed:", err)
				msg.SetRcode(r, dns.RcodeServerFailure)
				w.WriteMsg(msg)
				return
			}
			msg.Answer = append(msg.Answer, answers...)
		} else {
			answers, err := h.resolveExternal(name, q.Qtype)
			if err != nil {
				log.Println("dnsd: forward failed:", err)
				msg.SetRcode(r, dns.RcodeServerFailure)
				w.WriteMsg(msg)
				return
			}
			msg.Answer = append(msg.Answer, answers...)
		}
	}

	if len(msg.Answer) == 0 && msg.Authoritative {
		msg.SetRcode(r, dns.RcodeNameError)
	}
	w.WriteMsg(msg)
}

func MakeServer(argv *args.Arguments, db *sql.DB) *dns.Server {
	handler := &Handler{
		DB:        db,
		Zone:      dns.Fqdn(argv.DnsZone),
		ServerIP:  argv.ServerIP,
		Resolvers: argv.DnsResolvers,
	}
	return &dns.Server{
		Addr:    fmt.Sprintf(":%d", argv.DnsPort),
		Net:     "udp",
		Handler: handler,
		UDPSize: 65535,
	}
}