/*
 * 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 test

import (
	"context"
	"fmt"
	"io"
	"net"
	"reflect"
	"strings"
	"sync"
	"testing"

	dproto "github.com/cloudwego/dynamicgo/proto"
	"github.com/tidwall/gjson"

	"github.com/cloudwego/kitex/client"
	"github.com/cloudwego/kitex/client/genericclient"
	"github.com/cloudwego/kitex/internal/mocks/proto/kitex_gen/pbapi/mock"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/generic"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/streaming"
	"github.com/cloudwego/kitex/server"
)

func TestClientStreaming(t *testing.T) {
	ctx := context.Background()
	addr := test.GetLocalAddress()

	svr := initMockTestServer(new(StreamingTestImpl), addr)
	defer svr.Stop()

	cli := initStreamingClient(t, ctx, addr, "./idl/pbapi.proto")
	streamCli, err := genericclient.NewClientStreaming(ctx, cli, "ClientStreamingTest")
	test.Assert(t, err == nil, err)
	for i := 0; i < 3; i++ {
		req := fmt.Sprintf(`{"message": "grpc client streaming generic %dth request"}`, i)
		err = streamCli.Send(req)
		test.Assert(t, err == nil)
	}
	resp, err := streamCli.CloseAndRecv()
	test.Assert(t, err == nil)
	strResp, ok := resp.(string)
	test.Assert(t, ok)
	fmt.Printf("clientStreaming message received: %v\n", strResp)
	test.Assert(t, reflect.DeepEqual(gjson.Get(strResp, "message").String(),
		"all message: grpc client streaming generic 0th request, grpc client streaming generic 1th request, grpc client streaming generic 2th request"))
}

func TestServerStreaming(t *testing.T) {
	ctx := context.Background()
	addr := test.GetLocalAddress()

	svr := initMockTestServer(new(StreamingTestImpl), addr)
	defer svr.Stop()

	cli := initStreamingClient(t, ctx, addr, "./idl/pbapi.proto")
	streamCli, err := genericclient.NewServerStreaming(ctx, cli, "ServerStreamingTest", `{"message": "grpc server streaming generic request"}`)
	test.Assert(t, err == nil, err)
	for {
		resp, err := streamCli.Recv()
		if err != nil {
			test.Assert(t, err == io.EOF)
			fmt.Println("serverStreaming message receive done")
			break
		} else {
			strResp, ok := resp.(string)
			test.Assert(t, ok)
			fmt.Printf("serverStreaming message received: %s\n", strResp)
			test.Assert(t, strings.Contains(strResp, "grpc server streaming generic request ->"))
		}
	}
}

func TestBidirectionalStreaming(t *testing.T) {
	ctx := context.Background()
	addr := test.GetLocalAddress()

	svr := initMockTestServer(new(StreamingTestImpl), addr)
	defer svr.Stop()

	cli := initStreamingClient(t, ctx, addr, "./idl/pbapi.proto")
	streamCli, err := genericclient.NewBidirectionalStreaming(ctx, cli, "BidirectionalStreamingTest")
	test.Assert(t, err == nil)

	wg := &sync.WaitGroup{}
	wg.Add(2)

	go func() {
		defer wg.Done()
		defer streamCli.Close()
		for i := 0; i < 3; i++ {
			req := fmt.Sprintf(`{"message": "grpc bidirectional streaming generic %dth request"}`, i)
			err = streamCli.Send(req)
			test.Assert(t, err == nil)
			fmt.Printf("BidirectionalStreamingTest send: req = %s\n", req)
		}
	}()

	go func() {
		defer wg.Done()
		for {
			resp, err := streamCli.Recv()
			if err != nil {
				test.Assert(t, err == io.EOF)
				fmt.Println("bidirectionalStreaming message receive done")
				break
			} else {
				strResp, ok := resp.(string)
				test.Assert(t, ok)
				fmt.Printf("bidirectionalStreaming message received: %s\n", strResp)
				test.Assert(t, strings.Contains(strResp, "th response"))
			}
		}
	}()
	wg.Wait()
}

func TestUnary(t *testing.T) {
	ctx := context.Background()
	addr := test.GetLocalAddress()

	svr := initMockTestServer(new(StreamingTestImpl), addr)
	defer svr.Stop()

	cli := initStreamingClient(t, ctx, addr, "./idl/pbapi.proto")
	resp, err := cli.GenericCall(ctx, "UnaryTest", `{"message": "unary request"}`)
	test.Assert(t, err == nil)
	strResp, ok := resp.(string)
	test.Assert(t, ok)
	test.Assert(t, reflect.DeepEqual(gjson.Get(strResp, "message").String(), "hello unary request"))
}

func initStreamingClient(t *testing.T, ctx context.Context, addr, idl string, cliOpts ...client.Option) genericclient.Client {
	dOpts := dproto.Options{}
	p, err := generic.NewPbFileProviderWithDynamicGo(idl, ctx, dOpts)
	test.Assert(t, err == nil)
	g, err := generic.JSONPbGeneric(p)
	test.Assert(t, err == nil)
	return newGenericClient(g, addr, cliOpts...)
}

func initMockTestServer(handler mock.Mock, address string) server.Server {
	addr, _ := net.ResolveTCPAddr("tcp", address)
	svr := newMockTestServer(handler, addr)
	test.WaitServerStart(addr.String())
	return svr
}

func Test_invocationContainsPackage(t *testing.T) {
	ctx := context.Background()
	addr := test.GetLocalAddress()

	svr := initMockTestServer(new(StreamingTestImpl), addr)
	defer svr.Stop()
	cli := initStreamingClient(t, ctx, addr, "./idl/pbapi.proto", client.WithMiddleware(func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, req, resp interface{}) (err error) {
			ri := rpcinfo.GetRPCInfo(ctx)
			// make sure the package name has been injected into Invocation
			test.Assert(t, ri.Invocation().PackageName() == "pbapi", ri.Invocation())
			resp.(*streaming.Result).ClientStream = mockClientStream{ctx: ctx}
			return nil
		}
	}))
	_, err := genericclient.NewClientStreaming(ctx, cli, "ClientStreamingTest")
	test.Assert(t, err == nil, err)
}

type mockClientStream struct {
	ctx context.Context
	streaming.ClientStream
}

func (m mockClientStream) GetGRPCStream() streaming.Stream {
	return &mockGRPCStream{ctx: m.ctx}
}

func (m mockClientStream) Context() context.Context {
	return m.ctx
}

type mockGRPCStream struct {
	ctx context.Context
	streaming.Stream
}

func (m mockGRPCStream) Context() context.Context {
	return m.ctx
}
