/*
 * Copyright 2024 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package container

import (
	"context"
	"errors"
	"io"
	"runtime"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/cloudwego/kitex/internal/test"
)

func TestPipeline(t *testing.T) {
	ctx := context.Background()
	pipe := NewPipe[int]()
	var recv int
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		items := make([]int, 10)
		for {
			n, err := pipe.Read(ctx, items)
			if err != nil {
				if err != io.EOF {
					t.Error(err)
				}
				return
			}
			for i := 0; i < n; i++ {
				recv += items[i]
			}
		}
	}()
	round := 10000
	itemsPerRound := []int{1, 1, 1, 1, 1}
	for i := 0; i < round; i++ {
		err := pipe.Write(ctx, itemsPerRound...)
		if err != nil {
			t.Fatal(err)
		}
	}
	t.Logf("Pipe closing")
	pipe.Close()
	t.Logf("Pipe closed")
	wg.Wait()
	if recv != len(itemsPerRound)*round {
		t.Fatalf("Pipe expect %d items, got %d", len(itemsPerRound)*round, recv)
	}
}

func TestPipelineWriteCloseAndRead(t *testing.T) {
	ctx := context.Background()
	pipe := NewPipe[int]()
	pipeSize := 10
	var pipeRead int32
	var readWG sync.WaitGroup
	readWG.Add(1)
	go func() {
		defer readWG.Done()
		readBuf := make([]int, 1)
		for {
			n, err := pipe.Read(ctx, readBuf)
			if err != nil {
				if !errors.Is(err, ErrPipeEOF) {
					t.Errorf("un except err: %v", err)
				}
				break
			}
			atomic.AddInt32(&pipeRead, int32(n))
		}
	}()
	time.Sleep(time.Millisecond * 10) // let read goroutine start first
	for i := 0; i < pipeSize; i++ {
		err := pipe.Write(ctx, i)
		if err != nil {
			t.Error(err)
		}
	}
	for atomic.LoadInt32(&pipeRead) != int32(pipeSize) {
		runtime.Gosched()
	}
	pipe.Close()
	readWG.Wait()
}

func TestPipelineCancelCallback(t *testing.T) {
	pipeWrapper := struct{ pipe *Pipe[int] }{}
	finalNum := 100
	firstNum := 1
	pipe := NewPipe[int](WithCtxDoneCallback(
		func(ctx context.Context) {
			_ = pipeWrapper.pipe.Write(context.Background(), finalNum)
			pipeWrapper.pipe.Close()
		},
	))
	pipeWrapper.pipe = pipe
	var readWG sync.WaitGroup
	readWG.Add(1)
	go func() {
		defer readWG.Done()
		var canceled bool
		readBuf := make([]int, 1)
		rCtx, cancel := context.WithCancel(context.Background())
		defer cancel()
		for {
			n, rErr := pipe.Read(rCtx, readBuf)
			if canceled {
				if rErr != nil {
					test.Assert(t, rErr == io.EOF, rErr)
					return
				}
				test.Assert(t, n == 1, n)
				test.Assert(t, readBuf[0] == finalNum, readBuf[0])
				continue
			}
			test.Assert(t, n == 1, n)
			test.Assert(t, readBuf[0] == firstNum, readBuf[0])
			cancel()
			canceled = true
		}
	}()
	wErr := pipe.Write(context.Background(), firstNum)
	test.Assert(t, wErr == nil, wErr)
	readWG.Wait()
}

func BenchmarkPipeline(b *testing.B) {
	ctx := context.Background()
	pipe := NewPipe[int]()
	readCache := make([]int, 8)
	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		for j := 0; j < len(readCache); j++ {
			go pipe.Write(ctx, 1)
		}
		total := 0
		for total < len(readCache) {
			n, _ := pipe.Read(ctx, readCache)
			total += n
		}
	}
}
