// Package localzone serves authoritative data from DNS master files: split
// DNS setups, internal service records and the odd "answer this one name
// locally" need, all without a second nameserver to run.
package localzone

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync/atomic"
	"time"

	"github.com/miekg/dns"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/zonefile"
)

// Store is an immutable, compiled set of local zones.
type Store struct {
	// zones maps a lowercased apex (with trailing dot) to its compiled zone.
	zones map[string]*zone
	// count is the number of records held, for the UI.
	count int
	// files lists the sources that produced this store.
	files []string
	// builtAt is the compile timestamp.
	builtAt time.Time
}

type zone struct {
	origin string
	// exact is keyed by "name|type".
	exact map[string][]dns.RR
	// wild holds "*.name" entries.
	wild map[string]map[uint16][]dns.RR
	// soa is the zone's SOA, used for NXDOMAIN and NODATA authority.
	soa *dns.SOA
	// ns holds the apex NS set.
	ns map[uint16][]dns.RR
	// delegations are the names served by NS records (a cut below the apex).
	delegations map[string]bool
	// files maps an index key ("name|TYPE", or "*./name|TYPE" for a
	// wildcard) to the file the record was parsed from. Only the admin
	// API's zone browser reads it, so the query path pays nothing for it.
	files map[string]string
	// sources are the files this zone was compiled from, in file order.
	sources []string
	minTTL  uint32
}

// wildKey is the index key a wildcard record is stored under.
func wildKey(name, rrtype string) string { return "*./" + name + "|" + rrtype }

// source reports the file a record was parsed from ("" when unknown).
func (z *zone) source(key string) string { return z.files[key] }

// Options tunes zone loading.
type Options struct {
	// Origins overrides the per-file origin derivation.
	Origins map[string]string
	// MaxFiles bounds how many files are read (glob results included).
	MaxFiles int
}

// Load compiles every master file in paths. A path may be a glob.
func Load(paths []string, opt Options) (*Store, error) {
	files, err := expand(paths, opt.MaxFiles)
	if err != nil {
		return nil, err
	}
	s := &Store{
		zones:   make(map[string]*zone, len(files)),
		files:   files,
		builtAt: time.Now(),
	}
	for _, f := range files {
		origin := zonefile.OriginFromPath(f)
		if o, ok := opt.Origins[f]; ok {
			origin = dns.Fqdn(o)
		}
		z, err := zonefile.ParseFile(f, origin)
		if err != nil {
			return nil, fmt.Errorf("localzone: %w", err)
		}
		apex := strings.ToLower(z.Origin)
		zz, ok := s.zones[apex]
		if !ok {
			zz = &zone{
				origin:      apex,
				exact:       make(map[string][]dns.RR),
				wild:        make(map[string]map[uint16][]dns.RR),
				delegations: make(map[string]bool),
				files:       make(map[string]string),
				ns:          make(map[uint16][]dns.RR),
				minTTL:      0xFFFFFFFF,
			}
			s.zones[apex] = zz
		}
		zz.noteSource(f)
		for _, rec := range z.Records {
			if add(zz, rec) {
				s.count++
			}
		}
	}
	for _, zz := range s.zones {
		if zz.minTTL == 0xFFFFFFFF {
			zz.minTTL = 60
		}
	}
	return s, nil
}

