package output

import (
	"encoding/json"
	"io"
	"sort"
	"strings"

	"gitea.hansenits.com.au/hits/ExploreDNS/internal/traverse"
)

// jsonFormatter accumulates the run into a single document emitted once by
// Flush: {domain, qtype, root, results, summary, servers}. Aggregated leaves
// appear exactly once (in results); progress events are not recorded.
type jsonFormatter struct {
	cfg     *Config
	w       io.Writer
	payload jsonDocument
}

type jsonDocument struct {
	Domain    string       `json:"domain"`
	QueryType string       `json:"qtype"`
	Root      *jsonRoot    `json:"root,omitempty"`
	Results   []jsonResult `json:"results,omitempty"`
	Summary   *jsonSummary `json:"summary,omitempty"`
	Servers   []jsonServer `json:"servers,omitempty"`
}

// jsonRoot is the initial root the traversal started from.
type jsonRoot struct {
	Name string `json:"name"`
	IP   string `json:"ip,omitempty"`
}

// jsonResult is one aggregated leaf outcome, emitted exactly once.
type jsonResult struct {
	RefID       string   `json:"refid,omitempty"`
	Probability float64  `json:"probability"`
	Status      string   `json:"status"`
	Server      string   `json:"server,omitempty"`
	IP          string   `json:"ip,omitempty"`
	Qname       string   `json:"qname,omitempty"`
	Qtype       string   `json:"qtype,omitempty"`
	Answers     []string `json:"answers,omitempty"`
	Message     string   `json:"message,omitempty"`
}

type jsonServer struct {
	Name    string   `json:"name"`
	IPs     []string `json:"ips"`
	Version string   `json:"version,omitempty"`
}

type jsonSummary struct {
	ByStatus map[string]float64 `json:"by_status,omitempty"`
	Answers  []jsonAnswerStat   `json:"answers,omitempty"`
}

type jsonAnswerStat struct {
	Probability float64  `json:"probability"`
	Records     []string `json:"records,omitempty"`
}

func newJSONFormatter(cfg *Config, w io.Writer) *jsonFormatter {
	return &jsonFormatter{
		cfg: cfg,
		w:   w,
		payload: jsonDocument{
			Domain:    cfg.Domain,
			QueryType: cfg.QueryType,
		},
	}
}

func (f *jsonFormatter) WriteHeader(roots []traverse.StartServer) error {
	if len(roots) == 0 {
		return nil
	}
	root := &jsonRoot{Name: roots[0].Name}
	if len(roots[0].IPs) > 0 {
		root.IP = roots[0].IPs[0]
	}
	f.payload.Root = root
	return nil
}

// WriteProgress is a no-op: the JSON document contains only the aggregated
// outcome, never per-event duplicates.
func (f *jsonFormatter) WriteProgress(traverse.TraversalEvent) error {
	return nil
}

func (f *jsonFormatter) WriteResolve(traverse.TraversalEvent) error {
	return nil
}

func (f *jsonFormatter) WriteSummary(root *traverse.Referral, servers map[string][]string) error {
	if root != nil && f.cfg.ShowResults {
		// StatsList is sorted by stats key: deterministic ordering.
		for _, leaf := range root.StatsList() {
			f.payload.Results = append(f.payload.Results, leafToJSON(leaf))
		}
	}

	if f.cfg.ShowServers {
		names := make([]string, 0, len(servers))
		for name := range servers {
			names = append(names, name)
		}
		sort.Slice(names, func(i, j int) bool {
			return reverseString(strings.ToLower(names[i])) < reverseString(strings.ToLower(names[j]))
		})
		for _, name := range names {
			srv := jsonServer{Name: name, IPs: servers[name]}
			if f.cfg.ShowVersions && f.cfg.Fingerprints != nil {
				for _, ip := range srv.IPs {
					if v := f.cfg.Fingerprints[ip]; v != "" {
						srv.Version = v
						break
					}
				}
			}
			f.payload.Servers = append(f.payload.Servers, srv)
		}
	}

	if f.cfg.ShowSummaryResults {
		if stats := root.SummaryStats(); stats != nil {
			summary := &jsonSummary{ByStatus: make(map[string]float64)}
			for status, prob := range stats.ByStatus {
				summary.ByStatus[string(status)] = prob
			}
			for _, answer := range stats.Answers {
				stat := jsonAnswerStat{Probability: answer.Prob}
				for _, rr := range answer.RRs {
					stat.Records = append(stat.Records, collapseWhitespace(rr.String()))
				}
				summary.Answers = append(summary.Answers, stat)
			}
			f.payload.Summary = summary
		}
	}

	return nil
}

func (f *jsonFormatter) Flush() error {
	enc := json.NewEncoder(f.w)
	enc.SetIndent("", "  ")
	return enc.Encode(f.payload)
}

func leafToJSON(leaf *traverse.StatsEntry) jsonResult {
	resp := leaf.Response
	item := jsonResult{
		Probability: leaf.Prob,
		Status:      string(resp.Status),
		IP:          resp.IP,
		Qname:       resp.Qname,
		Qtype:       traverse.TypeToString(resp.Qtype),
	}
	if leaf.Referral != nil {
		item.RefID = leaf.Referral.RefID
		item.Server = leaf.Referral.Server
	}
	if resp.DQ != nil {
		for _, rr := range resp.DQ.Answers {
			item.Answers = append(item.Answers, collapseWhitespace(rr.String()))
		}
		switch resp.Status {
		case traverse.StatusError:
			item.Message = resp.DQ.ErrorMessage
		case traverse.StatusException:
			item.Message = resp.DQ.ExceptionMessage
		}
	}
	return item
}