package nscheck

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

	"github.com/miekg/dns"
)

// fake is a DNS server on one loopback address that answers with fn.
type fake func(q dns.Question) *dns.Msg

// serve starts fake servers on 127.0.0.x addresses, all on the same port,
// and returns the port.
func serve(t *testing.T, servers map[string]fake) string {
	t.Helper()
	port := ""
	for addr, fn := range servers {
		hp := net.JoinHostPort(addr, "0")
		if port != "" {
			hp = net.JoinHostPort(addr, port)
		}
		pc, err := net.ListenPacket("udp", hp)
		if err != nil {
			t.Skipf("can't listen on %s: %v", hp, err)
		}
		if port == "" {
			_, port, _ = net.SplitHostPort(pc.LocalAddr().String())
		}
		fn := fn
		srv := &dns.Server{PacketConn: pc, Handler: dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
			resp := fn(req.Question[0])
			if resp == nil {
				return // no answer: a timeout
			}
			rcode, aa := resp.Rcode, resp.Authoritative
			resp.SetReply(req) // resets both
			resp.Rcode, resp.Authoritative = rcode, aa
			w.WriteMsg(resp)
		})}
		go srv.ActivateAndServe()
		t.Cleanup(func() { srv.Shutdown() })
	}
	return port
}

func rr(s string) dns.RR {
	r, err := dns.NewRR(s)
	if err != nil {
		panic(err)
	}
	return r
}

// authSOA answers SOA queries for zone authoritatively with serial.
func authSOA(zone string, serial uint32) fake {
	return func(q dns.Question) *dns.Msg {
		m := new(dns.Msg)
		m.Authoritative = true
		if q.Qtype == dns.TypeSOA && strings.EqualFold(q.Name, zone) {
			m.Answer = []dns.RR{rr(zone + " 3600 IN SOA ns1." + zone + " admin." + zone + " " + itoa(serial) + " 3600 900 1209600 300")}
		}
		return m
	}
}

func itoa(n uint32) string { return strconv.FormatUint(uint64(n), 10) }

func checker(port, resolver string) *Checker {
	return &Checker{
		Timeout:  300 * time.Millisecond,
		Port:     port,
		Resolver: resolver,
		IsLocal:  func(a netip.Addr) bool { return a.String() == "127.0.0.9" },
	}
}

func TestSerials(t *testing.T) {
	zone := "example.com."
	port := serve(t, map[string]fake{
		"127.0.0.2": authSOA(zone, 2026100805),
		"127.0.0.3": authSOA(zone, 2026100803),
		"127.0.0.4": func(dns.Question) *dns.Msg { return new(dns.Msg) }, // answers, not authoritative
		"127.0.0.5": func(dns.Question) *dns.Msg { m := new(dns.Msg); m.Rcode = dns.RcodeRefused; return m },
		"127.0.0.6": func(dns.Question) *dns.Msg { return nil }, // never answers
	})
	ref := uint32(2026100805)
	res := checker(port, "").Serials(context.Background(), zone, &ref, []Target{
		{Name: "ns1.example.com.", Addresses: []string{"127.0.0.9"}, Source: "NS record"},
		{Name: "ns2.example.com.", Addresses: []string{"127.0.0.2"}, Source: "NS record"},
		{Name: "ns3.example.com.", Addresses: []string{"127.0.0.3"}, Source: "NS record"},
		{Name: "ns4.example.com.", Addresses: []string{"127.0.0.4"}, Source: "NS record"},
		{Addresses: []string{"127.0.0.5"}, Source: "also-notify"},
		{Addresses: []string{"127.0.0.6", "127.0.0.2"}, Source: "also-notify"}, // 127.0.0.2 again: asked once
	})
	got := map[string]Result{}
	for _, r := range res {
		got[r.Address] = r
	}
	if len(res) != 6 {
		t.Errorf("%d results: %+v", len(res), res)
	}
	check := func(addr, status string) {
		t.Helper()
		if got[addr].Status != status {
			t.Errorf("%s: %+v, want %s", addr, got[addr], status)
		}
	}
	check("127.0.0.9", "self")
	check("127.0.0.2", "in_sync")
	check("127.0.0.3", "behind")
	check("127.0.0.4", "not_authoritative")
	check("127.0.0.5", "refused")
	check("127.0.0.6", "unreachable")
	if got["127.0.0.3"].Diff != -2 || *got["127.0.0.3"].Serial != 2026100803 {
		t.Errorf("behind: %+v", got["127.0.0.3"])
	}
	// Without a serial to compare with, an answer is just an answer.
	res = checker(port, "").Serials(context.Background(), zone, nil, []Target{{Addresses: []string{"127.0.0.2"}}})
	if res[0].Status != "answered" || !res[0].OK() {
		t.Errorf("no reference: %+v", res[0])
	}
}

func TestSerialWraparound(t *testing.T) {
	// RFC 1982: 1 is ahead of 4294967295.
	zone := "example.com."
	port := serve(t, map[string]fake{"127.0.0.2": authSOA(zone, 1)})
	ref := uint32(4294967295)
	res := checker(port, "").Serials(context.Background(), zone, &ref, []Target{{Addresses: []string{"127.0.0.2"}}})
	if res[0].Status != "ahead" || res[0].Diff != 2 {
		t.Errorf("%+v", res[0])
	}
}

