package traverse

import (
	"context"
	"fmt"
	"net"
	"strings"

	"github.com/hits/ExploreDNS/internal/dns"
	miekgdns "github.com/miekg/dns"
	"golang.org/x/net/idna"
)

// ResolutionState tracks whether a Referral's nameserver addresses have been resolved.
type ResolutionState int

// Resolution state constants.
const (
	StateUnresolved ResolutionState = iota // nameserver addresses are not yet known
	StateResolving                         // address resolution is in progress
	StateResolved                          // addresses are available in Addresses
)

func (s ResolutionState) String() string {
	switch s {
	case StateUnresolved:
		return "unresolved"
	case StateResolving:
		return "resolving"
	case StateResolved:
		return "resolved"
	default:
		return "unknown"
	}
}

// Referral represents a pending DNS query: a (name, qtype) pair delegated to a
// set of nameserver addresses.  Referrals form a linked-list chain through
// Parent, enabling loop and depth detection.
type Referral struct {
	// Name is the fully-qualified domain name being queried.
	Name string
	// Qtype is the DNS record type being queried.
	Qtype uint16
	// Qclass is the DNS class (always ClassINET in practice).
	Qclass uint16
	// Bailiwick is the zone that delegated this referral.
	Bailiwick string

	// Addresses holds the resolved IP addresses for this nameserver referral.
	Addresses []net.IP
	// State tracks whether Addresses have been resolved.
	State ResolutionState

	// NSName is the nameserver hostname (before IP resolution).
	NSName string
	// Parent is the Referral that triggered this one, or nil for the root.
	Parent *Referral
	// Depth is the number of referral hops from the root.
	Depth int
	// Prob is the probability weight for this branch (product of 1/fanout at each step).
	Prob float64
}

// idnaLookup is the IDN lookup profile used to convert internationalised domain
// names (unicode labels) to their ACE/punycode equivalents before querying.
var idnaLookup = idna.New(
	idna.MapForLookup(),
	idna.BidiRule(),
	idna.StrictDomainName(false),
)

// toASCII converts a domain name that may contain unicode labels to its
// punycode (ACE) representation. Pure-ASCII names are returned unchanged.
// On conversion errors the original name is returned so the caller can still
// attempt a query (the server will reject it if truly invalid).
func toASCII(name string) string {
	if name == "" || name == "." {
		return name
	}
	ascii, err := idnaLookup.ToASCII(name)
	if err != nil {
		return name
	}
	return ascii
}

// NewReferral creates a Referral for (name, qtype) within bailiwick, at the
// given depth and probability.  The name and bailiwick are normalised to
// lowercase FQDN, and internationalised labels are converted to punycode.
func NewReferral(name string, qtype uint16, bailiwick string, depth int, prob float64, parent *Referral) *Referral {
	return &Referral{
		Name:      miekgdns.Fqdn(strings.ToLower(toASCII(name))),
		Qtype:     qtype,
		Qclass:    miekgdns.ClassINET,
		Bailiwick: miekgdns.Fqdn(strings.ToLower(toASCII(bailiwick))),
		Depth:     depth,
		Prob:      prob,
		Parent:    parent,
		State:     StateUnresolved,
	}
}

// InBailiwick reports whether name is within this referral's bailiwick zone.
// A root bailiwick ("." or "") is treated as matching everything.
func (r *Referral) InBailiwick(name string) bool {
	if r.Bailiwick == "" || r.Bailiwick == "." {
		return true
	}
	fqdn := miekgdns.Fqdn(strings.ToLower(name))
	return miekgdns.IsSubDomain(r.Bailiwick, fqdn)
}

// HasAddresses reports whether at least one nameserver IP address is known.
func (r *Referral) HasAddresses() bool {
	return len(r.Addresses) > 0
}

// SetAddresses stores addrs and updates State accordingly.
func (r *Referral) SetAddresses(addrs []net.IP) {
	r.Addresses = addrs
	if len(addrs) > 0 {
		r.State = StateResolved
	} else {
		r.State = StateUnresolved
	}
}

// CircularReferralError is returned when a referral chain revisits a nameserver,
// indicating a circular delegation.
type CircularReferralError struct {
	Name  string
	Chain []string
}

func (e *CircularReferralError) Error() string {
	return fmt.Sprintf("circular referral detected for %s: %v", e.Name, e.Chain)
}

