package resolver

import (
	"context"
	"fmt"
	"math"
	"regexp"
	"strings"
	"sync/atomic"
	"time"

	"github.com/0xERR0R/blocky/cache"
	"github.com/0xERR0R/blocky/config"
	"github.com/0xERR0R/blocky/evt"
	"github.com/0xERR0R/blocky/metrics"
	"github.com/0xERR0R/blocky/model"
	"github.com/0xERR0R/blocky/util"
	expirationcache "github.com/0xERR0R/expiration-cache"

	"github.com/0xERR0R/blocky/cache/prefetching"
	"github.com/miekg/dns"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/sirupsen/logrus"
)

const (
	defaultCachingCleanUpInterval = 5 * time.Second

	// cachedReason is the response reason used for cache hits.
	cachedReason = "CACHED"
)

// CacheDecorator optionally wraps the result cache (e.g., with Redis sync).
type CacheDecorator func(cache.ExpiringCache[[]byte]) (cache.ExpiringCache[[]byte], error)

//nolint:gochecknoglobals
var (
	cacheHits = promauto.With(metrics.Reg).NewCounter(
		prometheus.CounterOpts{
			Name: "blocky_cache_hits_total",
			Help: "Cache hit counter",
		},
	)
	cacheMisses = promauto.With(metrics.Reg).NewCounter(
		prometheus.CounterOpts{
			Name: "blocky_cache_misses_total",
			Help: "Cache miss counter",
		},
	)
)

// CachingResolver caches answers from dns queries with their TTL time,
// to avoid external resolver calls for recurrent queries
type CachingResolver struct {
	configurable[*config.Caching]
	NextResolver
	typed

	emitMetricEvents bool // disabled by Bootstrap

	// prefetchDO makes prefetch reloads request DNSSEC records (DO bit) so a
	// reload can't replace a signed entry with an unsigned one. Set from
	// dnssec.validate — unsigned caches shouldn't grow with RRSIGs.
	prefetchDO bool

	resultCache cache.ExpiringCache[[]byte]

	compiledExclusions []*regexp.Regexp
}

// NewCachingResolver creates a new resolver instance
func NewCachingResolver(ctx context.Context,
	cfg config.Caching,
	dnssecCfg config.DNSSEC,
	decorator CacheDecorator,
) (*CachingResolver, error) {
	return newCachingResolver(ctx, cfg, dnssecCfg, decorator, true)
}

func newCachingResolver(ctx context.Context,
	cfg config.Caching,
	dnssecCfg config.DNSSEC,
	decorator CacheDecorator,
	emitMetricEvents bool,
) (*CachingResolver, error) {
	c := &CachingResolver{
		configurable:     withConfig(&cfg),
		typed:            withType("caching"),
		emitMetricEvents: emitMetricEvents,
		prefetchDO:       dnssecCfg.Validate,
	}

	configureCaches(ctx, c, &cfg)
	err := configureExclusions(c, &cfg)

	if decorator != nil {
		decorated, err := decorator(c.resultCache)
		if err != nil {
			return nil, fmt.Errorf("failed to apply cache decorator: %w", err)
		}

		c.resultCache = decorated
	}

	return c, err
}

func configureCaches(ctx context.Context, c *CachingResolver, cfg *config.Caching) {
	options := expirationcache.Options{
		CleanupInterval: defaultCachingCleanUpInterval,
		MaxSize:         uint(cfg.MaxItemsCount),
		Shards:          cache.ShardCount(),
		OnCacheHitFn: func(key string) {
			cacheHits.Inc()
		},
		OnCacheMissFn: func(key string) {
			cacheMisses.Inc()
		},
		OnAfterPutFn: func(newSize int) {
			c.publishMetricsIfEnabled(evt.CachingResultCacheChanged, newSize)
		},
	}

	if cfg.Prefetching {
		prefetchingOptions := prefetching.PrefetchingOptions[[]byte]{
			Options:               options,
			PrefetchExpires:       time.Duration(cfg.PrefetchExpires),
			PrefetchThreshold:     cfg.PrefetchThreshold,
			PrefetchMaxItemsCount: cfg.PrefetchMaxItemsCount,
			ReloadFn:              c.reloadCacheEntry,
			OnPrefetchAfterPut: func(newSize int) {
				c.publishMetricsIfEnabled(evt.CachingDomainsToPrefetchCountChanged, newSize)
			},
			OnPrefetchEntryReloaded: func(key string) {
				c.publishMetricsIfEnabled(evt.CachingDomainPrefetched, key)
			},
			OnPrefetchCacheHit: func(key string) {
				c.publishMetricsIfEnabled(evt.CachingPrefetchCacheHit, key)
			},
		}

		c.resultCache = prefetching.NewPrefetchingCache(ctx, prefetchingOptions)
	} else {
		c.resultCache = expirationcache.NewCache[[]byte](ctx, options)
	}
}

