/*
 * Copyright 2021 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 server .
package server

import (
	"context"
	"errors"
	"fmt"
	"net"
	"reflect"
	"runtime/debug"
	"sync"
	"time"

	"github.com/cloudwego/localsession/backup"

	internal_server "github.com/cloudwego/kitex/internal/server"
	"github.com/cloudwego/kitex/pkg/acl"
	"github.com/cloudwego/kitex/pkg/diagnosis"
	"github.com/cloudwego/kitex/pkg/discovery"
	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/endpoint/sep"
	"github.com/cloudwego/kitex/pkg/gofunc"
	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/limiter"
	"github.com/cloudwego/kitex/pkg/registry"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/bound"
	"github.com/cloudwego/kitex/pkg/remote/remotesvr"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/stats"
	"github.com/cloudwego/kitex/pkg/streaming"
)

// Server is an abstraction of an RPC server. It accepts connections and dispatches them to the service
// registered to it.
type Server interface {
	RegisterService(svcInfo *serviceinfo.ServiceInfo, handler interface{}, opts ...RegisterOption) error
	GetServiceInfos() map[string]*serviceinfo.ServiceInfo
	Run() error
	Stop() error
}

type server struct {
	opt  *internal_server.Options
	svcs *services

	// actual rpc service implement of biz
	eps     endpoint.Endpoint
	svr     remotesvr.Server
	stopped sync.Once
	isInit  bool
	isRun   bool

	sync.Mutex
}

// NewServer creates a server with the given Options.
func NewServer(ops ...Option) Server {
	s := &server{
		opt:  internal_server.NewOptions(ops),
		svcs: newServices(),
	}
	return s
}

func (s *server) init() {
	if s.isInit {
		return
	}
	s.isInit = true
	ctx := fillContext(s.opt)
	if ds := s.opt.DebugService; ds != nil {
		ds.RegisterProbeFunc(diagnosis.OptionsKey, diagnosis.WrapAsProbeFunc(s.opt.DebugInfo))
		ds.RegisterProbeFunc(diagnosis.ChangeEventsKey, s.opt.Events.Dump)
	}
	backup.Init(s.opt.BackupOpt)

	// init invoker chain here since we need to get some svc information to add MW
	// register stream recv/send middlewares
	s.buildInvokeChain(ctx)
}

func fillContext(opt *internal_server.Options) context.Context {
	ctx := context.Background()
	ctx = context.WithValue(ctx, endpoint.CtxEventBusKey, opt.Bus)
	ctx = context.WithValue(ctx, endpoint.CtxEventQueueKey, opt.Events)
	return ctx
}

func (s *server) initOrResetRPCInfoFunc() func(rpcinfo.RPCInfo, net.Addr) rpcinfo.RPCInfo {
	return func(ri rpcinfo.RPCInfo, rAddr net.Addr) rpcinfo.RPCInfo {
		// Reset existing rpcinfo to improve performance for long connections (PR #584).
		if ri != nil && rpcinfo.PoolEnabled() {
			fi := rpcinfo.AsMutableEndpointInfo(ri.From())
			fi.Reset()
			fi.SetAddress(rAddr)
			rpcinfo.AsMutableEndpointInfo(ri.To()).ResetFromBasicInfo(s.opt.Svr)
			if setter, ok := ri.Invocation().(rpcinfo.InvocationSetter); ok {
				setter.Reset()
			}
			rpcinfo.AsMutableRPCConfig(ri.Config()).CopyFrom(s.opt.Configs)
			rpcStats := rpcinfo.AsMutableRPCStats(ri.Stats())
			rpcStats.Reset()
			if s.opt.StatsLevel != nil {
				rpcStats.SetLevel(*s.opt.StatsLevel)
			}
			return ri
		}

		// allocate a new rpcinfo if it's the connection's first request or rpcInfoPool is disabled
		rpcStats := rpcinfo.AsMutableRPCStats(rpcinfo.NewRPCStats())
		if s.opt.StatsLevel != nil {
			rpcStats.SetLevel(*s.opt.StatsLevel)
		}

		// Export read-only views to external users and keep a mapping for internal users.
		ri = rpcinfo.NewRPCInfo(
			rpcinfo.EmptyEndpointInfo(),
			rpcinfo.FromBasicInfo(s.opt.Svr),
			rpcinfo.NewServerInvocation(),
			rpcinfo.AsMutableRPCConfig(s.opt.Configs).Clone().ImmutableView(),
			rpcStats.ImmutableView(),
		)
		rpcinfo.AsMutableEndpointInfo(ri.From()).SetAddress(rAddr)
		return ri
	}
}

func (s *server) buildMiddlewares(ctx context.Context) []endpoint.Middleware {
	// 1. unary middleware
	s.opt.UnaryOptions.InitMiddlewares(ctx)

	// 2. stream middleware
	s.opt.Streaming.InitMiddlewares(ctx)
	s.opt.StreamOptions.EventHandler = s.opt.TracerCtl.GetStreamEventHandler()
	s.opt.StreamOptions.InitMiddlewares(ctx)

	// 3. universal middleware
	var mws []endpoint.Middleware
	mws = append(mws, s.wrapStreamMiddleware())
	// register server timeout middleware
	// prepend for adding timeout to the context for all middlewares and the handler
	if s.opt.EnableContextTimeout {
		mws = append(mws, serverTimeoutMW)
	}
	// register server middlewares
	for i := range s.opt.MWBs {
		if mw := s.opt.MWBs[i](ctx); mw != nil {
			mws = append(mws, mw)
		}
	}
	// register core middleware,
	// core middleware MUST be the last middleware
	mws = append(mws, s.buildCoreMiddleware())
	return mws
}

func (s *server) buildInvokeChain(ctx context.Context) {
	mws := s.buildMiddlewares(ctx)
	s.eps = endpoint.Chain(mws...)(s.unaryOrStreamEndpoint(ctx))
}

// RegisterService should not be called by users directly.
func (s *server) RegisterService(svcInfo *serviceinfo.ServiceInfo, handler interface{}, opts ...RegisterOption) error {
	s.Lock()
	defer s.Unlock()
	if s.isRun {
		panic("service cannot be registered while server is running")
	}
	if svcInfo == nil {
		panic("svcInfo is nil. please specify non-nil svcInfo")
	}
	if handler == nil || reflect.ValueOf(handler).IsNil() {
		panic("handler is nil. please specify non-nil handler")
	}

	registerOpts := internal_server.NewRegisterOptions(opts)
	// register service
	if err := s.svcs.addService(svcInfo, handler, registerOpts); err != nil {
		panic(err.Error())
	}
	return nil
}

func (s *server) GetServiceInfos() map[string]*serviceinfo.ServiceInfo {
	return s.svcs.getKnownSvcInfoMap()
}

// Run runs the server.
func (s *server) Run() (err error) {
	s.Lock()
	s.isRun = true
	s.Unlock()
	s.init()
	if err = s.check(); err != nil {
		return err
	}
	svrCfg := s.opt.RemoteOpt
	addr := svrCfg.Address // should not be nil
	if s.opt.Proxy != nil {
		svrCfg.Address, err = s.opt.Proxy.Replace(addr)
		if err != nil {
			return
		}
	}

	s.registerDebugInfo()
	s.richRemoteOption()
	transHdlr, err := s.newSvrTransHandler()
	if err != nil {
		return err
	}
	svr, err := remotesvr.NewServer(s.opt.RemoteOpt, transHdlr)
	if err != nil {
		return err
	}
	s.Lock()
	s.svr = svr
	s.Unlock()

	// start profiler
	if s.opt.RemoteOpt.Profiler != nil {
		gofunc.GoFunc(context.Background(), func() {
			klog.Info("KITEX: server starting profiler")
			err := s.opt.RemoteOpt.Profiler.Run(context.Background())
			if err != nil {
				klog.Errorf("KITEX: server started profiler error: error=%s", err.Error())
			}
		})
	}

	errCh := svr.Start()
	select {
	case err = <-errCh:
		klog.Errorf("KITEX: server start error: error=%s", err.Error())
		return err
	default:
	}
	muStartHooks.Lock()
	for i := range onServerStart {
		go onServerStart[i]()
	}
	muStartHooks.Unlock()
	s.Lock()
	s.buildRegistryInfo(svr.Address())
	s.Unlock()

	if err = s.waitExit(errCh); err != nil {
		klog.Errorf("KITEX: received error and exit: error=%s", err.Error())
	}
	if e := s.Stop(); e != nil && err == nil {
		err = e
		klog.Errorf("KITEX: stop server error: error=%s", e.Error())
	}
	return
}

// Stop stops the server gracefully.
func (s *server) Stop() (err error) {
	s.stopped.Do(func() {
		s.Lock()
		defer s.Unlock()

		muShutdownHooks.Lock()
		for i := range onShutdown {
			onShutdown[i]()
		}
		muShutdownHooks.Unlock()

		if s.opt.RegistryInfo != nil {
			err = s.opt.Registry.Deregister(s.opt.RegistryInfo)
			s.opt.RegistryInfo = nil
		}
		if s.svr != nil {
			if e := s.svr.Stop(); e != nil {
				err = e
			}
			s.svr = nil
		}
	})
	return
}

// buildCoreMiddleware build the core middleware that include some default framework logic like error handler and ACL.
// the `next` function for core middleware should be the service handler itself without any wrapped middleware.
func (s *server) buildCoreMiddleware() endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		serviceHandler := next
		return func(ctx context.Context, req, resp interface{}) (err error) {
			// --- Before Next ---
			if len(s.opt.ACLRules) > 0 {
				if err = acl.ApplyRules(ctx, req, s.opt.ACLRules); err != nil {
					return err
				}
			}

			// --- Run Next ---
			// run service handler
			err = serviceHandler(ctx, req, resp)

			// --- After Next ---
			// error handler only catches the server handler's error.
			if s.opt.ErrHandle != nil && err != nil {
				err = s.opt.ErrHandle(ctx, err)
			}

			return err
		}
	}
}

func (s *server) unaryOrStreamEndpoint(ctx context.Context) endpoint.Endpoint {
	unaryEp := endpoint.UnaryChain(s.opt.UnaryOptions.UnaryMiddlewares...)(s.invokeHandleEndpoint())
	streamEp := sep.StreamChain(s.opt.StreamOptions.StreamMiddlewares...)(s.streamHandleEndpoint())

	return func(ctx context.Context, req, resp interface{}) (err error) {
		if st, ok := req.(*streaming.Args); ok {
			return streamEp(ctx, st.ServerStream)
		} else {
			return unaryEp(ctx, req, resp)
		}
	}
}

func (s *server) invokeHandleEndpoint() endpoint.UnaryEndpoint {
	return func(ctx context.Context, args, resp interface{}) (err error) {
		ri := rpcinfo.GetRPCInfo(ctx)
		ink := ri.Invocation()
		methodName := ink.MethodName()
		serviceName := ink.ServiceName()
		svc := s.svcs.getService(serviceName)
		svcInfo := svc.svcInfo
		if methodName == "" && svcInfo.ServiceName != serviceinfo.GenericService {
			return errors.New("method name is empty in rpcinfo, should not happen")
		}
		defer func() {
			if handlerErr := recover(); handlerErr != nil {
				err = kerrors.ErrPanic.WithCauseAndStack(
					fmt.Errorf(
						"[happened in biz handler, method=%s.%s, please check the panic at the server side] %s",
						svcInfo.ServiceName, methodName, handlerErr),
					string(debug.Stack()))
				rpcStats := rpcinfo.AsMutableRPCStats(ri.Stats())
				rpcStats.SetPanicked(err)
			}
			rpcinfo.Record(ctx, ri, stats.ServerHandleFinish, err)
			// clear session
			backup.ClearCtx()
		}()
		implHandlerFunc := ink.MethodInfo().Handler()
		rpcinfo.Record(ctx, ri, stats.ServerHandleStart, nil)
		// set session
		backup.BackupCtx(ctx)
		err = implHandlerFunc(ctx, svc.getHandler(methodName), args, resp)
		if err != nil {
			if bizErr, ok := kerrors.FromBizStatusError(err); ok {
				if setter, ok := ri.Invocation().(rpcinfo.InvocationSetter); ok {
					setter.SetBizStatusErr(bizErr)
					return nil
				}
			}
			err = kerrors.ErrBiz.WithCause(err)
		}
		return err
	}
}

func (s *server) streamHandleEndpoint() sep.StreamEndpoint {
	return func(ctx context.Context, st streaming.ServerStream) (err error) {
		ri := rpcinfo.GetRPCInfo(ctx)
		ink := ri.Invocation()
		methodName := ink.MethodName()
		serviceName := ink.ServiceName()
		svc := s.svcs.getService(serviceName)
		svcInfo := svc.svcInfo
		if methodName == "" && svcInfo.ServiceName != serviceinfo.GenericService {
			return errors.New("method name is empty in rpcinfo, should not happen")
		}
		defer func() {
			if handlerErr := recover(); handlerErr != nil {
				err = kerrors.ErrPanic.WithCauseAndStack(
					fmt.Errorf(
						"[happened in biz handler, method=%s.%s, please check the panic at the server side] %s",
						svcInfo.ServiceName, methodName, handlerErr),
					string(debug.Stack()))
				rpcStats := rpcinfo.AsMutableRPCStats(ri.Stats())
				rpcStats.SetPanicked(err)
			}
			rpcinfo.Record(ctx, ri, stats.ServerHandleFinish, err)
			// clear session
			backup.ClearCtx()
		}()
		implHandlerFunc := ink.MethodInfo().Handler()
		rpcinfo.Record(ctx, ri, stats.ServerHandleStart, nil)
		// set session
		backup.BackupCtx(ctx)
		args := &streaming.Args{ServerStream: st}
		if grpcStreamGetter, ok := st.(streaming.GRPCStreamGetter); ok {
			// for compatible with gRPC
			if grpcStream := grpcStreamGetter.GetGRPCStream(); grpcStream != nil {
				// use contextStream to wrap the original Stream and rewrite Context()
				// so that we can get this ctx by Stream.Context()
				args.Stream = contextStream{
					Stream: grpcStream,
					ctx:    ctx,
				}
			}
		}
		err = implHandlerFunc(ctx, svc.getHandler(methodName), args, nil)
		if err != nil {
			if bizErr, ok := kerrors.FromBizStatusError(err); ok {
				if setter, ok := ri.Invocation().(rpcinfo.InvocationSetter); ok {
					setter.SetBizStatusErr(bizErr)
					return nil
				}
			}
			err = kerrors.ErrBiz.WithCause(err)
		}
		return err
	}
}

func (s *server) initBasicRemoteOption() {
	remoteOpt := s.opt.RemoteOpt
	remoteOpt.SvcSearcher = s.svcs
	remoteOpt.InitOrResetRPCInfoFunc = s.initOrResetRPCInfoFunc()
	remoteOpt.TracerCtl = s.opt.TracerCtl
	remoteOpt.ReadWriteTimeout = s.opt.Configs.ReadWriteTimeout()
}

func (s *server) richRemoteOption() {
	s.initBasicRemoteOption()

	s.addBoundHandlers(s.opt.RemoteOpt)
}

func (s *server) addBoundHandlers(opt *remote.ServerOption) {
	// add profiler meta handler, which should be exec after other MetaHandlers
	if opt.Profiler != nil && opt.ProfilerMessageTagging != nil {
		s.opt.MetaHandlers = append(s.opt.MetaHandlers,
			remote.NewProfilerMetaHandler(opt.Profiler, opt.ProfilerMessageTagging),
		)
	}
	// for server trans info handler
	if len(s.opt.MetaHandlers) > 0 {
		transInfoHdlr := bound.NewTransMetaHandler(s.opt.MetaHandlers)
		// meta handler exec before boundHandlers which add with option
		doAddBoundHandlerToHead(transInfoHdlr, opt)
		for _, h := range s.opt.MetaHandlers {
			if shdlr, ok := h.(remote.StreamingMetaHandler); ok {
				opt.StreamingMetaHandlers = append(opt.StreamingMetaHandlers, shdlr)
			}
		}
	}

	limitHdlr := s.buildLimiterWithOpt()
	if limitHdlr != nil {
		doAddBoundHandler(limitHdlr, opt)
	}
}

/*
 * There are two times when the rate limiter can take effect for a non-multiplexed server,
 * which are the OnRead and OnMessage callback. OnRead is called before request decoded
 * and OnMessage is called after.
 * Therefore, the optimization point is that we can make rate limiter take effect in OnRead as
 * possible to save computational cost of decoding.
 * The implementation is that when using the default rate limiter to launching a non-multiplexed
 * service, use the `serverLimiterOnReadHandler` whose rate limiting takes effect in the OnRead
 * callback.
 */