func add(z *zone, rec zonefile.Record) bool {
	hdr := rec.RR.Header()
	name := strings.ToLower(dns.Fqdn(hdr.Name))
	// SOA is zone metadata, not a queryable record.
	if soa, ok := rec.RR.(*dns.SOA); ok {
		z.soa = soa
		return false
	}
	key := name
	isWild := false
	if strings.HasPrefix(name, "*.") {
		key = name[2:]
		isWild = true
	}
	// A record must live inside its zone; silently dropping out-of-zone data
	// would be a surprise, so refuse it loudly instead.
	if !inZone(key, z.origin) {
		return false
	}
	// NS at or below the apex marks a delegation: the data below belongs to
	// another server, so we must not answer for it.
	if hdr.Rrtype == dns.TypeNS && !isWild {
		if key == z.origin {
			z.ns[hdr.Rrtype] = append(z.ns[hdr.Rrtype], rec.RR)
		} else {
			z.delegations[key] = true
		}
	}
	if hdr.Ttl < z.minTTL {
		z.minTTL = hdr.Ttl
	}
	if isWild {
		m, ok := z.wild[key]
		if !ok {
			m = make(map[uint16][]dns.RR)
			z.wild[key] = m
		}
		m[hdr.Rrtype] = append(m[hdr.Rrtype], rec.RR)
		z.files[wildKey(key, dns.TypeToString[hdr.Rrtype])] = rec.File
		return true
	}
	k := name + "|" + dns.TypeToString[hdr.Rrtype]
	z.exact[k] = append(z.exact[k], rec.RR)
	z.files[k] = rec.File
	return true
}

// noteSource records a contributing file, keeping the order files were read.
func (z *zone) noteSource(file string) {
	if file == "" {
		return
	}
	for _, f := range z.sources {
		if f == file {
			return
		}
	}
	z.sources = append(z.sources, file)
}

func inZone(name, origin string) bool {
	if origin == "." {
		return true
	}
	return name == origin || strings.HasSuffix(name, "."+origin)
}

// expand turns configured paths (possibly globs) into a concrete file list.
func expand(paths []string, max int) ([]string, error) {
	var out []string
	seen := map[string]bool{}
	for _, p := range paths {
		if strings.ContainsAny(p, "*?[") {
			matches, err := filepath.Glob(p)
			if err != nil {
				return nil, fmt.Errorf("localzone: bad glob %q: %w", p, err)
			}
			sort.Strings(matches)
			for _, m := range matches {
				if !seen[m] {
					seen[m] = true
					out = append(out, m)
				}
			}
			continue
		}
		if !seen[p] {
			seen[p] = true
			out = append(out, p)
		}
	}
	// A path that is a directory contributes its *.zone files, which is the
	// natural thing to point at.
	var expanded []string
	for _, p := range out {
		fi, err := os.Stat(p)
		if err != nil {
			return nil, fmt.Errorf("localzone: %w", err)
		}
		if fi.IsDir() {
			entries, err := filepath.Glob(filepath.Join(p, "*.zone"))
			if err != nil {
				return nil, err
			}
			sort.Strings(entries)
			expanded = append(expanded, entries...)
			continue
		}
		expanded = append(expanded, p)
	}
	if max > 0 && len(expanded) > max {
		expanded = expanded[:max]
	}
	return expanded, nil
}

// Result is the outcome of a local lookup.
type Result struct {
	// Found is true when the store has an answer (including a CNAME).
	Found bool
	// Records is the RRset for the queried type, or a CNAME when the name
	// only has one.
	Records []dns.RR
	// Type is the type the records are for (may differ from the question for
	// a CNAME answer).
	Type uint16
	// Delegated is true when the name is served by another server.
	Delegated bool
	// Authority is the SOA for NXDOMAIN/NODATA responses.
	Authority *dns.SOA
	// TTL is the smallest TTL in the answer.
	TTL uint32
	// Zone is the apex of the zone that answered.
	Zone string
}

