package stream

import (
	"fmt"
	"testing"

	"github.com/onsi/gomega"
	"github.com/pingcap/errors"
)

func Test_OfMapEach(t *testing.T) {
	m := make(map[int]int)
	m[1] = 1
	m[2] = 2
	m[3] = 3

	expectM := make(map[int]int)
	OfMap(m).Each(func(entry MapEntry) {
		expectM[int(entry.Key.Int())] = int(entry.Value.Int())
	})

	// stream.ofmap().map().each
	g := gomega.NewGomegaWithT(t)
	g.Expect(len(expectM)).To(gomega.Equal(len(m)))
	g.Expect(expectM[1]).To(gomega.Equal(1))
	g.Expect(expectM[2]).To(gomega.Equal(2))
	g.Expect(expectM[3]).To(gomega.Equal(3))
}

func Test_MapMap(t *testing.T) {
	m := make(map[int]int)
	m[1] = 1
	m[2] = 2
	m[3] = 3

	expectM := make([]int, 0)
	OfMap(m).Map(func(entry MapEntry) int {
		return int(entry.Key.Int() + entry.Value.Int())
	}).Each(func(v int) {
		fmt.Printf("===Test_MapMap=%d \n", v)
		expectM = append(expectM, v)
	})

	g := gomega.NewGomegaWithT(t)
	g.Expect(len(expectM)).To(gomega.Equal(len(m)))
	g.Expect(expectM).To(gomega.ContainElement(2))
	g.Expect(expectM).To(gomega.ContainElement(4))
	g.Expect(expectM).To(gomega.ContainElement(6))
}

func Test_MapCollect(t *testing.T) {
	m := make(map[int]int)
	m[1] = 1
	m[2] = 2
	m[3] = 3

	expectM := make([]int, 0)
	OfMap(m).Map(func(entry MapEntry) int {
		return int(entry.Key.Int()+entry.Value.Int()) * 2
	}).Collect(&expectM)
	g := gomega.NewGomegaWithT(t)
	g.Expect(len(expectM)).To(gomega.Equal(len(m)))
	g.Expect(expectM).To(gomega.ContainElement(4))
	g.Expect(expectM).To(gomega.ContainElement(8))
	g.Expect(expectM).To(gomega.ContainElement(12))
}

func Test_Consume(t *testing.T) {
	dataLen := 100
	data := make([]string, dataLen)
	for i := 0; i < dataLen; i++ {
		data[i] = fmt.Sprintf("data%d", i)
	}

	ite := &TestIterator{
		len:   len(data),
		index: 0,
		data:  data,
	}

	g := gomega.NewWithT(t)

	actualResults := make([]string, 0)
	f := func(next string) {
		g.Expect(data).To(gomega.ContainElements(next))
		actualResults = append(actualResults, next)
	}

	err := OfIterator(ite).Consume(f)
	g.Expect(err).To(gomega.Succeed())

	g.Expect(len(actualResults)).To(gomega.Equal(len(data)))

	for _, d := range actualResults {
		g.Expect(data).To(gomega.ContainElements(d))
	}
}

func Test_SinkBatchUnEvenNumber(t *testing.T) {
	dataLen := 999
	data := make([]string, dataLen)
	for i := 0; i < dataLen; i++ {
		data[i] = fmt.Sprintf("data%d", i)
	}

	ite := &TestIterator{
		len:   len(data),
		index: 0,
		data:  data,
	}

	g := gomega.NewWithT(t)

	actualResults := make([]string, 0)
	f := func(arr []string) error {
		for _, d := range arr {
			g.Expect(data).To(gomega.ContainElements(d))
		}
		actualResults = append(actualResults, arr...)

		if len(arr) < 9 {
			return errors.New("unexpect")
		}

		return nil
	}

	err := OfIterator(ite).BatchSink(f, 10)
	g.Expect(err).To(gomega.Succeed())

	g.Expect(len(actualResults)).To(gomega.Equal(len(data)))

	for _, d := range actualResults {
		g.Expect(data).To(gomega.ContainElements(d))
	}
}

func Test_SinkBatchEvenNumber(t *testing.T) {
	dataLen := 1000
	data := make([]string, dataLen)
	for i := 0; i < dataLen; i++ {
		data[i] = fmt.Sprintf("data%d", i)
	}

	ite := &TestIterator{
		len:   len(data),
		index: 0,
		data:  data,
	}

	g := gomega.NewWithT(t)

	actualResults := make([]string, 0)
	f := func(arr []string) error {
		for _, d := range arr {
			g.Expect(data).To(gomega.ContainElements(d))
		}
		actualResults = append(actualResults, arr...)

		if len(arr) != 10 {
			return errors.New("unexpect")
		}

		return nil
	}

	err := OfIterator(ite).BatchSink(f, 10)
	g.Expect(err).To(gomega.Succeed())

	g.Expect(len(actualResults)).To(gomega.Equal(len(data)))

	for _, d := range actualResults {
		g.Expect(data).To(gomega.ContainElements(d))
	}
}

func Test_ConsumeErr(t *testing.T) {
	dataLen := 100
	data := make([]string, dataLen)
	for i := 0; i < dataLen; i++ {
		data[i] = fmt.Sprintf("data%d", i)
	}

	ite := &TestIterator{
		len:   len(data),
		index: 0,
		data:  data,
	}

	g := gomega.NewWithT(t)

	actualResults := make([]string, 0)
	f := func(next string) error {
		g.Expect(data).To(gomega.ContainElements(next))
		actualResults = append(actualResults, next)

		if next == "data0" || next == "data1" {
			return errors.New(next)
		}

		return nil
	}

	err := OfIterator(ite).Consume(f)

	g.Expect(err).To(gomega.MatchError("data0; data1"))

	g.Expect(len(actualResults)).To(gomega.Equal(len(data)))

	for _, d := range actualResults {
		g.Expect(data).To(gomega.ContainElements(d))
	}
}

func Test_Sink(t *testing.T) {
	dataLen := 20
	data := make([]string, dataLen)
	for i := 0; i < dataLen; i++ {
		data[i] = fmt.Sprintf("data%d", i)
	}

	ite := &TestIterator{
		len:   len(data),
		index: 0,
		data:  data,
	}

	g := gomega.NewGomegaWithT(t)
	f := func(result []string) error {
		g.Expect(len(result)).To(gomega.Equal(dataLen))
		for _, d := range data {
			g.Expect(result).To(gomega.ContainElements(d))
		}

		return nil
	}
	err := OfIterator(ite).Sink(f)
	g.Expect(err).To(gomega.Succeed())
}

func Test_SinkError(t *testing.T) {
	dataLen := 100
	data := make([]string, dataLen)
	for i := 0; i < dataLen; i++ {
		data[i] = fmt.Sprintf("data%d", i)
	}

	ite := &TestIterator{
		len:   len(data),
		index: 0,
		data:  data,
	}

	g := gomega.NewGomegaWithT(t)
	f := func(result []string) error {
		g.Expect(len(result)).To(gomega.Equal(dataLen))
		for _, d := range data {
			g.Expect(result).To(gomega.ContainElements(d))
		}

		return errors.New("Test error")
	}

	err := OfIterator(ite).Sink(f)
	g.Expect(err).To(gomega.MatchError("Test error"))
}

type TestIterator struct {
	len   int
	index int
	data  []string
}

func (ti *TestIterator) Next() (any, bool) {
	if ti.index > ti.len-1 {
		return nil, false
	}

	d := ti.data[ti.index]
	ti.index++

	return d, true
}

func (ti *TestIterator) Len() (int, bool) {
	return ti.len, true
}