func configureExclusions(c *CachingResolver, cfg *config.Caching) error {
	compiled := []*regexp.Regexp{}
	for _, expStr := range cfg.Exclude {
		if !strings.HasPrefix(expStr, "/") || !strings.HasSuffix(expStr, "/") {
			return fmt.Errorf("cache exclusion configuration '%s' fail because of missing slashes", expStr)
		}
		re, err := regexp.Compile(strings.TrimSpace(expStr[1 : len(expStr)-1]))
		if err != nil {
			return fmt.Errorf("cache exclusion configuration '%s' fail because '%s'", expStr, err.Error())
		}
		compiled = append(compiled, re)
	}
	c.compiledExclusions = compiled

	return nil
}

func (r *CachingResolver) reloadCacheEntry(ctx context.Context, cacheKey string) (*[]byte, time.Duration) {
	qType, domainName := util.ExtractCacheKey(cacheKey)
	ctx, logger := r.log(ctx)

	logger.Debugf("prefetching '%s' (%s)", util.Obfuscate(domainName), qType)

	req := newRequest(dns.Fqdn(domainName), qType)

	// Prefetch reloads bypass the resolvers above the cache, including the DNSSEC
	// resolver that sets DO on normal queries. When validation is enabled, request
	// DNSSEC records so a reload can't replace a signed entry with an unsigned one
	// (which would then fail re-validation as bogus on the next hit).
	if r.prefetchDO {
		req.Req.SetEdns0(ednsUDPSize, true)
	}

	response, err := r.next.Resolve(ctx, req)
	if err != nil {
		util.LogOnError(ctx, fmt.Sprintf("can't prefetch '%s' ", domainName), err)

		return nil, 0
	}

	// only refresh entries the normal put path would cache: upstream-derived,
	// successful, not truncated and without the CD flag (mirrors putInCache).
	if !isCacheableResponseType(response.RType) ||
		response.Res.Rcode != dns.RcodeSuccess || !isResponseCacheable(response.Res) {
		return nil, 0
	}

	// clamp the record TTLs before packing so the stored (and Redis-synced) bytes
	// carry the same TTLs the normal put path produces.
	ttl := r.adjustTTLs(response.Res.Answer)

	packed, err := packForCache(ctx, response.Res)
	if err != nil {
		return nil, 0
	}

	return &packed, ttl
}

// LogConfig implements `config.Configurable`.
func (r *CachingResolver) LogConfig(logger *logrus.Entry) {
	r.cfg.LogConfig(logger)

	logger.Infof("cache entries = %d", r.resultCache.TotalCount())
}

