package rpcserver

import (
	"context"
	"errors"
	"fmt"
	"net"

	"github.com/go-logr/logr"
	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	"google.golang.org/grpc"
	"google.golang.org/grpc/health"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/reflection"

	"github.com/tidbcloud/aisupport/packages/sdk-go/infra/rpcserver/middlewares"
)

type RPCHandler interface {
	Name() string
	RegisterRPC(server *grpc.Server) error
}

type RPCServerParam struct {
	// Address is the listening address for the RPC server, e.g. ":9090".
	Address string
	// Logger is used to log the RPC server internal info, e.g. error and warning.
	Logger logr.Logger
	// ServerOptions is the list of options used to create the gRPC server.
	ServerOptions []grpc.ServerOption
	// Tracing enables gRPC server tracing.
	Tracing bool
	// Handlers is the list of handlers that will be registered to the RPC server.
	Handlers []RPCHandler
}

type RPCServer struct {
	address      string
	server       *grpc.Server
	healthServer *health.Server
	logger       logr.Logger
	handlers     []RPCHandler
}

func NewRPCServer(param RPCServerParam) (*RPCServer, error) {
	serverOptions := param.ServerOptions
	serverOptions = append(serverOptions, middlewares.WrapErrorMiddleware())
	if param.Tracing {
		serverOptions = append(serverOptions, grpc.StatsHandler(otelgrpc.NewServerHandler()))
	}

	s := &RPCServer{
		address:  param.Address,
		server:   grpc.NewServer(serverOptions...),
		logger:   param.Logger.WithName("RPCServer"),
		handlers: param.Handlers,
	}

	s.healthServer = health.NewServer()
	healthpb.RegisterHealthServer(s.server, s.healthServer)
	reflection.Register(s.server)

	if err := s.registerHandlers(); err != nil {
		return nil, fmt.Errorf("init rpc server failed: %s", err)
	}

	return s, nil
}

func (s *RPCServer) registerHandlers() error {
	for _, handler := range s.handlers {
		s.logger.Info("register a handler", "handler", handler.Name())
		if err := handler.RegisterRPC(s.server); err != nil {
			return fmt.Errorf("register handler %s failed: %s", handler.Name(), err)
		}
	}

	return nil
}

// Run runs the RPC server and blocks until the RPC server exits.
func (s *RPCServer) Run() error {
	listener, err := net.Listen("tcp", s.address)
	if err != nil {
		return fmt.Errorf("failed to listen on %s: %s", s.address, err)
	}

	err = s.server.Serve(listener)
	if err != nil && !errors.Is(err, grpc.ErrServerStopped) {
		s.logger.Error(err, "rpc server is shutdown but got a error")
		return fmt.Errorf("failed to run rpc server: %s", err)
	}

	s.logger.Info("rpc server is shutdown")
	return nil
}

func (s *RPCServer) shutdownRPCServer(ctx context.Context) error {
	if err := ctx.Err(); err != nil {
		s.server.Stop()
		return err
	}

	stopped := make(chan struct{})
	go func() {
		s.server.GracefulStop()
		close(stopped)
	}()

	select {
	case <-stopped:
		return nil
	case <-ctx.Done():
		s.server.Stop()
		<-stopped
		return ctx.Err()
	}
}

// Shutdown stops the RPC server and waits for server to exit.
func (s *RPCServer) Shutdown(ctx context.Context) error {
	s.logger.Info("rpc server is shutting down by caller")
	s.healthServer.Shutdown()
	return s.shutdownRPCServer(ctx)
}