func (s *server) buildLimiterWithOpt() (handler remote.InboundHandler) {
	limits := s.opt.Limit.Limits
	connLimit := s.opt.Limit.ConLimit
	qpsLimit := s.opt.Limit.QPSLimit
	if limits == nil && connLimit == nil && qpsLimit == nil {
		return
	}

	if connLimit == nil {
		if limits != nil {
			connLimit = limiter.NewConnectionLimiter(limits.MaxConnections)
		} else {
			connLimit = &limiter.DummyConcurrencyLimiter{}
		}
	}

	if qpsLimit == nil {
		if limits != nil {
			interval := time.Millisecond * 100 // FIXME: should not care this implementation-specific parameter
			qpsLimit = limiter.NewQPSLimiter(interval, limits.MaxQPS)
		} else {
			qpsLimit = &limiter.DummyRateLimiter{}
		}
	} else {
		s.opt.Limit.QPSLimitPostDecode = true
	}

	if limits != nil && limits.UpdateControl != nil {
		updater := limiter.NewLimiterWrapper(connLimit, qpsLimit)
		limits.UpdateControl(updater)
	}

	handler = bound.NewServerLimiterHandler(connLimit, qpsLimit, s.opt.Limit.LimitReporter, s.opt.Limit.QPSLimitPostDecode)
	// TODO: gRPC limiter
	return
}

