package traverse

import (
	"sort"
	"strings"

	miekgdns "github.com/miekg/dns"
)

// AnswerStat is one distinct answered RRset with its accumulated probability.
type AnswerStat struct {
	// Key groups identical RRset content: sorted rdata strings joined with
	// "@@@" (summary_stats.rb get_answer_stats).
	Key  string
	Prob float64
	RRs  []miekgdns.RR
}

// SummaryStats groups the aggregated leaves by status; answered leaves are
// additionally grouped by RRset content, one entry per distinct RRset
// (summary_stats.rb). The ByStatus probabilities sum to 1.0 and the answer
// probabilities sum to ByStatus[StatusAnswered].
type SummaryStats struct {
	ByStatus map[Status]float64
	Answers  []AnswerStat
}

// SummaryStats computes (and memoises) the summary grouping of this node's
// aggregated leaf statistics (referral.rb summary_stats). It returns nil
// until the node's statistics have been calculated.
func (r *Referral) SummaryStats() *SummaryStats {
	if r == nil || !r.calculated || len(r.Stats) == 0 {
		return nil
	}
	if r.summaryStats != nil {
		return r.summaryStats
	}

	stats := &SummaryStats{ByStatus: make(map[Status]float64)}
	answers := make(map[string]*AnswerStat)
	for _, leaf := range r.StatsList() {
		status := leaf.Response.Status
		stats.ByStatus[status] += leaf.Prob
		if status != StatusAnswered {
			continue
		}
		rdatas := make([]string, 0, len(leaf.Response.DQ.Answers))
		for _, rr := range leaf.Response.DQ.Answers {
			rdatas = append(rdatas, rrData(rr))
		}
		sort.Strings(rdatas)
		key := strings.Join(rdatas, "@@@")
		if e, ok := answers[key]; ok {
			e.Prob += leaf.Prob
		} else {
			answers[key] = &AnswerStat{Key: key, Prob: leaf.Prob, RRs: leaf.Response.DQ.Answers}
		}
	}

	for _, e := range answers {
		stats.Answers = append(stats.Answers, *e)
	}
	sort.Slice(stats.Answers, func(i, j int) bool { return stats.Answers[i].Key < stats.Answers[j].Key })
	r.summaryStats = stats
	return stats
}

// rrData extracts the rdata portion of a record (dnsruby rdata_to_string):
// everything after owner/TTL/class/type in presentation format.
func rrData(rr miekgdns.RR) string {
	s := rr.String()
	h := rr.Header().String()
	return strings.TrimPrefix(s, h)
}