// Package integration provides end-to-end tests for ExploreDNS using a mock
// DNS exchange that allows deterministic, network-independent testing of the
// full engine through its exported API.
package integration

import (
	"context"
	"math"
	"net"
	"testing"
	"time"

	dnsinternal "gitea.hansenits.com.au/hits/ExploreDNS/internal/dns"
	"gitea.hansenits.com.au/hits/ExploreDNS/internal/traverse"
	"github.com/miekg/dns"
)

// mockNet maps (server IP, qname, qtype) to a canned response, mirroring how
// distinct nameservers answer differently for the same question.
type mockNet struct {
	responses map[string]*dns.Msg
}

func newMockNet() *mockNet {
	return &mockNet{responses: make(map[string]*dns.Msg)}
}

func key(server, qname string, qtype uint16) string {
	return server + "|" + dns.Fqdn(qname) + "|" + dns.TypeToString[qtype]
}

func (m *mockNet) on(server, qname string, qtype uint16, msg *dns.Msg) {
	m.responses[key(server, qname, qtype)] = msg
}

func (m *mockNet) exchange(_ context.Context, server string, msg *dns.Msg, _ bool) (*dns.Msg, error) {
	host := server
	if h, _, err := net.SplitHostPort(server); err == nil {
		host = h
	}
	q := msg.Question[0]
	resp, ok := m.responses[key(host, q.Name, q.Qtype)]
	if !ok {
		// Unknown question: NXDOMAIN, like an authoritative miss.
		out := new(dns.Msg)
		out.SetReply(msg)
		out.Rcode = dns.RcodeNameError
		return out, nil
	}
	out := resp.Copy()
	out.SetReply(msg)
	out.Answer, out.Ns, out.Extra = resp.Answer, resp.Ns, resp.Extra
	out.Rcode = resp.Rcode
	return out, nil
}

func aRR(name, ip string) dns.RR {
	return &dns.A{
		Hdr: dns.RR_Header{Name: dns.Fqdn(name), Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 300},
		A:   net.ParseIP(ip).To4(),
	}
}

func nsRR(zone, target string) dns.RR {
	return &dns.NS{
		Hdr: dns.RR_Header{Name: dns.Fqdn(zone), Rrtype: dns.TypeNS, Class: dns.ClassINET, Ttl: 300},
		Ns:  dns.Fqdn(target),
	}
}

func cnameRR(name, target string) dns.RR {
	return &dns.CNAME{
		Hdr:    dns.RR_Header{Name: dns.Fqdn(name), Rrtype: dns.TypeCNAME, Class: dns.ClassINET, Ttl: 300},
		Target: dns.Fqdn(target),
	}
}

func answerMsg(rrs ...dns.RR) *dns.Msg {
	m := new(dns.Msg)
	m.Answer = rrs
	return m
}

func referralMsg(nsRRs []dns.RR, glue ...dns.RR) *dns.Msg {
	m := new(dns.Msg)
	m.Ns = nsRRs
	m.Extra = glue
	return m
}

func newTraverser(maxDepth int) *traverse.Traverser {
	return traverse.NewTraverser(&traverse.TraverserConfig{
		MaxDepth:  maxDepth,
		QueryType: dnsinternal.TypeA,
		RootAddrs: []net.IP{net.ParseIP("198.41.0.4")},
		QueryConfig: &dnsinternal.QueryConfig{
			Retries:    1,
			Timeout:    time.Second,
			RetryDelay: time.Millisecond,
		},
	})
}

func run(t *testing.T, tr *traverse.Traverser, m *mockNet, qname string) *traverse.Referral {
	t.Helper()
	tr.SetExchange(m.exchange)
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	root, err := tr.Run(ctx, qname)
	if err != nil {
		t.Fatalf("Run(%q): %v", qname, err)
	}
	assertProbabilityInvariant(t, root)
	return root
}

