/*
 * 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 netpoll contains server and client implementation for netpoll.

package netpoll

import (
	"context"
	"errors"
	"net"
	"runtime/debug"
	"sync"
	"syscall"
	"time"

	"github.com/cloudwego/netpoll"

	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/trans"
	"github.com/cloudwego/kitex/pkg/utils"
)

// NewTransServerFactory creates a default netpoll transport server factory.
func NewTransServerFactory() remote.TransServerFactory {
	return &netpollTransServerFactory{}
}

type netpollTransServerFactory struct{}

// NewTransServer implements the remote.TransServerFactory interface.
func (f *netpollTransServerFactory) NewTransServer(opt *remote.ServerOption, transHdlr remote.ServerTransHandler) remote.TransServer {
	return &transServer{
		opt:       opt,
		transHdlr: transHdlr,
		lncfg:     trans.NewListenConfig(opt),
	}
}

type transServer struct {
	opt       *remote.ServerOption
	transHdlr remote.ServerTransHandler

	evl       netpoll.EventLoop
	ln        net.Listener
	lncfg     net.ListenConfig
	connCount utils.AtomicInt
	sync.Mutex
}

var _ remote.TransServer = &transServer{}

// CreateListener implements the remote.TransServer interface.
func (ts *transServer) CreateListener(addr net.Addr) (net.Listener, error) {
	if addr.Network() == "unix" {
		syscall.Unlink(addr.String())
	}
	// The network must be "tcp", "tcp4", "tcp6" or "unix".
	ln, err := ts.lncfg.Listen(context.Background(), addr.Network(), addr.String())
	return ln, err
}

// BootstrapServer implements the remote.TransServer interface.
func (ts *transServer) BootstrapServer(ln net.Listener) (err error) {
	if ln == nil {
		return errors.New("listener is nil in netpoll transport server")
	}
	ts.ln = ln
	opts := []netpoll.Option{
		netpoll.WithIdleTimeout(ts.opt.MaxConnectionIdleTime),
		netpoll.WithReadTimeout(ts.opt.ReadWriteTimeout),
	}

	ts.Lock()
	opts = append(opts, netpoll.WithOnPrepare(ts.onConnActive))
	ts.evl, err = netpoll.NewEventLoop(ts.onConnRead, opts...)
	ts.Unlock()

	if err == nil {
		// Convert the listener so that closing it also stops the
		// event loop in netpoll.
		ts.Lock()
		ts.ln, err = netpoll.ConvertListener(ts.ln)
		ts.Unlock()
	}
	if err != nil {
		return err
	}
	return ts.evl.Serve(ts.ln)
}

// Shutdown implements the remote.TransServer interface.
func (ts *transServer) Shutdown() (err error) {
	ts.Lock()
	defer ts.Unlock()

	ctx, cancel := context.WithTimeout(context.Background(), ts.opt.ExitWaitTime)
	defer cancel()
	if g, ok := ts.transHdlr.(remote.GracefulShutdown); ok {
		if ts.ln != nil {
			// 1. stop listener
			ts.ln.Close()

			// 2. signal all active connections to close gracefully
			err = g.GracefulShutdown(ctx)
			if err != nil {
				klog.Warnf("KITEX: server graceful shutdown error: %v", err)
			}
			// 3. wait some time to receive requests before closing idle conns
			/*
			   When the netpoll eventloop shutdown, all idle connections will be closed.
			   At this time, these connections may just receive requests, and then the peer side will report an EOF error.
			   To reduce such cases, wait for some time to try to receive these requests as much as possible,
			   so that the closing of connections can be controlled by the upper-layer protocol and the EOF problem can be reduced.
			*/
			time.Sleep(100 * time.Millisecond)
		}
	}
	if ts.evl != nil {
		err = ts.evl.Shutdown(ctx)
	}
	return
}

// ConnCount implements the remote.TransServer interface.
func (ts *transServer) ConnCount() utils.AtomicInt {
	return utils.AtomicInt(ts.connCount.Value())
}

// TODO: Is it necessary to init RPCInfo when connection is first created?
// Now the strategy is the same as the original one.
// 1. There's no request when the connection is first created.
// 2. Doesn't need to init RPCInfo if it's not RPC request, such as heartbeat.
func (ts *transServer) onConnActive(conn netpoll.Connection) context.Context {
	ctx := context.Background()
	defer transRecover(ctx, conn, "OnActive", false)
	conn.AddCloseCallback(func(connection netpoll.Connection) error {
		ts.onConnInactive(ctx, conn)
		return nil
	})
	ts.connCount.Inc()
	ctx, err := ts.transHdlr.OnActive(ctx, conn)
	if err != nil {
		ts.onError(ctx, err, conn)
		conn.Close()
		return ctx
	}
	return ctx
}

func (ts *transServer) onConnRead(ctx context.Context, conn netpoll.Connection) error {
	// in case it's panicked, it may caused by framework,
	// try to propagate the err and let it crash.
	// we mainly use transRecover for logging
	defer transRecover(ctx, conn, "onConnRead", true)
	err := ts.transHdlr.OnRead(ctx, conn)
	if err != nil {
		ts.onError(ctx, err, conn)
		if conn != nil {
			// close the connection if OnRead return error
			conn.Close()
		}
	}
	return nil
}

func (ts *transServer) onConnInactive(ctx context.Context, conn netpoll.Connection) {
	defer transRecover(ctx, conn, "OnInactive", false)
	ts.connCount.Dec()
	ts.transHdlr.OnInactive(ctx, conn)
}

func (ts *transServer) onError(ctx context.Context, err error, conn netpoll.Connection) {
	ts.transHdlr.OnError(ctx, err, conn)
}

func transRecover(ctx context.Context, conn netpoll.Connection, funcName string, propagatePanic bool) {
	panicErr := recover()
	if panicErr != nil {
		if conn != nil {
			klog.CtxErrorf(ctx, "KITEX: panic happened in %s, remoteAddress=%s, error=%v\nstack=%s", funcName, conn.RemoteAddr(), panicErr, string(debug.Stack()))
		} else {
			klog.CtxErrorf(ctx, "KITEX: panic happened in %s, error=%v\nstack=%s", funcName, panicErr, string(debug.Stack()))
		}
		if propagatePanic {
			panic(panicErr)
		}
	}
}