func (s *server) check() error {
	return s.svcs.check(s.opt.RefuseTrafficWithoutServiceName)
}

func doAddBoundHandlerToHead(h remote.BoundHandler, opt *remote.ServerOption) {
	add := false
	if ih, ok := h.(remote.InboundHandler); ok {
		handlers := []remote.InboundHandler{ih}
		opt.Inbounds = append(handlers, opt.Inbounds...)
		add = true
	}
	if oh, ok := h.(remote.OutboundHandler); ok {
		handlers := []remote.OutboundHandler{oh}
		opt.Outbounds = append(handlers, opt.Outbounds...)
		add = true
	}
	if !add {
		panic("invalid BoundHandler: must implement InboundHandler or OutboundHandler")
	}
}

func doAddBoundHandler(h remote.BoundHandler, opt *remote.ServerOption) {
	add := false
	if ih, ok := h.(remote.InboundHandler); ok {
		opt.Inbounds = append(opt.Inbounds, ih)
		add = true
	}
	if oh, ok := h.(remote.OutboundHandler); ok {
		opt.Outbounds = append(opt.Outbounds, oh)
		add = true
	}
	if !add {
		panic("invalid BoundHandler: must implement InboundHandler or OutboundHandler")
	}
}

func (s *server) newSvrTransHandler() (handler remote.ServerTransHandler, err error) {
	transHdlrFactory := s.opt.RemoteOpt.SvrHandlerFactory
	transHdlr, err := transHdlrFactory.NewTransHandler(s.opt.RemoteOpt)
	if err != nil {
		return nil, err
	}
	if setter, ok := transHdlr.(remote.InvokeHandleFuncSetter); ok {
		setter.SetInvokeHandleFunc(s.eps)
	}
	transPl := remote.NewTransPipeline(transHdlr)

	for _, ib := range s.opt.RemoteOpt.Inbounds {
		transPl.AddInboundHandler(ib)
	}
	for _, ob := range s.opt.RemoteOpt.Outbounds {
		transPl.AddOutboundHandler(ob)
	}
	return transPl, nil
}