// assertProbabilityInvariant checks the engine ground rule: aggregated leaf
// probabilities at the root sum to 1.0.
func assertProbabilityInvariant(t *testing.T, root *traverse.Referral) {
	t.Helper()
	sum := 0.0
	for _, leaf := range root.StatsList() {
		sum += leaf.Prob
	}
	if math.Abs(sum-1.0) > 1e-9 {
		t.Errorf("leaf probabilities sum to %v, want 1.0", sum)
	}
}

func statuses(root *traverse.Referral) map[traverse.Status]float64 {
	out := make(map[traverse.Status]float64)
	for _, leaf := range root.StatsList() {
		out[leaf.Response.Status] += leaf.Prob
	}
	return out
}

// TestIntegrationReferralChain verifies the classic delegation walk:
// root → com → example.com with per-server responses.
func TestIntegrationReferralChain(t *testing.T) {
	m := newMockNet()
	m.on("198.41.0.4", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("com", "a.gtld-servers.test")},
		aRR("a.gtld-servers.test", "192.5.6.30"),
	))
	m.on("192.5.6.30", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("example.com", "ns1.example.com")},
		aRR("ns1.example.com", "1.2.3.4"),
	))
	m.on("1.2.3.4", "www.example.com", dns.TypeA, answerMsg(aRR("www.example.com", "93.184.216.34")))

	root := run(t, newTraverser(10), m, "www.example.com")
	got := statuses(root)
	if math.Abs(got[traverse.StatusAnswered]-1.0) > 1e-9 {
		t.Errorf("statuses = %v, want 100%% answered", got)
	}
}

// TestIntegrationCNAMERestart verifies that an out-of-zone CNAME target
// restarts the traversal from the branch cache.
func TestIntegrationCNAMERestart(t *testing.T) {
	m := newMockNet()
	m.on("198.41.0.4", "www.example.com", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("example.com", "ns1.example.com")},
		aRR("ns1.example.com", "1.2.3.4"),
	))
	m.on("1.2.3.4", "www.example.com", dns.TypeA, answerMsg(cnameRR("www.example.com", "cdn.example.net")))
	m.on("198.41.0.4", "cdn.example.net", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("example.net", "ns1.example.net")},
		aRR("ns1.example.net", "5.6.7.8"),
	))
	m.on("5.6.7.8", "cdn.example.net", dns.TypeA, answerMsg(aRR("cdn.example.net", "93.184.216.35")))

	root := run(t, newTraverser(10), m, "www.example.com")
	got := statuses(root)
	if math.Abs(got[traverse.StatusAnswered]-1.0) > 1e-9 {
		t.Errorf("statuses = %v, want 100%% answered via restart", got)
	}
	for _, leaf := range root.StatsList() {
		if leaf.Response.Status == traverse.StatusAnswered && leaf.Response.Qname != "cdn.example.net" {
			t.Errorf("answered qname = %q, want the CNAME target", leaf.Response.Qname)
		}
	}
}

// TestIntegrationNXDOMAIN verifies rcode errors surface as error leaves with
// the reference wording.
func TestIntegrationNXDOMAIN(t *testing.T) {
	m := newMockNet()
	// mock returns NXDOMAIN for anything unmocked
	root := run(t, newTraverser(10), m, "nonexistent.example.test")
	got := statuses(root)
	if math.Abs(got[traverse.StatusError]-1.0) > 1e-9 {
		t.Errorf("statuses = %v, want 100%% error", got)
	}
	for _, leaf := range root.StatsList() {
		if leaf.Response.DQ.ErrorMessage != "No such domain (NXDOMAIN)" {
			t.Errorf("error message = %q", leaf.Response.DQ.ErrorMessage)
		}
	}
}

