package utils

import (
	"bytes"
	"context"
	"encoding/csv"
	"github.com/pingcap/log"
	"go.uber.org/zap"
	"io"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

type Iterator interface {
	Header() []string
	Open() error
	Next() []string
	Close() error
}

type Stream struct {
	buf         *bytes.Buffer
	csvWriter   *csv.Writer
	it          Iterator
	bufItemSize int
}

//NewCSVStream return a k8s csv stream and open it
func NewCSVStream(it Iterator, bufSize int) (*Stream, error) {
	buf := &bytes.Buffer{}
	csvWriter := csv.NewWriter(buf)
	err := csvWriter.Write(it.Header())
	csvWriter.Flush()
	if err != nil {
		return nil, errors.NewInternalError(err)
	}
	err = it.Open()
	if err != nil {
		log.Error("open iterator failed", zap.Error(err))
		return nil, errors.NewInternalError(err)
	}
	return &Stream{
		buf:         buf,
		csvWriter:   csvWriter,
		it:          it,
		bufItemSize: bufSize,
	}, 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 stream.it.Close()
}

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
		size := 0
		for c := stream.it.Next(); c != nil; {
			err = stream.csvWriter.Write(c)
			if size++; size >= stream.bufItemSize {
				break
			}
			c = stream.it.Next()
		}
		if size > 0 {
			stream.csvWriter.Flush()
			//read again, if EOF return it directly
			return stream.buf.Read(p)
		}
	}
	return n, err
}