package rfc2136

import (
	"strings"
	"sync"
	"sync/atomic"

	"github.com/miekg/dns"
)

// recordStore is the in-memory store of dynamic records. It is keyed
// by canonical (lowercased, trailing-dot) owner name, then by RR type,
// to a slice of RRs forming the RRset.
//
// Concurrency: a single RWMutex covers all mutations. Read paths
// (queries) are far more frequent than writes (UPDATEs), so the
// RW split is the right choice. The lock is held only across the
// O(1) map operations; nothing slow happens under the lock.
//
// Serial: every successful mutation bumps the atomic generation
// counter. The synthetic SOA uses this counter so resolvers see a
// fresh serial after every UPDATE — even though we don't run AXFR
// against secondaries, future use-cases (DNS-NOTIFY, debugging,
// cache invalidation) benefit from a monotonic serial.
type recordStore struct {
	mu sync.RWMutex
	// rrs[name][rrtype] = RRset (slice). The map-of-maps shape means
	// per-type lookups are O(1), and "name exists at all" is O(1) on
	// the outer map.
	rrs map[string]map[uint16][]dns.RR

	gen atomic.Uint64
}

func newStore() *recordStore {
	return &recordStore{rrs: make(map[string]map[uint16][]dns.RR)}
}

// generation returns the current monotonic counter. Used for SOA serial.
func (s *recordStore) generation() uint64 {
	return s.gen.Load()
}

// canon normalises a name to the store's internal form: lowercase +
// trailing dot. miekg/dns sometimes returns names without the dot;
// always passing through this keeps everything consistent.
func canon(name string) string {
	return strings.ToLower(dns.Fqdn(name))
}

// Add inserts a single RR into the store. Duplicate RRs (same owner,
// type, AND rdata) are silently de-duplicated — matches RFC 2136 §3.4.2.2
// behavior for `add` semantics.
func (s *recordStore) Add(rr dns.RR) {
	s.mu.Lock()
	defer s.mu.Unlock()

	name := canon(rr.Header().Name)
	rtype := rr.Header().Rrtype

	byType, ok := s.rrs[name]
	if !ok {
		byType = make(map[uint16][]dns.RR)
		s.rrs[name] = byType
	}

	// De-duplicate by string comparison of the rdata-bearing form.
	rrStr := rr.String()
	for _, existing := range byType[rtype] {
		if existing.String() == rrStr {
			return
		}
	}
	byType[rtype] = append(byType[rtype], rr)
	s.gen.Add(1)
}

// RemoveRRset deletes ALL records of the given (name, type) — i.e.
// "delete the RRset" semantics from RFC 2136 §3.4.2.3. If the name
// has no records left after this, the name entry is reaped so
// NameExists returns false.
func (s *recordStore) RemoveRRset(name string, rtype uint16) {
	s.mu.Lock()
	defer s.mu.Unlock()

	name = canon(name)
	byType, ok := s.rrs[name]
	if !ok {
		return
	}
	if _, hadType := byType[rtype]; !hadType {
		return
	}
	delete(byType, rtype)
	if len(byType) == 0 {
		delete(s.rrs, name)
	}
	s.gen.Add(1)
}

// RemoveRR deletes one specific RR (matching owner, type, and rdata).
// This implements RFC 2136 §3.4.2.4 "delete an RR from an RRset".
func (s *recordStore) RemoveRR(rr dns.RR) {
	s.mu.Lock()
	defer s.mu.Unlock()

	name := canon(rr.Header().Name)
	rtype := rr.Header().Rrtype
	byType, ok := s.rrs[name]
	if !ok {
		return
	}
	rrs := byType[rtype]
	target := rr.String()
	for i, existing := range rrs {
		if existing.String() == target {
			byType[rtype] = append(rrs[:i], rrs[i+1:]...)
			if len(byType[rtype]) == 0 {
				delete(byType, rtype)
			}
			if len(byType) == 0 {
				delete(s.rrs, name)
			}
			s.gen.Add(1)
			return
		}
	}
}

// RemoveName deletes all records for an owner name (§3.4.2.3 "delete
// all RRsets from a name").
func (s *recordStore) RemoveName(name string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	name = canon(name)
	if _, ok := s.rrs[name]; !ok {
		return
	}
	delete(s.rrs, name)
	s.gen.Add(1)
}

// Lookup returns the RRset for (name, rtype). Returns nil for both
// "name doesn't exist" and "name exists with other types but not this
// one" — use NameExists to distinguish NODATA from NXDOMAIN.
//
// The returned slice is a copy so callers can freely mutate it without
// affecting store state.
func (s *recordStore) Lookup(name string, rtype uint16) []dns.RR {
	s.mu.RLock()
	defer s.mu.RUnlock()

	byType, ok := s.rrs[canon(name)]
	if !ok {
		return nil
	}
	rrs := byType[rtype]
	if len(rrs) == 0 {
		return nil
	}
	out := make([]dns.RR, len(rrs))
	copy(out, rrs)
	return out
}

// NameExists reports whether ANY records exist for the given name.
// Used to distinguish NODATA (name exists, no records of asked type)
// from NXDOMAIN (name doesn't exist at all).
func (s *recordStore) NameExists(name string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	_, ok := s.rrs[canon(name)]
	return ok
}