// TestIntegrationSERVFAIL verifies SERVFAIL classification.
func TestIntegrationSERVFAIL(t *testing.T) {
	m := newMockNet()
	sf := new(dns.Msg)
	sf.Rcode = dns.RcodeServerFailure
	m.on("198.41.0.4", "example.com", dns.TypeA, sf)

	root := run(t, newTraverser(10), m, "example.com")
	for _, leaf := range root.StatsList() {
		if leaf.Response.Status != traverse.StatusError {
			t.Errorf("status = %q, want error", leaf.Response.Status)
		}
		if leaf.Response.DQ.ErrorMessage != "Server failure (SERVFAIL)" {
			t.Errorf("error message = %q", leaf.Response.DQ.ErrorMessage)
		}
	}
}

// TestIntegrationCNAMELoop verifies cross-response CNAME loops terminate as
// cname_loop leaves.
func TestIntegrationCNAMELoop(t *testing.T) {
	m := newMockNet()
	m.on("198.41.0.4", "www.a.test", dns.TypeA, answerMsg(cnameRR("www.a.test", "www.b.test")))
	m.on("198.41.0.4", "www.b.test", dns.TypeA, answerMsg(cnameRR("www.b.test", "www.a.test")))

	root := run(t, newTraverser(10), m, "www.a.test")
	got := statuses(root)
	if math.Abs(got[traverse.StatusCNAMELoop]-1.0) > 1e-9 {
		t.Errorf("statuses = %v, want 100%% cname_loop", got)
	}
}

// TestIntegrationMaxDepthExceeded verifies that an endless referral chain is
// cut off with a "Maxdepth N exceeded" exception leaf.
func TestIntegrationMaxDepthExceeded(t *testing.T) {
	m := newMockNet()
	// Each hop delegates one label deeper: the node at depth 3 (refid 1.1.1)
	// is never queried because MaxDepth 3 injects the exception first.
	m.on("198.41.0.4", "www.d2.d1", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("d1", "ns.d1")},
		aRR("ns.d1", "10.0.0.1"),
	))
	m.on("10.0.0.1", "www.d2.d1", dns.TypeA, referralMsg(
		[]dns.RR{nsRR("d2.d1", "ns.d2.d1")},
		aRR("ns.d2.d1", "10.0.0.2"),
	))

	root := run(t, newTraverser(3), m, "www.d2.d1")

	foundMaxdepth := false
	for _, leaf := range root.StatsList() {
		if leaf.Response.Status == traverse.StatusException &&
			leaf.Response.DQ.ExceptionMessage == "Maxdepth 3 exceeded" {
			foundMaxdepth = true
		}
	}
	if !foundMaxdepth {
		t.Errorf("expected a Maxdepth 3 exceeded exception leaf, got %v", statuses(root))
	}
}

// TestIntegrationHooksReceiveEvents verifies start/answer events pair up.
func TestIntegrationHooksReceiveEvents(t *testing.T) {
	m := newMockNet()
	m.on("198.41.0.4", "example.com", dns.TypeA, answerMsg(aRR("example.com", "93.184.216.34")))

	tr := newTraverser(10)
	var startEvents, answerEvents int
	tr.SetHooks(&traverse.TraverserHooks{
		OnEvent: func(e traverse.TraversalEvent) {
			switch e.Stage {
			case traverse.StageStart:
				startEvents++
			case traverse.StageAnswer:
				answerEvents++
			}
		},
	})
	run(t, tr, m, "example.com")

	if startEvents == 0 {
		t.Error("expected at least one start event")
	}
	if startEvents != answerEvents {
		t.Errorf("start events (%d) != answer events (%d)", startEvents, answerEvents)
	}
}

// TestIntegrationContextCancellation verifies that the traversal respects
// context cancellation and returns an appropriate error.
func TestIntegrationContextCancellation(t *testing.T) {
	m := newMockNet()
	m.on("198.41.0.4", "example.com", dns.TypeA, answerMsg(aRR("example.com", "93.184.216.34")))

	tr := newTraverser(10)
	tr.SetExchange(m.exchange)
	ctx, cancel := context.WithCancel(context.Background())
	cancel() // cancel before traversal starts

	if _, err := tr.Run(ctx, "example.com"); err == nil {
		t.Fatal("expected error when context is cancelled")
	}
}