package output

import (
	"encoding/json"
	"io"

	"github.com/hits/ExploreDNS/internal/dns"
	"github.com/hits/ExploreDNS/internal/traverse"
)

type jsonFormatter struct {
	cfg     *Config
	w       io.Writer
	payload jsonDocument
}

type jsonDocument struct {
	Domain    string              `json:"domain"`
	QueryType string              `json:"query_type"`
	Progress  []jsonProgressEvent `json:"progress,omitempty"`
	Resolves  []jsonProgressEvent `json:"resolves,omitempty"`
	Results   []jsonResult        `json:"results,omitempty"`
	Servers   []jsonServer        `json:"servers,omitempty"`
	Summary   jsonSummary         `json:"summary,omitempty"`
}

type jsonProgressEvent struct {
	Stage     string `json:"stage"`
	Depth     int    `json:"depth"`
	Name      string `json:"name"`
	QType     string `json:"qtype"`
	Server    string `json:"server,omitempty"`
	Bailiwick string `json:"bailiwick,omitempty"`
	Resolving bool   `json:"resolving,omitempty"`
}

type jsonResult struct {
	Depth        int      `json:"depth"`
	Probability  float64  `json:"probability"`
	ResponseType string   `json:"response_type"`
	Server       string   `json:"server,omitempty"`
	Answers      []string `json:"answers,omitempty"`
	CNAMEChain   []string `json:"cname_chain,omitempty"`
}

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

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

type jsonAnswerStat struct {
	RData       string   `json:"rdata"`
	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) WriteProgress(event traverse.TraversalEvent) error {
	if !f.cfg.ShowProgress {
		return nil
	}
	f.payload.Progress = append(f.payload.Progress, f.eventToJSON(event))
	return nil
}

func (f *jsonFormatter) WriteResolve(event traverse.TraversalEvent) error {
	if !f.cfg.ShowResolves {
		return nil
	}
	f.payload.Resolves = append(f.payload.Resolves, f.eventToJSON(event))
	return nil
}

func (f *jsonFormatter) WriteResult(result traverse.TraversalResult) error {
	if !f.cfg.ShowAllStats {
		return nil
	}
	f.payload.Results = append(f.payload.Results, f.resultToJSON(result))
	return nil
}

func (f *jsonFormatter) WriteSummary(results []traverse.TraversalResult) error {
	if f.cfg.ShowResults {
		for _, result := range terminalResults(results) {
			f.payload.Results = append(f.payload.Results, f.resultToJSON(result))
		}
	}

	if f.cfg.ShowServers {
		servers := collectServers(results)
		for name, ips := range servers {
			srv := jsonServer{Name: name, IPs: ips}
			if f.cfg.ShowVersions && f.cfg.Fingerprints != nil {
				for _, ip := range ips {
					if v := f.cfg.Fingerprints[ip]; v != "" {
						srv.Version = v
						break
					}
				}
			}
			f.payload.Servers = append(f.payload.Servers, srv)
		}
	}

	if f.cfg.ShowSummaryResults {
		stats := ComputeSummary(results)
		if stats != nil {
			f.payload.Summary = jsonSummary{
				ByType: stats.ByType,
			}
			for _, answer := range stats.Answers {
				f.payload.Summary.Answers = append(f.payload.Summary.Answers, jsonAnswerStat{
					RData:       answer.RData,
					Probability: answer.Prob,
					Records:     answer.RRs,
				})
			}
		}
	}

	return nil
}

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

func (f *jsonFormatter) eventToJSON(event traverse.TraversalEvent) jsonProgressEvent {
	ref := event.Result.Referral
	if ref == nil {
		return jsonProgressEvent{}
	}

	item := jsonProgressEvent{
		Stage:     stageName(event.Stage),
		Depth:     ref.Depth,
		Name:      trimDomain(ref.Name),
		QType:     dns.QNameType(ref.Qtype),
		Bailiwick: trimDomain(ref.Bailiwick),
		Resolving: !ref.HasAddresses(),
	}
	if event.Result.Response != nil && event.Result.Response.Server != nil {
		item.Server = event.Result.Response.Server.String()
	} else {
		item.Server = referralServerLabel(ref, event.Result.Response)
	}
	return item
}

func (f *jsonFormatter) resultToJSON(result traverse.TraversalResult) jsonResult {
	item := jsonResult{}
	if result.Referral != nil {
		item.Depth = result.Referral.Depth
		item.Probability = result.Referral.Prob
	}
	if result.Response != nil {
		item.ResponseType = result.Response.Type.String()
		if result.Response.Server != nil {
			item.Server = result.Response.Server.String()
		}
		if result.Response.Decoded != nil {
			for _, rr := range result.Response.Decoded.Answers {
				item.Answers = append(item.Answers, dns.FormatRecord(rr))
			}
			item.CNAMEChain = append(item.CNAMEChain, result.Response.Decoded.CNAMEChain...)
		}
	}
	return item
}

func stageName(stage traverse.EventStage) string {
	switch stage {
	case traverse.EventStart:
		return "start"
	case traverse.EventComplete:
		return "complete"
	default:
		return "unknown"
	}
}