// Resolve checks if the current query should use the cache and if the result is already in
// the cache and returns it or delegates to the next resolver
func (r *CachingResolver) Resolve(ctx context.Context, request *model.Request) (response *model.Response, err error) {
	ctx, logger := r.log(ctx)

	if !r.IsEnabled() || !r.isRequestCacheable(request) {
		logger.Debug("skip cache")

		return r.next.Resolve(ctx, request)
	}

	for _, question := range request.Req.Question {
		domain := util.ExtractDomain(question)
		cacheKey := util.GenerateCacheKey(dns.Type(question.Qtype), domain)
		logger := logger.WithField(logFieldDomain, util.Obfuscate(domain))

		val, ttl := r.getFromCache(logger, cacheKey)

		if val != nil {
			logger.Debug("domain is cached")

			val.SetRcode(request.Req, val.Rcode)

			// Adjust TTL
			r.setTTLInCachedResponse(val, ttl)

			if val.Rcode == dns.RcodeSuccess {
				return &model.Response{Res: val, RType: model.ResponseTypeCACHED, Reason: cachedReason}, nil
			}

			return &model.Response{Res: val, RType: model.ResponseTypeCACHED, Reason: "CACHED NEGATIVE"}, nil
		}

		logger.WithField("next_resolver", Name(r.next)).Trace("not in cache: go to next resolver")
		response, err = r.next.Resolve(ctx, request)

		if err == nil {
			cacheTTL := r.adjustTTLs(response.Res.Answer)
			r.putInCache(ctx, cacheKey, response, cacheTTL)
		} else {
			return nil, err
		}
	}

	return response, nil
}

func (r *CachingResolver) getFromCache(logger *logrus.Entry, key string) (*dns.Msg, time.Duration) {
	val, ttl := r.resultCache.Get(key)
	if val == nil {
		return nil, 0
	}

	res := new(dns.Msg)

	err := res.Unpack(*val)
	if err != nil {
		logger.Error("can't unpack cached entry. Cache malformed?", err)

		return nil, 0
	}

	return res, ttl
}

// setTTLInCachedResponse ages all records of a cached message by the time the entry has
// already spent in the cache, so that every section counts down instead of repeating the
// TTLs the upstream sent. `ttl` is the entry's remaining lifetime.
func (r *CachingResolver) setTTLInCachedResponse(resp *dns.Msg, ttl time.Duration) {
	// Mirrors putInCache: only a successful, non-empty answer is stored with the smallest
	// answer TTL (see adjustTTLs). Everything else -- NXDOMAIN (even one carrying a CNAME
	// chain) and NODATA -- is stored with the negative cache time. That baseline is what
	// the remaining lifetime has been counting down from.
	baseTTL := r.cfg.CacheTimeNegative.SecondsU32()

	if resp.Rcode == dns.RcodeSuccess && len(resp.Answer) > 0 {
		baseTTL = uint32(math.MaxInt32)
		for _, rr := range resp.Answer {
			baseTTL = min(baseTTL, rr.Header().Ttl)
		}
	}

	var elapsed uint32
	if remaining := uint32(ttl.Seconds()); baseTTL > remaining {
		elapsed = baseTTL - remaining
	}

	for _, section := range [][]dns.RR{resp.Answer, resp.Ns, resp.Extra} {
		for _, rr := range section {
			// An OPT record's TTL field carries flags and the extended rcode, not a lifetime.
			if rr.Header().Rrtype == dns.TypeOPT {
				continue
			}

			rr.Header().Ttl = max(rr.Header().Ttl, elapsed) - elapsed
		}
	}
}

// isRequestCacheable returns true if the request should be cached
func (r *CachingResolver) isRequestCacheable(request *model.Request) bool {
	// don't cache response if name ends with any exclution
	if questionsMatchAnyExcludedElement(request.Req.Question, r.compiledExclusions) {
		return false
	}
	// don't cache responses with EDNS Client Subnet option with masks that include more than one client
	if so := util.GetEdns0Option[*dns.EDNS0_SUBNET](request.Req); so != nil {
		if (so.Family == ecsFamilyIPv4 && so.SourceNetmask != ecsMaskIPv4) ||
			(so.Family == ecsFamilyIPv6 && so.SourceNetmask != ecsMaskIPv6) {
			return false
		}
	}

	return true
}

func questionsMatchAnyExcludedElement(questions []dns.Question, exclutions []*regexp.Regexp) bool {
	for _, q := range questions {
		if matchAnyElementOfArray(q.Name[:len(q.Name)-1], exclutions) {
			return true
		}
	}

	return false
}

