package engine

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

	"github.com/miekg/dns"
)

// ReverseCache maps client IPs to reverse-DNS names so the ACL can match on
// names without a synchronous PTR lookup on the query path.
//
// A miss schedules a background lookup and, for the duration of that lookup,
// matches on nothing. That is the right trade: a blocking PTR on the hot path
// would make a broken reverse zone a latency incident, whereas a miss only
// means an ACL rule keyed on a name does not fire for the first few queries.
type ReverseCache struct {
	mu      sync.RWMutex
	entries map[string]reverseEntry
	// negative caches the "no PTR exists" answer so we do not re-query for
	// every query from the same client.
	max      int
	ttl      time.Duration
	negative time.Duration
	resolver *dns.Client
	names    string
	timeout  time.Duration
	inflight map[string]bool
}

type reverseEntry struct {
	names    []string
	fetched  time.Time
	negative bool
}

// ReverseOptions configures a ReverseCache.
type ReverseOptions struct {
	// Max is the number of cached clients.
	Max int
	// TTL is how long a successful PTR is trusted.
	TTL time.Duration
	// NegativeTTL is how long a "no PTR" answer is trusted.
	NegativeTTL time.Duration
	// Nameserver is the server to ask, e.g. "9.9.9.9:53". Empty uses the
	// system resolver via net.LookupAddr.
	Nameserver string
	// Timeout bounds one PTR lookup.
	Timeout time.Duration
}

// NewReverseCache builds a cache.
func NewReverseCache(opt ReverseOptions) *ReverseCache {
	if opt.Max <= 0 {
		opt.Max = 16384
	}
	if opt.TTL <= 0 {
		opt.TTL = 10 * time.Minute
	}
	if opt.NegativeTTL <= 0 {
		opt.NegativeTTL = 30 * time.Minute
	}
	if opt.Timeout <= 0 {
		opt.Timeout = 2 * time.Second
	}
	rc := &ReverseCache{
		entries:  make(map[string]reverseEntry, 1024),
		max:      opt.Max,
		ttl:      opt.TTL,
		negative: opt.NegativeTTL,
		names:    opt.Nameserver,
		timeout:  opt.Timeout,
		inflight: make(map[string]bool),
	}
	if opt.Nameserver != "" {
		rc.resolver = &dns.Client{Net: "udp", Timeout: opt.Timeout}
	}
	return rc
}

// Names returns the cached reverse names for an IP, and whether the entry was
// present. The second return is false for a miss (including a cached negative).
func (r *ReverseCache) Names(ip string) ([]string, bool) {
	if r == nil {
		return nil, false
	}
	r.mu.RLock()
	e, ok := r.entries[ip]
	r.mu.RUnlock()
	if !ok {
		r.Prime(ip)
		return nil, false
	}
	ttl := r.ttl
	if e.negative {
		ttl = r.negative
	}
	if time.Since(e.fetched) > ttl {
		r.Prime(ip)
		return e.names, true // serve the stale value while refreshing
	}
	return e.names, !e.negative
}

// Prime schedules a background PTR lookup for ip, at most one in flight.
func (r *ReverseCache) Prime(ip string) {
	if r == nil || ip == "" {
		return
	}
	r.mu.Lock()
	if r.inflight[ip] {
		r.mu.Unlock()
		return
	}
	r.inflight[ip] = true
	r.mu.Unlock()
	go r.lookup(ip)
}

func (r *ReverseCache) lookup(ip string) {
	defer func() {
		r.mu.Lock()
		delete(r.inflight, ip)
		r.mu.Unlock()
	}()
	names, err := r.query(ip)
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.entries) >= r.max {
		// Simple bound: drop the oldest quarter. A full LRU here would cost
		// more than it saves, since the working set is usually small.
		cutoff := time.Now().Add(-r.ttl)
		removed := 0
		for k, e := range r.entries {
			if e.fetched.Before(cutoff) || removed >= r.max/4 {
				delete(r.entries, k)
				removed++
			}
		}
		if removed == 0 {
			for k := range r.entries {
				delete(r.entries, k)
				if removed >= r.max/4 {
					break
				}
			}
		}
	}
	r.entries[ip] = reverseEntry{
		names:    names,
		fetched:  time.Now(),
		negative: err != nil || len(names) == 0,
	}
}

func (r *ReverseCache) query(ip string) ([]string, error) {
	if r.resolver == nil {
		ctx, cancel := context.WithTimeout(context.Background(), r.timeout)
		defer cancel()
		return net.DefaultResolver.LookupAddr(ctx, ip)
	}
	name, err := dns.ReverseAddr(ip)
	if err != nil {
		return nil, err
	}
	m := new(dns.Msg)
	m.SetQuestion(name, dns.TypePTR)
	m.RecursionDesired = true
	resp, _, err := r.resolver.Exchange(m, r.names)
	if err != nil {
		return nil, err
	}
	var out []string
	for _, rr := range resp.Answer {
		if ptr, ok := rr.(*dns.PTR); ok {
			out = append(out, ptr.Ptr)
		}
	}
	return out, nil
}

// Len is the number of cached clients.
func (r *ReverseCache) Len() int {
	if r == nil {
		return 0
	}
	r.mu.RLock()
	defer r.mu.RUnlock()
	return len(r.entries)
}

// Flush empties the cache.
func (r *ReverseCache) Flush() {
	if r == nil {
		return
	}
	r.mu.Lock()
	r.entries = make(map[string]reverseEntry, 1024)
	r.mu.Unlock()
}