func (s *server) buildRegistryInfo(lAddr net.Addr) {
	if s.opt.RegistryInfo == nil {
		s.opt.RegistryInfo = &registry.Info{}
	}
	info := s.opt.RegistryInfo
	if !info.SkipListenAddr {
		// notice: lAddr may be nil when listen failed
		info.Addr = lAddr
	}
	if info.ServiceName == "" {
		info.ServiceName = s.opt.Svr.ServiceName
	}
	if info.PayloadCodec == "" {
		if svc := s.svcs.getTargetSvcInfo(); svc != nil {
			info.PayloadCodec = svc.PayloadCodec.String()
		}
	}
	if info.Weight == 0 {
		info.Weight = discovery.DefaultWeight
	}
	if info.Tags == nil {
		info.Tags = s.opt.Svr.Tags
	}
}

func (s *server) registerDebugInfo() {
	diagnosis.RegisterProbeFunc(s.opt.DebugService, diagnosis.ServiceInfosKey, diagnosis.WrapAsProbeFunc(s.svcs.getKnownSvcInfoMap()))
	if s.svcs.fallbackSvc != nil {
		diagnosis.RegisterProbeFunc(s.opt.DebugService, diagnosis.FallbackServiceKey, diagnosis.WrapAsProbeFunc(s.svcs.fallbackSvc.svcInfo.ServiceName))
	}
	if s.svcs.unknownSvc != nil {
		diagnosis.RegisterProbeFunc(s.opt.DebugService, diagnosis.UnknownServiceKey, diagnosis.WrapAsProbeFunc(nil))
	}
}

func (s *server) waitExit(errCh chan error) error {
	exitSignal := s.opt.ExitSignal()

	// service may not be available as soon as startup.
	delayRegister := time.After(1 * time.Second)
	for {
		select {
		case err := <-exitSignal:
			return err
		case err := <-errCh:
			return err
		case <-delayRegister:
			s.Lock()
			if err := s.opt.Registry.Register(s.opt.RegistryInfo); err != nil {
				s.Unlock()
				return err
			}
			s.Unlock()
		}
	}
}
