package dns

import (
	"sync"
	"time"

	"github.com/miekg/dns"
)

type cacheKey struct {
	Name  string
	Qtype uint16
	Class uint16
}

type cacheEntry struct {
	msg       *dns.Msg
	expiresAt time.Time
}

// Cache is a thread-safe in-memory DNS response cache with TTL-based expiry.
type Cache struct {
	mu      sync.RWMutex
	entries map[cacheKey]*cacheEntry
}

// NewCache creates an empty DNS cache.
func NewCache() *Cache {
	return &Cache{
		entries: make(map[cacheKey]*cacheEntry),
	}
}

// Get returns a cached response if it exists and has not expired.
func (c *Cache) Get(name string, qtype, class uint16) *dns.Msg {
	c.mu.RLock()
	defer c.mu.RUnlock()

	key := cacheKey{Name: name, Qtype: qtype, Class: class}
	entry, ok := c.entries[key]
	if !ok || time.Now().After(entry.expiresAt) {
		return nil
	}
	return entry.msg
}

// Set stores a DNS response in the cache with the given TTL.
func (c *Cache) Set(name string, qtype, class uint16, msg *dns.Msg, ttl time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()

	key := cacheKey{Name: name, Qtype: qtype, Class: class}
	c.entries[key] = &cacheEntry{
		msg:       msg.Copy(),
		expiresAt: time.Now().Add(ttl),
	}

	// Lazy eviction: clean up expired entries if cache is growing.
	if len(c.entries) > 1000 {
		now := time.Now()
		for k, v := range c.entries {
			if now.After(v.expiresAt) {
				delete(c.entries, k)
			}
		}
	}
}