package traverse

import (
	"net"

	"github.com/hits/ExploreDNS/internal/dns"
	miekgdns "github.com/miekg/dns"
)

type ResponseType int

const (
	RespReferral  ResponseType = iota
	RespAnswer
	RespCNAMEFollow
	RespNODATA
	RespNXDOMAIN
	RespSERVFAIL
	RespError
)

func (rt ResponseType) String() string {
	switch rt {
	case RespReferral:
		return "referral"
	case RespAnswer:
		return "answer"
	case RespCNAMEFollow:
		return "cname_follow"
	case RespNODATA:
		return "nodata"
	case RespNXDOMAIN:
		return "nxdomain"
	case RespSERVFAIL:
		return "servfail"
	case RespError:
		return "error"
	default:
		return "unknown"
	}
}

type Response struct {
	Referral *Referral
	Server   net.IP
	Cache    *InfoCache
	Decoded  *dns.DecodedResponse
	Type     ResponseType
}

func NewResponse(ref *Referral, server net.IP, cache *InfoCache) *Response {
	return &Response{
		Referral: ref,
		Server:   server,
		Cache:    cache,
	}
}

func (r *Response) Process(msg *miekgdns.Msg) *Response {
	if msg == nil {
		r.Type = RespError
		return r
	}

	r.Decoded = dns.DecodeResponse(msg)
	if r.Decoded == nil {
		r.Type = RespError
		return r
	}

	r.Type = r.classify()
	return r
}

func (r *Response) classify() ResponseType {
	switch r.Decoded.Classification {
	case dns.ResponseNXDOMAIN:
		return RespNXDOMAIN
	case dns.ResponseSERVFAIL:
		return RespSERVFAIL
	case dns.ResponseAnswer:
		if len(r.Decoded.CNAMEChain) > 0 && !r.hasFinalAnswer() {
			return RespCNAMEFollow
		}
		return RespAnswer
	case dns.ResponseReferral:
		return RespReferral
	case dns.ResponseNODATA:
		return RespNODATA
	default:
		return RespError
	}
}

func (r *Response) hasFinalAnswer() bool {
	for _, rr := range r.Decoded.Answers {
		if _, ok := rr.(*miekgdns.CNAME); ok {
			continue
		}
		return true
	}
	return false
}

func (r *Response) ChildReferrals() []*Referral {
	if r.Type != RespReferral {
		return nil
	}
	if r.Referral == nil {
		return nil
	}

	var nameservers []string
	for _, rr := range r.Decoded.Authority {
		if ns, ok := rr.(*miekgdns.NS); ok {
			if r.Referral.InBailiwick(ns.Ns) {
				nameservers = append(nameservers, ns.Ns)
			}
		}
	}

	if len(nameservers) == 0 {
		for _, rr := range r.Decoded.Authority {
			if ns, ok := rr.(*miekgdns.NS); ok {
				nameservers = append(nameservers, ns.Ns)
			}
		}
	}

	r.storeAuthority(nameservers)

	prob := r.childProb(len(nameservers))
	var children []*Referral
	for _, ns := range nameservers {
		child := NewReferral(
			r.Referral.Name,
			r.Referral.Qtype,
			ns,
			r.Referral.Depth+1,
			prob,
			r.Referral,
		)
		r.resolveGlue(child)
		children = append(children, child)
	}
	return children
}

func (r *Response) CNAMEFollowReferral() *Referral {
	if r.Type != RespCNAMEFollow || len(r.Decoded.CNAMEChain) == 0 {
		return nil
	}
	target := r.Decoded.CNAMEChain[len(r.Decoded.CNAMEChain)-1]
	follow := NewReferral(
		target,
		r.Referral.Qtype,
		r.Referral.Bailiwick,
		r.Referral.Depth+1,
		r.Referral.Prob,
		r.Referral,
	)
	if len(r.Referral.Addresses) > 0 {
		follow.Addresses = make([]net.IP, len(r.Referral.Addresses))
		copy(follow.Addresses, r.Referral.Addresses)
		follow.State = StateResolved
	}
	return follow
}

func (r *Response) storeAuthority(nameservers []string) {
	if r.Cache == nil {
		return
	}
	zone := r.Referral.Name
	r.Cache.StoreNS(zone, nameservers)
}

func (r *Response) resolveGlue(child *Referral) {
	if r.Cache == nil {
		return
	}
	nsName := child.Bailiwick
	for _, rr := range r.Decoded.Additional {
		switch v := rr.(type) {
		case *miekgdns.A:
			if normalize(v.Header().Name) == normalize(nsName) {
				child.Addresses = append(child.Addresses, v.A)
			}
		case *miekgdns.AAAA:
			if normalize(v.Header().Name) == normalize(nsName) {
				child.Addresses = append(child.Addresses, v.AAAA)
			}
		}
	}
	if child.HasAddresses() {
		child.State = StateResolved
	}
	r.Cache.StoreGlue(nsName, child.Addresses)
}

func (r *Response) IsTerminal() bool {
	switch r.Type {
	case RespAnswer, RespNODATA, RespNXDOMAIN, RespSERVFAIL, RespError:
		return true
	default:
		return false
	}
}

func (r *Response) childProb(n int) float64 {
	if n <= 0 {
		return 0
	}
	if r.Referral == nil {
		return 1.0 / float64(n)
	}
	return r.Referral.Prob / float64(n)
}