/*
 * 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"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/cloudwego/kitex/client"
	"github.com/cloudwego/kitex/client/genericclient"
	kt "github.com/cloudwego/kitex/internal/mocks/thrift"
	"github.com/cloudwego/kitex/pkg/generic"
	"github.com/cloudwego/kitex/pkg/kerrors"
	kitex "github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/streaming"
	"github.com/cloudwego/kitex/server"
	"github.com/cloudwego/kitex/transport"
)

func newGenericStreamingClient(g generic.Generic, targetIPPort string) genericclient.Client {
	cli, err := genericclient.NewStreamingClient("destService", g,
		client.WithTransportProtocol(transport.GRPC),
		client.WithHostPorts(targetIPPort),
	)
	if err != nil {
		panic(err)
	}
	return cli
}

func newGenericClient(g generic.Generic, targetIPPort string) genericclient.Client {
	cli, err := genericclient.NewClient("destService", g,
		client.WithHostPorts(targetIPPort),
		client.WithTransportProtocol(transport.TTHeader))
	if err != nil {
		panic(err)
	}
	return cli
}

var _ kt.TestService = &StreamingTestImpl{}

type StreamingTestImpl struct{}

func (s *StreamingTestImpl) Echo(stream kt.TestService_EchoServer) (err error) {
	fmt.Println("Echo called")
	wg := &sync.WaitGroup{}
	wg.Add(2)

	go func() {
		defer func() {
			if p := recover(); p != nil {
				err = fmt.Errorf("panic: %v", p)
			}
			wg.Done()
		}()
		defer stream.Close()
		for {
			msg, recvErr := stream.Recv()
			if recvErr == io.EOF {
				return
			} else if recvErr != nil {
				err = recvErr
				return
			}
			fmt.Printf("Echo: received message = %s\n", msg.Message)
			time.Sleep(100 * time.Millisecond)
		}
	}()

	go func() {
		defer func() {
			if p := recover(); p != nil {
				err = fmt.Errorf("panic: %v", p)
			}
			wg.Done()
		}()
		resp := &kt.Response{}
		for i := 0; i < 3; i++ {
			resp.Message = fmt.Sprintf("%dth response", i)
			if sendErr := stream.Send(resp); sendErr != nil {
				err = sendErr
				return
			}
			fmt.Printf("Echo: sent message = %s\n", resp)
			time.Sleep(100 * time.Millisecond)
		}
	}()
	wg.Wait()
	return
}

func (s *StreamingTestImpl) EchoClient(stream kt.TestService_EchoClientServer) (err error) {
	fmt.Println("EchoClient called")
	var msgs []string
	for {
		req, err := stream.Recv()
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		fmt.Printf("Recv: %s\n", req.Message)
		msgs = append(msgs, req.Message)
		time.Sleep(100 * time.Millisecond)
	}
	return stream.SendAndClose(&kt.Response{Message: "all message: " + strings.Join(msgs, ", ")})
}

func (s *StreamingTestImpl) EchoServer(req *kt.Request, stream kt.TestService_EchoServerServer) (err error) {
	fmt.Println("EchoServer called")
	resp := &kt.Response{}
	for i := 0; i < 3; i++ {
		resp.Message = fmt.Sprintf("%v -> %dth response", req.Message, i)
		err := stream.Send(resp)
		if err != nil {
			return err
		}
		time.Sleep(100 * time.Millisecond)
	}
	return
}

func (s *StreamingTestImpl) EchoUnary(ctx context.Context, req *kt.Request) (resp *kt.Response, err error) {
	fmt.Println("EchoUnary called")
	resp = &kt.Response{}
	resp.Message = "hello " + req.Message
	return
}

func (s *StreamingTestImpl) EchoPingPong(ctx context.Context, req *kt.Request) (resp *kt.Response, err error) {
	fmt.Println("EchoPingPong called")
	resp = &kt.Response{}
	resp.Message = "hello " + req.Message
	return
}

func (s *StreamingTestImpl) EchoBizException(stream kt.TestService_EchoBizExceptionServer) error {
	fmt.Println("EchoBizException called")
	for {
		req, err := stream.Recv()
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		fmt.Printf("Recv: %s\n", req.Message)
		time.Sleep(100 * time.Millisecond)
	}
	return kerrors.NewBizStatusError(int32(404), "not found")
}

// normal server
func newMockServer(handler kt.TestService, addr net.Addr, opts ...server.Option) server.Server {
	var options []server.Option

	opts = append(opts, server.WithServiceAddr(addr))
	options = append(options, opts...)
	options = append(options, server.WithCompatibleMiddlewareForUnary())

	svr := server.NewServer(options...)
	if err := svr.RegisterService(serviceInfo(), handler); err != nil {
		panic(err)
	}
	go func() {
		err := svr.Run()
		if err != nil {
			panic(err)
		}
	}()
	return svr
}

var serviceMethods = map[string]kitex.MethodInfo{
	"Echo": kitex.NewMethodInfo(
		echoHandler,
		newTestServiceEchoArgs,
		newTestServiceEchoResult,
		false,
		kitex.WithStreamingMode(kitex.StreamingBidirectional),
	),
	"EchoClient": kitex.NewMethodInfo(
		echoClientHandler,
		newTestServiceEchoClientArgs,
		newTestServiceEchoClientResult,
		false,
		kitex.WithStreamingMode(kitex.StreamingClient),
	),
	"EchoServer": kitex.NewMethodInfo(
		echoServerHandler,
		newTestServiceEchoServerArgs,
		newTestServiceEchoServerResult,
		false,
		kitex.WithStreamingMode(kitex.StreamingServer),
	),
	"EchoUnary": kitex.NewMethodInfo(
		echoUnaryHandler,
		newTestServiceEchoUnaryArgs,
		newTestServiceEchoUnaryResult,
		false,
		kitex.WithStreamingMode(kitex.StreamingUnary),
	),
	"EchoBizException": kitex.NewMethodInfo(
		echoBizExceptionHandler,
		newTestServiceEchoBizExceptionArgs,
		newTestServiceEchoBizExceptionResult,
		false,
		kitex.WithStreamingMode(kitex.StreamingClient),
	),
	"EchoPingPong": kitex.NewMethodInfo(
		echoPingPongHandler,
		newTestServiceEchoPingPongArgs,
		newTestServiceEchoPingPongResult,
		false,
		kitex.WithStreamingMode(kitex.StreamingNone),
	),
}

func serviceInfo() *kitex.ServiceInfo {
	serviceName := "TestService"
	handlerType := (*kt.TestService)(nil)
	methods := map[string]kitex.MethodInfo{}
	for name, m := range serviceMethods {
		methods[name] = m
	}
	extra := map[string]interface{}{
		"PackageName": "thrift",
	}
	extra["streaming"] = true

	svcInfo := &kitex.ServiceInfo{
		ServiceName:     serviceName,
		HandlerType:     handlerType,
		Methods:         methods,
		PayloadCodec:    kitex.Thrift,
		KiteXGenVersion: "v0.11.0",
		Extra:           extra,
	}
	return svcInfo
}

func echoHandler(ctx context.Context, handler, arg, result interface{}) error {
	st, ok := arg.(*streaming.Args)
	if !ok {
		return errors.New("TestService.Echo is a thrift streaming method, please call with Kitex StreamClient")
	}
	stream := &testServiceEchoServer{st.Stream}
	return handler.(kt.TestService).Echo(stream)
}

type testServiceEchoServer struct {
	streaming.Stream
}

func (x *testServiceEchoServer) Send(m *kt.Response) error {
	return x.Stream.SendMsg(m)
}

func (x *testServiceEchoServer) Recv() (*kt.Request, error) {
	m := new(kt.Request)
	return m, x.Stream.RecvMsg(m)
}

func newTestServiceEchoArgs() interface{} {
	return kt.NewTestServiceEchoArgs()
}

func newTestServiceEchoResult() interface{} {
	return kt.NewTestServiceEchoResult()
}

func echoClientHandler(ctx context.Context, handler, arg, result interface{}) error {
	st, ok := arg.(*streaming.Args)
	if !ok {
		return errors.New("TestService.EchoClient is a thrift streaming method, please call with Kitex StreamClient")
	}
	stream := &testServiceEchoClientServer{st.Stream}
	return handler.(kt.TestService).EchoClient(stream)
}

type testServiceEchoClientServer struct {
	streaming.Stream
}

func (x *testServiceEchoClientServer) SendAndClose(m *kt.Response) error {
	return x.Stream.SendMsg(m)
}

func (x *testServiceEchoClientServer) Recv() (*kt.Request, error) {
	m := new(kt.Request)
	return m, x.Stream.RecvMsg(m)
}

func newTestServiceEchoClientArgs() interface{} {
	return kt.NewTestServiceEchoClientArgs()
}

func newTestServiceEchoClientResult() interface{} {
	return kt.NewTestServiceEchoClientResult()
}

func echoServerHandler(ctx context.Context, handler, arg, result interface{}) error {
	st, ok := arg.(*streaming.Args)
	if !ok {
		return errors.New("TestService.EchoServer is a thrift streaming method, please call with Kitex StreamClient")
	}
	stream := &testServiceEchoServerServer{st.Stream}
	req := new(kt.Request)
	if err := st.Stream.RecvMsg(req); err != nil {
		return err
	}
	return handler.(kt.TestService).EchoServer(req, stream)
}

type testServiceEchoServerServer struct {
	streaming.Stream
}

func (x *testServiceEchoServerServer) Send(m *kt.Response) error {
	return x.Stream.SendMsg(m)
}

func newTestServiceEchoServerArgs() interface{} {
	return kt.NewTestServiceEchoServerArgs()
}

func newTestServiceEchoServerResult() interface{} {
	return kt.NewTestServiceEchoServerResult()
}

func echoUnaryHandler(ctx context.Context, handler, arg, result interface{}) error {
	if streaming.GetStream(ctx) == nil {
		return errors.New("TestService.EchoUnary is a thrift streaming unary method, please call with Kitex StreamClient or remove the annotation streaming.mode")
	}
	realArg := arg.(*kt.TestServiceEchoUnaryArgs)
	realResult := result.(*kt.TestServiceEchoUnaryResult)
	success, err := handler.(kt.TestService).EchoUnary(ctx, realArg.Req)
	if err != nil {
		return err
	}
	realResult.Success = success
	return nil
}

func newTestServiceEchoUnaryArgs() interface{} {
	return kt.NewTestServiceEchoUnaryArgs()
}

func newTestServiceEchoUnaryResult() interface{} {
	return kt.NewTestServiceEchoUnaryResult()
}

func echoBizExceptionHandler(ctx context.Context, handler, arg, result interface{}) error {
	st, ok := arg.(*streaming.Args)
	if !ok {
		return errors.New("TestService.EchoBizException is a thrift streaming method, please call with Kitex StreamClient")
	}
	stream := &testServiceEchoBizExceptionServer{st.Stream}
	return handler.(kt.TestService).EchoBizException(stream)
}

type testServiceEchoBizExceptionServer struct {
	streaming.Stream
}

func (x *testServiceEchoBizExceptionServer) SendAndClose(m *kt.Response) error {
	return x.Stream.SendMsg(m)
}

func (x *testServiceEchoBizExceptionServer) Recv() (*kt.Request, error) {
	m := new(kt.Request)
	return m, x.Stream.RecvMsg(m)
}

func newTestServiceEchoBizExceptionArgs() interface{} {
	return kt.NewTestServiceEchoBizExceptionArgs()
}

func newTestServiceEchoBizExceptionResult() interface{} {
	return kt.NewTestServiceEchoBizExceptionResult()
}

func echoPingPongHandler(ctx context.Context, handler, arg, result interface{}) error {
	realArg := arg.(*kt.TestServiceEchoPingPongArgs)
	realResult := result.(*kt.TestServiceEchoPingPongResult)
	success, err := handler.(kt.TestService).EchoPingPong(ctx, realArg.Req)
	if err != nil {
		return err
	}
	realResult.Success = success
	return nil
}

func newTestServiceEchoPingPongArgs() interface{} {
	return kt.NewTestServiceEchoPingPongArgs()
}

func newTestServiceEchoPingPongResult() interface{} {
	return kt.NewTestServiceEchoPingPongResult()
}
