package nscheck

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

	"github.com/miekg/dns"

	"jarl/internal/dnssec"
)

// 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, nil)
	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, nil)
	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, 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])
	}
}

func TestDelegationDS(t *testing.T) {
	const dsData = "22642 13 2 0B7A2F5E1C8E2F6A0B7A2F5E1C8E2F6A0B7A2F5E1C8E2F6A0B7A2F5E1C8E2F6A"
	parent := func(withDS bool) fake {
		return func(q dns.Question) *dns.Msg {
			m := new(dns.Msg)
			switch {
			case 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")}
			case q.Qtype == dns.TypeDS && q.Name == "adbraun.lol.":
				m.Authoritative = true
				if withDS {
					m.Answer = []dns.RR{rr("adbraun.lol. 3600 IN DS " + dsData)}
				}
			}
			return m
		}
	}
	ours := dnssec.DS{KeyTag: 22642, Algorithm: 13, DigestType: 2, Digest: strings.Fields(dsData)[3]}
	other := ours
	other.KeyTag, other.Digest = 1111, strings.Repeat("AB", 32)
	ref := uint32(7)
	for _, c := range []struct {
		name   string
		withDS bool
		sig    Signing
		level  string
		want   string
	}{
		{"secure", true, Signing{Policy: "default", DS: []dnssec.DS{ours}}, "ok", "matches the zone's key 22642"},
		{"wrong key", true, Signing{Policy: "default", DS: []dnssec.DS{other}}, "error", "matches none of the zone's keys"},
		{"unsigned", true, Signing{}, "error", "but the zone isn't signed"},
		{"turning off", true, Signing{Policy: "insecure", DS: []dnssec.DS{ours}}, "warning", "remove it at your registrar"},
		{"no DS yet", false, Signing{Policy: "default", DS: []dnssec.DS{ours}}, "info", "the parent has no DS record for it"},
		{"DS gone", false, Signing{Policy: "insecure"}, "ok", "turning DNSSEC off can go ahead"},
	} {
		port, res := internet(t, parent(c.withDS))
		sig := c.sig
		rep := checker(port, res).Delegation(context.Background(), "adbraun.lol.", []string{"ns1.adbraun.lol."}, map[string][]string{"ns1.adbraun.lol.": {"127.0.0.12"}}, &ref, &sig)
		found := false
		for _, f := range rep.Findings {
			if f.Level == c.level && strings.Contains(f.Message, c.want) {
				found = true
			}
		}
		if !found {
			t.Errorf("%s: no %s finding %q in %+v", c.name, c.level, c.want, rep.Findings)
		}
		if c.withDS && (len(rep.ParentDS) != 1 || rep.ParentDS[0].KeyTag != 22642) {
			t.Errorf("%s: parent DS %+v", c.name, rep.ParentDS)
		}
	}
	// Unsigned and no DS: nothing about DNSSEC.
	port, res := internet(t, parent(false))
	rep := checker(port, res).Delegation(context.Background(), "adbraun.lol.", []string{"ns1.adbraun.lol."}, map[string][]string{"ns1.adbraun.lol.": {"127.0.0.12"}}, &ref, &Signing{})
	for _, f := range rep.Findings {
		if strings.Contains(f.Message, "DS") {
			t.Errorf("unsigned zone: %+v", f)
		}
	}
}

// TestWatch: with Watch set, every query is reported as it starts and
// ends, as a dig command line with a summary of the answer.
func TestWatch(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
	})
	var mu sync.Mutex
	var events []Event
	c := checker(port, res)
	c.Watch = func(e Event) {
		mu.Lock()
		events = append(events, e)
		mu.Unlock()
	}
	ref := uint32(7)
	c.Delegation(context.Background(), "adbraun.lol.", []string{"ns1.adbraun.lol."}, map[string][]string{"ns1.adbraun.lol.": {"127.0.0.12"}}, &ref, nil)
	open := map[int64]string{}
	var queries, answers []string
	for _, e := range events {
		switch e.Kind {
		case "start":
			if !strings.HasPrefix(e.Query, "dig @") {
				t.Errorf("query %q", e.Query)
			}
			open[e.ID] = e.Query
			queries = append(queries, e.Query)
		case "end":
			q, ok := open[e.ID]
			if !ok {
				t.Errorf("end without start: %+v", e)
			}
			delete(open, e.ID)
			a := e.Answer + e.Error
			if e.Result != nil {
				a = e.Result.Status
			}
			answers = append(answers, q+" => "+a)
		}
	}
	if len(open) != 0 {
		t.Errorf("never ended: %v", open)
	}
	all := strings.Join(answers, "\n")
	for _, want := range []string{
		" adbraun.lol. NS +norec => referral to ns1.adbraun.lol.",        // asking a parent server
		"@127.0.0.12 -p " + port + " adbraun.lol. SOA +norec => in_sync", // asking the delegated server
	} {
		if !strings.Contains(all, want) {
			t.Errorf("no %q in\n%s", want, all)
		}
	}
	if len(queries) < 4 {
		t.Errorf("only %d queries:\n%s", len(queries), all)
	}
}