package esstream

import (
	"bytes"
	"context"
	"encoding/csv"
	"fmt"
	"github.com/elastic/go-elasticsearch/v6"
	"github.com/pingcap/log"
	"github.com/pingcap/tep/pkg/utils"
	"github.com/pingcap/tep/pkg/utils/es/query"
	"go.uber.org/zap"
	"io"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"strings"
	"time"
)

type Stream struct {
	client         *elasticsearch.Client
	indexPrefix    string
	buf            *bytes.Buffer
	csvWriter      *csv.Writer
	scrollID       string
	lastResultSize int
	csvLog         ESToCSV
}

//New return a stream and open it
func New(client *elasticsearch.Client, indexPrefix string, queryStr string, csvLog ESToCSV) (*Stream, error) {
	buf := &bytes.Buffer{}
	csvWriter := csv.NewWriter(buf)
	err := csvWriter.Write(csvLog.Header())
	if err != nil {
		return nil, err
	}
	res, err := client.Search(
		client.Search.WithContext(context.Background()),
		client.Search.WithIndex(fmt.Sprintf("%s*", indexPrefix)),
		client.Search.WithBody(strings.NewReader(queryStr)),
		client.Search.WithPretty(),
		client.Search.WithScroll(time.Duration(time.Minute)),
	)
	if err != nil {
		log.Error("Query from elasticsearch failed", zap.String("query", queryStr), zap.Error(err))
		return nil, err
	}

	searchResponse := query.NewResponse(res.Body)
	defer utils.CallDeferMethod(res.Body.Close, "close body failed")
	err = csvWriter.WriteAll(csvLog.ToCSV(searchResponse))
	if err != nil {
		return nil, err
	}
	return &Stream{
		client:         client,
		indexPrefix:    indexPrefix,
		buf:            buf,
		csvWriter:      csvWriter,
		scrollID:       searchResponse.ScrollID,
		lastResultSize: len(searchResponse.Result()),
		csvLog:         csvLog,
	}, nil
}

func (stream *Stream) GetObjectKind() schema.ObjectKind {
	return schema.EmptyObjectKind
}
func (stream *Stream) DeepCopyObject() runtime.Object {
	panic("rest.LocationStreamer does not implement DeepCopyObject")
}

func (stream *Stream) InputStream(ctx context.Context, apiVersion, acceptHeader string) (io.ReadCloser, bool, string, error) {
	return stream, true, "text/csv", nil
}

func (stream *Stream) Close() error {
	return nil
}

func (stream *Stream) Read(p []byte) (int, error) {
	n, err := stream.buf.Read(p)
	if err == io.EOF {
		//we read all buf, try to get another one if there is one
		if stream.scrollID != "" && stream.lastResultSize > 0 {
			res, scanErr := stream.client.Scroll(
				stream.client.Scroll.WithContext(context.Background()),
				stream.client.Scroll.WithPretty(),
				stream.client.Scroll.WithScroll(time.Duration(time.Minute)),
				stream.client.Scroll.WithScrollID(stream.scrollID),
			)
			if scanErr != nil {
				return 0, scanErr
			}
			searchResponse := query.NewResponse(res.Body)
			defer utils.CallDeferMethod(res.Body.Close, "close body failed")
			err = stream.csvWriter.WriteAll(stream.csvLog.ToCSV(searchResponse))
			if err != nil {
				stream.lastResultSize = 0
				stream.scrollID = ""
				return 0, err
			}
			stream.scrollID = searchResponse.ScrollID
			stream.lastResultSize = len(searchResponse.Result())
			//read again, if EOF return it directly
			return stream.buf.Read(p)
		}
	}
	return n, err
}