// internet builds a small fake DNS tree: a resolver, a parent server for
// "lol." and the child's servers.
func internet(t *testing.T, parent fake) (port, resolver string) {
	port = serve(t, map[string]fake{
		// The recursive resolver.
		"127.0.0.10": func(q dns.Question) *dns.Msg {
			m := new(dns.Msg)
			m.RecursionAvailable = true
			switch {
			case q.Qtype == dns.TypeSOA && q.Name == "lol.":
				m.Answer = []dns.RR{rr("lol. 300 IN SOA ns.nic.lol. hostmaster.lol. 1 3600 900 1209600 300")}
			case q.Qtype == dns.TypeNS && q.Name == "lol.":
				m.Answer = []dns.RR{rr("lol. 300 IN NS ns.nic.lol.")}
			case q.Qtype == dns.TypeA && q.Name == "ns.nic.lol.":
				m.Answer = []dns.RR{rr("ns.nic.lol. 300 IN A 127.0.0.11")}
			case q.Qtype == dns.TypeA && q.Name == "ns.elsewhere.net.":
				m.Answer = []dns.RR{rr("ns.elsewhere.net. 300 IN A 127.0.0.13")}
			}
			return m
		},
		"127.0.0.11": parent,
		// ns1.adbraun.lol: serves the zone.
		"127.0.0.12": authSOA("adbraun.lol.", 7),
		// ns.elsewhere.net: doesn't serve it (lame).
		"127.0.0.13": func(dns.Question) *dns.Msg { m := new(dns.Msg); m.Rcode = dns.RcodeRefused; return m },
	})
	return port, net.JoinHostPort("127.0.0.10", port)
}

func levels(rep *DelegationReport) map[string][]string {
	out := map[string][]string{}
	for _, f := range rep.Findings {
		out[f.Level] = append(out[f.Level], f.Message)
	}
	return out
}

func TestDelegationMismatch(t *testing.T) {
	port, res := internet(t, func(q dns.Question) *dns.Msg {
		m := new(dns.Msg)
		if q.Qtype == dns.TypeNS && q.Name == "adbraun.lol." {
			m.Ns = []dns.RR{rr("adbraun.lol. 3600 IN NS ns1.adbraun.lol."), rr("adbraun.lol. 3600 IN NS ns.elsewhere.net.")}
			m.Extra = []dns.RR{rr("ns1.adbraun.lol. 3600 IN A 127.0.0.12")}
		}
		return m
	})
	ref := uint32(7)
	rep := checker(port, res).Delegation(context.Background(), "adbraun.lol.",
		[]string{"ns1.adbraun.lol.", "ns2.adbraun.lol."},
		map[string][]string{"ns1.adbraun.lol.": {"127.0.0.99"}, "ns2.adbraun.lol.": {"127.0.0.14"}}, &ref)
	if rep.Parent != "lol." || !rep.Delegated {
		t.Fatalf("report: %+v", rep)
	}
	l := levels(rep)
	all := strings.Join(append(append(l["error"], l["warning"]...), l["info"]...), "\n")
	for _, want := range []string{
		"delegates to ns.elsewhere.net, which the zone's NS records don't list",
		"list ns2.adbraun.lol, which the parent doesn't delegate to",
		"glue for ns1.adbraun.lol is 127.0.0.12, but the zone says 127.0.0.99",
		"ns.elsewhere.net (127.0.0.13): refuses queries",
	} {
		if !strings.Contains(all, want) {
			t.Errorf("missing finding %q in:\n%s", want, all)
		}
	}
	if len(rep.Glue) != 1 || rep.Glue[0].Match {
		t.Errorf("glue: %+v", rep.Glue)
	}
}

func TestDelegationGood(t *testing.T) {
	port, res := internet(t, func(q dns.Question) *dns.Msg {
		m := new(dns.Msg)
		if q.Qtype == dns.TypeNS && q.Name == "adbraun.lol." {
			m.Ns = []dns.RR{rr("adbraun.lol. 3600 IN NS ns1.adbraun.lol.")}
			m.Extra = []dns.RR{rr("ns1.adbraun.lol. 3600 IN A 127.0.0.12")}
		}
		return m
	})
	ref := uint32(7)
	rep := checker(port, res).Delegation(context.Background(), "adbraun.lol.", []string{"ns1.adbraun.lol."}, map[string][]string{"ns1.adbraun.lol.": {"127.0.0.12"}}, &ref)
	l := levels(rep)
	if len(l["error"])+len(l["warning"])+len(l["info"]) != 0 || len(l["ok"]) != 2 {
		t.Errorf("findings: %v", l)
	}
}

func TestNotDelegated(t *testing.T) {
	// An internal-only zone: one calm note, no warnings.
	port, res := internet(t, func(q dns.Question) *dns.Msg {
		m := new(dns.Msg)
		m.Authoritative = true
		m.Rcode = dns.RcodeNameError
		m.Ns = []dns.RR{rr("lol. 300 IN SOA ns.nic.lol. hostmaster.lol. 1 3600 900 1209600 300")}
		return m
	})
	rep := checker(port, res).Delegation(context.Background(), "adbraun.lol.", []string{"ns1.adbraun.lol."}, nil, nil)
	l := levels(rep)
	if rep.Delegated || len(l["info"]) != 1 || len(l["warning"])+len(l["error"]) != 0 {
		t.Errorf("findings: %v", l)
	}
	if !strings.Contains(l["info"][0], "adbraun.lol isn't delegated from the parent zone lol, so the internet can't see it. That's expected for a zone used only on your own network") {
		t.Errorf("note: %s", l["info"][0])
	}
}