// UnresolvableNameserverError is returned when a nameserver hostname cannot be
// resolved to any IP address.
type UnresolvableNameserverError struct {
	Name   string
	Reason string
}

func (e *UnresolvableNameserverError) Error() string {
	return fmt.Sprintf("unresolvable nameserver %s: %s", e.Name, e.Reason)
}

// Resolve attempts to resolve the nameserver addresses for this Referral by
// performing a fresh iterative traversal.  It stores the found addresses and
// updates State.  Returns an error when resolution fails.
func (r *Referral) Resolve(ctx context.Context, traverser *Traverser, cache *InfoCache, visited map[string]bool, depth int) error {
	if r.HasAddresses() {
		r.State = StateResolved
		return nil
	}

	if cache != nil {
		if addrs := cache.LookupGlue(r.Name); len(addrs) > 0 {
			r.Addresses = addrs
			r.State = StateResolved
			return nil
		}
	}

	if visited != nil {
		if visited[r.Name] {
			return &CircularReferralError{
				Name:  r.Name,
				Chain: getVisitedNames(visited),
			}
		}
		visited[r.Name] = true
	}

	if depth > DefaultMaxDepth {
		return &UnresolvableNameserverError{
			Name:   r.Name,
			Reason: "max depth exceeded",
		}
	}

	roots, err := traverser.discoverRoots(ctx)
	if err != nil {
		return fmt.Errorf("root discovery: %w", err)
	}

	initial := NewReferral(r.Name, dns.TypeA, ".", 0, 1.0, nil)
	initial.Addresses = roots
	initial.State = StateResolved

	stack := NewStack(DefaultMaxDepth)
	stack.Push(initial)

	traversalCache := NewInfoCache(nil)
	if visited != nil {
		for name := range visited {
			traversalCache.StoreGlue(name, []net.IP{})
		}
	}

	var lastErr error
	for {
		select {
		case <-ctx.Done():
			return fmt.Errorf("resolution cancelled: %w", ctx.Err())
		default:
		}

		ref := stack.Pop()
		if ref == nil {
			break
		}

		cacheForStep := traversalCache
		if ref.Parent != nil {
			cacheForStep = traversalCache.Child()
		}

		resp := traverser.processReferral(ctx, ref, cacheForStep)

		if resp.Type == RespAnswer {
			if len(resp.Decoded.Answers) > 0 {
				var addrs []net.IP
				for _, rr := range resp.Decoded.Answers {
					if a, ok := rr.(*miekgdns.A); ok {
						addrs = append(addrs, a.A)
					}
					if aaaa, ok := rr.(*miekgdns.AAAA); ok {
						addrs = append(addrs, aaaa.AAAA)
					}
				}
				if len(addrs) > 0 {
					r.Addresses = addrs
					r.State = StateResolved
					if cache != nil {
						cache.StoreGlue(r.Name, addrs)
					}
					return nil
				}
			}
		}

		if resp.Type == RespNXDOMAIN {
			lastErr = &UnresolvableNameserverError{
				Name:   r.Name,
				Reason: "NXDOMAIN",
			}
			break
		}

		if resp.Type == RespSERVFAIL || resp.Type == RespError {
			lastErr = fmt.Errorf("server error resolving %s: %s", r.Name, resp.Type)
			continue
		}

		if resp.Type == RespReferral {
			children := resp.ChildReferrals()
			for _, child := range children {
				if visited != nil && visited[child.Name] {
					continue
				}
				if !stack.Push(child) {
					lastErr = &UnresolvableNameserverError{
						Name:   r.Name,
						Reason: "max depth exceeded during resolution",
					}
				}
			}
		}
	}

	if lastErr != nil {
		return lastErr
	}

	return &UnresolvableNameserverError{
		Name:   r.Name,
		Reason: "resolution exhausted without answer",
	}
}

func getVisitedNames(visited map[string]bool) []string {
	var names []string
	for name := range visited {
		names = append(names, name)
	}
	return names
}

// IsNameInChain reports whether name appears anywhere in this referral's ancestor
// chain, including this referral itself. Used for CNAME loop detection.
func (r *Referral) IsNameInChain(name string) bool {
	n := miekgdns.Fqdn(strings.ToLower(name))
	curr := r
	for curr != nil {
		if curr.Name == n {
			return true
		}
		curr = curr.Parent
	}
	return false
}