// Package fingerprint identifies DNS server software by querying the
// version.bind name in the CHAOS class.  Many authoritative and recursive DNS
// servers respond with a version string (e.g. "BIND 9.18.1-1") that can be
// used to identify the software and version in use.
//
// Queries are cached per server IP so that repeated lookups within a single
// traversal run do not incur extra network round-trips.
package fingerprint

import (
	"context"
	"net"
	"sync"
	"time"

	miekgdns "github.com/miekg/dns"
)

const defaultTimeout = 2 * time.Second

// Fingerprinter queries DNS servers for their software version via the
// version.bind CHAOS TXT query. Results are cached per server IP.
type Fingerprinter struct {
	mu       sync.Mutex
	cache    map[string]string
	timeout  time.Duration
	exchange func(ctx context.Context, addr string, m *miekgdns.Msg) (*miekgdns.Msg, error)
}

// New returns a Fingerprinter with a 2-second per-query timeout.
func New() *Fingerprinter {
	return NewWithTimeout(defaultTimeout)
}

// NewWithTimeout returns a Fingerprinter using the given per-query timeout.
func NewWithTimeout(timeout time.Duration) *Fingerprinter {
	return &Fingerprinter{
		cache:   make(map[string]string),
		timeout: timeout,
	}
}

// Query returns the version string for ip, or "" if the server doesn't
// respond or doesn't support the version.bind CHAOS query.
// Results are cached: subsequent calls for the same IP return immediately.
func (f *Fingerprinter) Query(ctx context.Context, ip net.IP) string {
	key := ip.String()

	f.mu.Lock()
	if v, ok := f.cache[key]; ok {
		f.mu.Unlock()
		return v
	}
	f.mu.Unlock()

	version := f.probe(ctx, ip)

	f.mu.Lock()
	f.cache[key] = version
	f.mu.Unlock()

	return version
}

// FingerprintAll queries all ips concurrently and returns a map of
// IP string → version string. IPs that don't respond map to "".
// Already-cached IPs are returned from cache without a network round-trip.
func (f *Fingerprinter) FingerprintAll(ctx context.Context, ips []net.IP) map[string]string {
	results := make(map[string]string, len(ips))

	var (
		wg      sync.WaitGroup
		mu      sync.Mutex
		toQuery []net.IP
	)

	f.mu.Lock()
	for _, ip := range ips {
		key := ip.String()
		if v, ok := f.cache[key]; ok {
			results[key] = v
		} else {
			toQuery = append(toQuery, ip)
		}
	}
	f.mu.Unlock()

	for _, ip := range toQuery {
		wg.Add(1)
		go func(ip net.IP) {
			defer wg.Done()
			version := f.probe(ctx, ip)
			key := ip.String()

			f.mu.Lock()
			f.cache[key] = version
			f.mu.Unlock()

			mu.Lock()
			results[key] = version
			mu.Unlock()
		}(ip)
	}

	wg.Wait()
	return results
}

// probe sends a version.bind CHAOS TXT query and returns the version string,
// or "" on any error or non-success response.
func (f *Fingerprinter) probe(ctx context.Context, ip net.IP) string {
	m := new(miekgdns.Msg)
	m.SetQuestion("version.bind.", miekgdns.TypeTXT)
	m.Question[0].Qclass = miekgdns.ClassCHAOS
	m.RecursionDesired = false

	target := net.JoinHostPort(ip.String(), "53")

	queryCtx, cancel := context.WithTimeout(ctx, f.timeout)
	defer cancel()

	var resp *miekgdns.Msg
	var err error

	if f.exchange != nil {
		resp, err = f.exchange(queryCtx, target, m)
	} else {
		client := &miekgdns.Client{
			Net:     "udp",
			Timeout: f.timeout,
		}
		resp, _, err = client.ExchangeContext(queryCtx, m, target)
	}

	if err != nil || resp == nil || resp.Rcode != miekgdns.RcodeSuccess {
		return ""
	}

	for _, rr := range resp.Answer {
		if txt, ok := rr.(*miekgdns.TXT); ok && len(txt.Txt) > 0 {
			return txt.Txt[0]
		}
	}
	return ""
}