// Package query ... the response from elasticsearch.
// The raw response of elasticsearch is json format, it is ugly and complex.
package query

import (
	"encoding/json"
	"github.com/wushilin/stream"
	"io"

	"github.com/pingcap/log"
	"go.uber.org/zap"
)

type QueryResponse struct {
	ScrollID     string           `json:"_scroll_id,omitempty"`
	Hits         *hits            `json:"hits,omitempty"`
	Error        *errorMsg        `json:"error,omitempty"`
	Aggregations *json.RawMessage `json:"aggregations,omitempty"`
}

type hits struct {
	Total int64  `json:"total,omitempty"`
	Hits  []*Hit `json:"hits,omitempty"`
}

type buckets struct {
	Buckets []*bucket `json:"buckets"`
}

type bucket struct {
	Key   string `json:"key,omitempty"`
	Count int    `json:"doc_count,omitempty"`
}

type Hit struct {
	ID     string           `json:"_id,omitempty"`
	Source *json.RawMessage `json:"_source,omitempty"`
}

type errorMsg struct {
	Reason   string                 `json:"reason,omitempty"`
	CausedBy map[string]interface{} `json:"caused_by,omitempty"`
}

//TotalHits return the number of hits for a search result.
func (r *QueryResponse) TotalHits() int64 {
	if r.Hits != nil {
		return r.Hits.Total
	}
	return 0
}

func (r *QueryResponse) printErrMsg() {
	if r.Error.Reason == "" {
		return
	}

	log.Error("failed to query from elasticsearch", zap.String("reason", r.Error.Reason), zap.Any("caused_by", r.Error.CausedBy))
}

// Result iterate over all hits. It skips empty result and return json list
func (r *QueryResponse) Result() []string {
	if r.Hits == nil || r.Hits.Hits == nil || len(r.Hits.Hits) == 0 {
		return make([]string, 0)
	}

	var results []string
	for _, hit := range r.Hits.Hits {
		if hit.Source == nil {
			continue
		}

		results = append(results, string(*hit.Source))
	}
	return results
}

// Result iterate over all hits. It skips empty result and return json list
func (r *QueryResponse) HitsSlice() []*Hit {
	if r.Hits == nil || r.Hits.Hits == nil || len(r.Hits.Hits) == 0 {
		return make([]*Hit, 0)
	}
	return r.Hits.Hits
}

// Agg iterate over all aggreations.
func (r *QueryResponse) Agg(from int) (map[string]int, error) {
	if r.Aggregations == nil {
		return make(map[string]int), nil
	}

	var agg map[string]*buckets
	d, err := r.Aggregations.MarshalJSON()
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(d, &agg)
	if err != nil {
		return nil, err
	}

	result := make(map[string]int)
	stream.FromMapValues(agg).Each(func(b *buckets) {
		if b.Buckets == nil || len(b.Buckets) == 0 {
			return
		}
		// skip offset
		stream.FromArray(b.Buckets).Skip(from).Each(func(b *bucket) {
			result[b.Key] = b.Count
		})

	})

	return result, nil
}

// NewResponse return response with structure elasticsearch's response
func NewResponse(data io.ReadCloser) *QueryResponse {
	r := new(QueryResponse)
	json.NewDecoder(data).Decode(r)
	return r
}