// Lookup resolves a question against the local store. It does not chase CNAMEs
// across the local store; the handler re-queries, which keeps the wildcard and
// delegation rules in one place.
func (s *Store) Lookup(q dns.Question) Result {
	if s == nil || len(s.zones) == 0 {
		return Result{}
	}
	name := strings.ToLower(dns.Fqdn(q.Name))
	qtype := q.Qtype
	// The cache and the resolver both skip the root and the local reverse
	// zones, but a local zone may legitimately define them, so check first.
	var best *zone
	for apex, z := range s.zones {
		if !inZone(name, apex) {
			continue
		}
		if best == nil || len(apex) > len(best.origin) {
			best = z
		}
	}
	if best == nil {
		return Result{}
	}
	if qtype == dns.TypeANY {
		return s.any(best, name)
	}
	// A delegation means "ask them, not me": the parent holds no data at or
	// below the cut, so it must neither answer nor claim NXDOMAIN. The NS
	// RRset at the cut is the one exception -- only the parent has it, and a
	// resolver needs it to learn the delegation (RFC 1034 §4.3.3).
	for d := name; d != ""; d = parentOf(d) {
		if !best.delegations[d] {
			continue
		}
		if d == name && qtype == dns.TypeNS {
			break
		}
		return Result{Found: false, Delegated: true, Zone: best.origin, Authority: best.soa, TTL: best.minTTL}
	}
	if rrs, ok := best.exact[name+"|"+dns.TypeToString[qtype]]; ok {
		return s.result(best, name, qtype, rrs)
	}
	// Wildcard: "*.zone" covers one label; RFC 1034 §4.3.3 wants the closest
	// enclosing wildcard, so try progressively shorter parents.
	if rrs, ok := closestWildcard(best, name, qtype); ok {
		// A wildcard answer is labelled with the name that was asked for.
		return s.result(best, name, qtype, synthesize(rrs, name))
	}
	// A CNAME at the name answers any other type.
	if qtype != dns.TypeCNAME {
		if rrs, ok := best.exact[name+"|CNAME"]; ok {
			return s.result(best, name, dns.TypeCNAME, rrs)
		}
	}
	// NODATA: the name exists but has no records of this type.
	if nameExists(best, name) {
		return Result{Found: true, Authority: best.soa, Zone: best.origin, TTL: best.minTTL}
	}
	return Result{Found: false, Authority: best.soa, Zone: best.origin, TTL: best.minTTL}
}

func (s *Store) result(z *zone, name string, qtype uint16, rrs []dns.RR) Result {
	min := uint32(0xFFFFFFFF)
	for _, rr := range rrs {
		if t := rr.Header().Ttl; t < min {
			min = t
		}
	}
	if min == 0xFFFFFFFF {
		min = z.minTTL
	}
	return Result{Found: true, Records: rrs, Type: qtype, TTL: min, Zone: z.origin, Authority: z.soa}
}

func (s *Store) any(z *zone, name string) Result {
	var out []dns.RR
	if rrs, ok := z.exact[name+"|A"]; ok {
		out = append(out, rrs...)
	}
	if rrs, ok := z.exact[name+"|AAAA"]; ok {
		out = append(out, rrs...)
	}
	if rrs, ok := z.exact[name+"|CNAME"]; ok {
		out = append(out, rrs...)
	}
	if rrs, ok := z.exact[name+"|MX"]; ok {
		out = append(out, rrs...)
	}
	if rrs, ok := z.exact[name+"|TXT"]; ok {
		out = append(out, rrs...)
	}
	if rrs, ok := z.exact[name+"|SRV"]; ok {
		out = append(out, rrs...)
	}
	if rrs, ok := z.exact[name+"|NS"]; ok {
		out = append(out, rrs...)
	}
	if len(out) == 0 {
		if rrs, ok := closestWildcardAny(z, name); ok {
			out = append(out, synthesize(rrs, name)...)
		}
	}
	if len(out) == 0 {
		return Result{Found: nameExists(z, name), Authority: z.soa, Zone: z.origin, TTL: z.minTTL}
	}
	return s.result(z, name, dns.TypeANY, out)
}