func matchAnyElementOfArray(givingText string, arr []*regexp.Regexp) bool {
	loweredGivingText := strings.ToLower(givingText)
	for _, s := range arr {
		if s.MatchString(loweredGivingText) {
			return true
		}
	}

	return false
}

// isResponseCacheable returns true if the response is not truncated and its CD flag isn't set.
func isResponseCacheable(msg *dns.Msg) bool {
	// we don't cache truncated responses and responses with CD flag
	return !msg.Truncated && !msg.CheckingDisabled
}

// packForCache copies the message, strips EDNS0 OPT records (which must never be
// cached) and returns the packed wire bytes. Both the normal put path and the
// prefetch reload path use it so cached and prefetched entries are byte-identical.
func packForCache(ctx context.Context, msg *dns.Msg) ([]byte, error) {
	msgCopy := msg.Copy()
	util.RemoveEdns0Record(msgCopy)

	packed, err := msgCopy.Pack()
	util.LogOnError(ctx, "error on packing", err)

	return packed, err
}

// isCacheableResponseType reports whether the response originates from the general
// upstreams, directly (RESOLVED) or via DNS64 synthesis (SYNTHESIZED). Only such
// answers may be cached: the rebinding protection resolver sits above the cache and
// inspects CACHED responses, so caching answers from trusted local sources
// (conditional upstream, special-use domains) would re-label them as CACHED on hits
// and subject them to inspection.
func isCacheableResponseType(rType model.ResponseType) bool {
	return rType == model.ResponseTypeRESOLVED || rType == model.ResponseTypeSYNTHESIZED
}

func (r *CachingResolver) putInCache(
	ctx context.Context, cacheKey string, response *model.Response, ttl time.Duration,
) {
	if !isCacheableResponseType(response.RType) {
		return
	}

	packed, err := packForCache(ctx, response.Res)
	if err != nil {
		return
	}

	if response.Res.Rcode == dns.RcodeSuccess && isResponseCacheable(response.Res) {
		// put value into cache
		r.resultCache.Put(cacheKey, &packed, ttl)
	} else if response.Res.Rcode == dns.RcodeNameError {
		if r.cfg.CacheTimeNegative.IsAboveZero() {
			// put negative cache if result code is NXDOMAIN
			r.resultCache.Put(cacheKey, &packed, r.cfg.CacheTimeNegative.ToDuration())
		}
	}
}

// adjustTTLs calculates and returns the min TTL (considers also the min and max cache time)
// for all records from answer or a negative cache time for empty answer
// adjust the TTL in the answer header accordingly
func (r *CachingResolver) adjustTTLs(answer []dns.RR) (ttl time.Duration) {
	minTTL := uint32(math.MaxInt32)

	if len(answer) == 0 {
		return r.cfg.CacheTimeNegative.ToDuration()
	}

	for _, a := range answer {
		// if TTL < mitTTL -> adjust the value, set minTTL
		if r.cfg.MinCachingTime.IsAboveZero() {
			if atomic.LoadUint32(&a.Header().Ttl) < r.cfg.MinCachingTime.SecondsU32() {
				atomic.StoreUint32(&a.Header().Ttl, r.cfg.MinCachingTime.SecondsU32())
			}
		}

		if r.cfg.MaxCachingTime.IsAboveZero() {
			if atomic.LoadUint32(&a.Header().Ttl) > r.cfg.MaxCachingTime.SecondsU32() {
				atomic.StoreUint32(&a.Header().Ttl, r.cfg.MaxCachingTime.SecondsU32())
			}
		}

		headerTTL := atomic.LoadUint32(&a.Header().Ttl)
		if minTTL > headerTTL {
			minTTL = headerTTL
		}
	}

	return time.Duration(minTTL) * time.Second
}

func (r *CachingResolver) publishMetricsIfEnabled(event string, val any) {
	if r.emitMetricEvents {
		evt.Bus().Publish(event, val)
	}
}

func (r *CachingResolver) FlushCaches(ctx context.Context) {
	_, logger := r.log(ctx)

	logger.Debug("flush caches")
	r.resultCache.Clear()
}