func closestWildcard(z *zone, name string, qtype uint16) ([]dns.RR, bool) {
	// A wildcard only fills in for a name that has no node of its own.
	if len(z.wild) == 0 || nameExists(z, name) {
		return nil, false
	}
	for parent := parentOf(name); parent != ""; parent = parentOf(parent) {
		if m, ok := z.wild[parent]; ok {
			if rrs, ok := m[qtype]; ok {
				return rrs, true
			}
		}
		// A real node (or a delegation) is the closest encloser, and only a
		// wildcard directly below it can answer: a wildcard never reaches
		// past a name that exists.
		if z.delegations[parent] || nameExists(z, parent) {
			return nil, false
		}
	}
	return nil, false
}

func closestWildcardAny(z *zone, name string) ([]dns.RR, bool) {
	if len(z.wild) == 0 || nameExists(z, name) {
		return nil, false
	}
	for parent := parentOf(name); parent != ""; parent = parentOf(parent) {
		if m, ok := z.wild[parent]; ok {
			var out []dns.RR
			for _, rrs := range m {
				out = append(out, rrs...)
			}
			if len(out) > 0 {
				return out, true
			}
		}
		if z.delegations[parent] || nameExists(z, parent) {
			return nil, false
		}
	}
	return nil, false
}

func parentOf(name string) string {
	i := strings.IndexByte(name, '.')
	if i < 0 {
		return ""
	}
	return name[i+1:]
}

// synthesize copies rrs with the owner name replaced by qname, because a
// wildcard answer is labelled with the name that was asked for. The store is
// shared read-only, so the records are copied instead of rewritten in place.
func synthesize(rrs []dns.RR, qname string) []dns.RR {
	out := make([]dns.RR, 0, len(rrs))
	for _, rr := range rrs {
		cp := dns.Copy(rr)
		if cp == nil {
			continue
		}
		cp.Header().Name = qname
		out = append(out, cp)
	}
	return out
}

func nameExists(z *zone, name string) bool {
	// The apex always exists: it carries the SOA and the NS set, and neither
	// is indexed as a queryable record.
	if name == z.origin {
		return true
	}
	prefix := name + "|"
	for k := range z.exact {
		if strings.HasPrefix(k, prefix) {
			return true
		}
	}
	if _, ok := z.wild[name]; ok {
		return true
	}
	return z.delegations[name]
}

// Count is the number of records in the store.
func (s *Store) Count() int {
	if s == nil {
		return 0
	}
	return s.count
}

// Zones lists the zone apexes.
func (s *Store) Zones() []string {
	if s == nil {
		return nil
	}
	out := make([]string, 0, len(s.zones))
	for z := range s.zones {
		out = append(out, z)
	}
	sort.Strings(out)
	return out
}

// Files lists the source files.
func (s *Store) Files() []string {
	if s == nil {
		return nil
	}
	return s.files
}

// BuiltAt is the compile timestamp.
func (s *Store) BuiltAt() time.Time {
	if s == nil {
		return time.Time{}
	}
	return s.builtAt
}

// Records returns every record in a zone as zone-file text, for the API and
// the UI's zone browser.
func (s *Store) Records(origin string) []string {
	if s == nil {
		return nil
	}
	z, ok := s.zones[strings.ToLower(dns.Fqdn(origin))]
	if !ok {
		return nil
	}
	var out []string
	for _, rrs := range z.exact {
		for _, rr := range rrs {
			out = append(out, rr.String())
		}
	}
	for name, m := range z.wild {
		for _, rrs := range m {
			for _, rr := range rrs {
				out = append(out, strings.Replace(rr.String(), "\t"+rr.Header().Name, "\t*."+name, 1))
			}
		}
	}
	sort.Strings(out)
	return out
}

// Holder keeps the current Store for the hot path.
type Holder struct{ cur atomic.Pointer[Store] }

// NewHolder wraps a store.
func NewHolder(s *Store) *Holder {
	h := &Holder{}
	if s != nil {
		h.cur.Store(s)
	}
	return h
}

// Get returns the current store (nil when nothing is configured).
func (h *Holder) Get() *Store { return h.cur.Load() }

// Set installs a new store.
func (h *Holder) Set(s *Store) { h.cur.Store(s) }