/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 getty

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"net"
	"net/http"
	"os"
	"runtime"
	"strings"
	"sync"
	"time"
)

import (
	gxnet "github.com/dubbogo/gost/net"
	gxsync "github.com/dubbogo/gost/sync"
	gxtime "github.com/dubbogo/gost/time"

	"github.com/gorilla/websocket"

	perrors "github.com/pkg/errors"

	uatomic "go.uber.org/atomic"
)

import (
	log "github.com/AlexStocks/getty/util"
)

var (
	errSelfConnect        = perrors.New("connect self!")
	errServerClosed       = perrors.New("server closed")
	serverFastFailTimeout = time.Second * 1

	serverID uatomic.Int32
)

// Server accepts connections and manages its registered event loops.
// Close normally waits for those loops to exit. While any TCP/UDP session
// NewSessionCallback or OnOpen callback is active, Close
// initiates shutdown without waiting for the loops, even when called by another
// goroutine. Callbacks must return so initialization and loop cleanup can finish;
// call Close again after callbacks finish to wait. This is endpoint-wide activity tracking,
// not detection of the calling goroutine. WS/WSS callbacks run in HTTP handlers
// outside this wait group and retain the normal event-loop wait.
type Server interface {
	EndPoint
}

// StreamServer is like tcp/websocket/wss server
type StreamServer interface {
	Server
	// Listener get the network listener
	Listener() net.Listener
}

// PacketServer is like udp listen endpoint
type PacketServer interface {
	Server
	// PacketConn get the network listener
	PacketConn() net.PacketConn
}

type server struct {
	ServerOptions

	// endpoint ID
	endPointID EndPointID

	// net
	pktListener    net.PacketConn
	streamListener net.Listener
	lock           sync.RWMutex // for server
	endPointType   EndPointType
	server         *http.Server // for ws or wss server
	sync.Once
	done chan struct{}
	wg   sync.WaitGroup
	// initializing counts active TCP/UDP initialization callbacks under lock.
	initializing int
}

func (s *server) init(opts ...ServerOption) {
	for _, opt := range opts {
		opt(&(s.ServerOptions))
	}
}

func newServer(t EndPointType, opts ...ServerOption) *server {
	s := &server{
		endPointID:   serverID.Add(1),
		endPointType: t,
		done:         make(chan struct{}),
	}

	s.init(opts...)

	if t != UDP_ENDPOINT && s.sslEnabled && s.tlsConfigBuilder == nil {
		// #125: every endpoint whose listener goes through listenTCP - tcp, ws
		// and wss - builds the listener TLS configuration from this builder when
		// sslEnabled is set, so a missing builder used to kill the process with a
		// nil interface call from inside RunEventLoop/listen(). The udp endpoint
		// never reads sslEnabled and is left alone; ws/wss normally take their
		// certificate from WithWebsocketServerCert/WithWebsocketServerPrivateKey
		// and do not set sslEnabled at all.
		panic(fmt.Sprintf("server type:%s, sslEnabled is true but no tlsConfigBuilder was supplied; use WithServerTlsConfigBuilder, or drop sslEnabled for ws/wss and configure their certificate instead", t))
	}

	if (t == WS_SERVER || t == WSS_SERVER) && s.path == "" {
		// Registering an empty pattern panics inside http.ServeMux, and that
		// registration happens in runWSEventLoop's goroutine, so the process
		// went down with "http: invalid pattern" without naming the option that
		// was missing. Fail where the caller can see it, the way the cert/key
		// check in NewWSSServer does.
		panic(fmt.Sprintf("server type:%s, no websocket path was supplied; use WithWebsocketServerPath", t))
	}

	return s
}

// NewTCPServer builds a tcp server.
func NewTCPServer(opts ...ServerOption) Server {
	return newServer(TCP_SERVER, opts...)
}

// NewUDPEndPoint builds a unconnected udp server.
func NewUDPEndPoint(opts ...ServerOption) Server {
	return newServer(UDP_ENDPOINT, opts...)
}

// NewWSServer builds a websocket server.
func NewWSServer(opts ...ServerOption) Server {
	return newServer(WS_SERVER, opts...)
}

// NewWSSServer builds a secure websocket server.
func NewWSSServer(opts ...ServerOption) Server {
	s := newServer(WSS_SERVER, opts...)

	if s.addr == "" || s.cert == "" || s.privateKey == "" {
		panic(fmt.Sprintf("@addr:%s, @cert:%s, @privateKey:%s, @caCert:%s",
			s.addr, s.cert, s.privateKey, s.caCert))
	}

	return s
}

func (s *server) ID() int32 {
	return s.endPointID
}

func (s *server) EndPointType() EndPointType {
	return s.endPointType
}

func (s *server) stop() {
	select {
	case <-s.done:
		return
	default:
		s.Do(func() {
			close(s.done)
			s.lock.Lock()
			if s.server != nil {
				ctx, cancel := context.WithTimeout(context.Background(), serverFastFailTimeout)
				if err := s.server.Shutdown(ctx); err != nil {
					// if the log output is "shutdown ctx: context deadline exceeded"， it means that
					// there are still some active connections.
					log.Errorf("server shutdown ctx:%s error:%v", ctx, err)
				}
				cancel()
			}
			s.server = nil
			// Snapshot the listeners under s.lock. Keep the published listener
			// objects in place so event loops cannot race with a nil assignment.
			streamListener := s.streamListener
			pktListener := s.pktListener
			s.lock.Unlock()
			// close outside the lock to avoid blocking other lock holders.
			if streamListener != nil {
				// let the server exit asap when got error from RunEventLoop.
				_ = streamListener.Close()
			}
			if pktListener != nil {
				_ = pktListener.Close()
			}
		})
	}
}

func (s *server) GetTaskPool() gxsync.GenericTaskPool {
	return s.tPool
}

func (s *server) IsClosed() bool {
	select {
	case <-s.done:
		return true
	default:
		return false
	}
}

// net.ipv4.tcp_max_syn_backlog
// net.ipv4.tcp_timestamps
// net.ipv4.tcp_tw_recycle
func (s *server) listenTCP() error {
	var (
		err            error
		streamListener net.Listener
		sslConfig      *tls.Config
	)
	if s.IsClosed() {
		return errServerClosed
	}

	if s.sslEnabled {
		sslConfig, err = s.tlsConfigBuilder.BuildTlsConfig()
		if err != nil {
			return perrors.Wrapf(err, "BuildTlsConfig(addr:%s)", s.addr)
		}
		if sslConfig == nil {
			return fmt.Errorf("BuildTlsConfig returned nil config without error for addr:%s", s.addr)
		}
		// Match tls.Listen's validation: tls.NewListener does not check that
		// the config can supply a certificate. Reject it before opening a socket.
		if len(sslConfig.Certificates) == 0 && sslConfig.GetCertificate == nil && sslConfig.GetConfigForClient == nil {
			return fmt.Errorf("tls: neither Certificates, GetCertificate, nor GetConfigForClient set in Config")
		}
	}

	if len(s.addr) == 0 || !strings.Contains(s.addr, ":") {
		streamListener, err = gxnet.ListenOnTCPRandomPort(s.addr)
		if err != nil {
			return perrors.Wrapf(err, "gxnet.ListenOnTCPRandomPort(addr:%s)", s.addr)
		}
	} else {
		streamListener, err = net.Listen("tcp", s.addr)
		if err != nil {
			return perrors.Wrapf(err, "net.Listen(tcp, addr:%s)", s.addr)
		}
	}

	// Apply TLS to both explicit addresses and automatically assigned ports.
	if s.sslEnabled {
		streamListener = tls.NewListener(streamListener, sslConfig)
	}

	return s.publishStreamListener(streamListener)
}

func (s *server) publishStreamListener(streamListener net.Listener) error {
	addr := streamListener.Addr().String()
	s.lock.Lock()
	select {
	case <-s.done:
		s.lock.Unlock()
		_ = streamListener.Close()
		return errServerClosed
	default:
	}
	s.streamListener = streamListener
	s.addr = addr
	s.lock.Unlock()

	return nil
}

func (s *server) listenUDP() error {
	var (
		err         error
		localAddr   *net.UDPAddr
		pktListener *net.UDPConn
	)
	if s.IsClosed() {
		return errServerClosed
	}

	if len(s.addr) == 0 || !strings.Contains(s.addr, ":") {
		pktListener, err = gxnet.ListenOnUDPRandomPort(s.addr)
		if err != nil {
			return perrors.Wrapf(err, "gxnet.ListenOnUDPRandomPort(addr:%s)", s.addr)
		}
	} else {
		localAddr, err = net.ResolveUDPAddr("udp", s.addr)
		if err != nil {
			return perrors.Wrapf(err, "net.ResolveUDPAddr(udp, addr:%s)", s.addr)
		}
		pktListener, err = net.ListenUDP("udp", localAddr)
		if err != nil {
			return perrors.Wrapf(err, "net.ListenUDP((udp, localAddr:%#v)", localAddr)
		}
	}

	return s.publishPacketListener(pktListener)
}

func (s *server) publishPacketListener(pktListener net.PacketConn) error {
	addr := pktListener.LocalAddr().String()
	s.lock.Lock()
	select {
	case <-s.done:
		s.lock.Unlock()
		_ = pktListener.Close()
		return errServerClosed
	default:
	}
	s.pktListener = pktListener
	s.addr = addr
	s.lock.Unlock()

	return nil
}

// Listen announces on the local network address.
func (s *server) listen() error {
	switch s.endPointType {
	case TCP_SERVER, WS_SERVER, WSS_SERVER:
		return perrors.WithStack(s.listenTCP())
	case UDP_ENDPOINT:
		return perrors.WithStack(s.listenUDP())
	}

	return nil
}

// callNewSession invokes the user's NewSessionCallback and turns a panic from it
// into an error. It is called from goroutines this library started - the tcp
// accept loop, the udp endpoint loop, the ws handler and the client's reconnect
// loop - where nothing else recovers a panic, so a bug in application code would
// otherwise kill the whole process instead of closing the one connection it was
// called for. handlePackage already contains a panic from the user's OnMessage
// the same way.
func callNewSession(cb NewSessionCallback, ss Session) (err error) {
	defer func() {
		if r := recover(); r != nil {
			const size = 64 << 10
			rBuf := make([]byte, size)
			rBuf = rBuf[:runtime.Stack(rBuf, false)]
			err = perrors.Errorf("[newSession] panic: err=%v\n%s", r, rBuf)
		}
	}()

	return cb(ss)
}

// accept accepts a connection, configures its session and runs session startup.
// If initialization fails before run, it closes the accepted connection.
// Session startup handles OnOpen failure and shutdown through its own lifecycle.
func (s *server) accept(newSession NewSessionCallback) (Session, error) {
	conn, err := s.streamListener.Accept()
	if err != nil {
		return nil, perrors.WithStack(err)
	}
	if gxnet.IsSameAddr(conn.RemoteAddr(), conn.LocalAddr()) {
		log.Warnf("conn.localAddr{%s} == conn.RemoteAddr{%s}", conn.LocalAddr().String(), conn.RemoteAddr().String())
		// #123: the connection was accepted, so it owns a descriptor. Returning
		// without closing it leaks that fd until the process exits, because the
		// caller just keeps accepting. The client side already closes before it
		// reports errSelfConnect.
		_ = conn.Close()
		return nil, perrors.WithStack(errSelfConnect)
	}

	ss := newTCPSession(conn, s)
	err = s.initSession(newSession, ss)
	if err != nil {
		_ = conn.Close()
		return nil, perrors.WithStack(err)
	}

	return ss, nil
}

// invokeInitCallback tracks a callback on an event-loop goroutine, so it can
// close the server without waiting for that same goroutine to exit.
// It is also used by session.callOnOpen for TCP/UDP server sessions.
func (s *server) invokeInitCallback(f func() error) error {
	s.lock.Lock()
	if s.IsClosed() {
		s.lock.Unlock()
		return errServerClosed
	}
	s.initializing++
	s.lock.Unlock()
	defer func() {
		s.lock.Lock()
		s.initializing--
		s.lock.Unlock()
	}()

	return f()
}

// initSession configures and starts a TCP/UDP session.
// The caller closes the connection if this method returns an error.
// Once run starts, the session lifecycle handles OnOpen failure and shutdown.
func (s *server) initSession(newSession NewSessionCallback, ss Session) error {
	if err := s.invokeInitCallback(func() error { return callNewSession(newSession, ss) }); err != nil {
		return err
	}
	// A callback may have closed the server; do not start OnOpen afterwards.
	if s.IsClosed() {
		return errServerClosed
	}
	ss.(*session).run()
	return nil
}

// startEventLoop registers one event loop goroutine. RunEventLoop publishes the
// listener first (under s.lock) and only then reaches the run*EventLoop call
// that used to do the wg.Add on its own, so a Close() landing in between closed
// done and returned from wg.Wait() while the counter was still zero - and the
// event loop was registered afterwards anyway. Testing done and doing the Add
// under the same lock makes both atomic with respect to stop(): either the loop
// is registered before stop() closes done (and Wait() then waits for it), or
// the closed server refuses the registration and no goroutine is started.
func (s *server) startEventLoop() bool {
	s.lock.Lock()
	defer s.lock.Unlock()
	select {
	case <-s.done:
		return false
	default:
	}
	s.wg.Add(1)

	return true
}

func (s *server) runTCPEventLoop(newSession NewSessionCallback) {
	if !s.startEventLoop() {
		return
	}
	go func() {
		defer s.wg.Done()
		var (
			err   error
			delay time.Duration
		)
		for {
			if s.IsClosed() {
				log.Infof("server{%s} stop accepting client connect request.", s.addr)
				return
			}
			if delay != 0 {
				// #130: do not sit out the back-off when the server is closing -
				// Close() waits for this goroutine, and a full delay would also
				// let the loop call Accept once more before checking IsClosed().
				select {
				case <-s.done:
					return
				case <-gxtime.After(delay):
				}
			}
			_, err = s.accept(newSession)
			if err != nil {
				if perrors.Cause(err) == errServerClosed {
					return
				}
				//	Accept timeouts (per https://github.com/golang/go/issues/45729,
				//  Timeout() captures the retryable subset of the old Temporary()) are
				//  retried quietly; every other error - EMFILE/ENFILE from fd exhaustion,
				//  for instance - is reported.
				if netErr, ok := perrors.Cause(err).(net.Error); !ok || !netErr.Timeout() {
					log.Warnf("server{%s}.Accept() = err {%+v}", s.addr, perrors.WithStack(err))
				}
				// #130: any persistent accept error must back off, not only a timeout.
				//  Re-calling Accept at full speed spins the CPU and floods the log
				//  exactly when the process is out of resources.
				if delay == 0 {
					delay = 5 * time.Millisecond
				} else {
					delay *= 2
				}
				if max := 1 * time.Second; delay > max {
					delay = max
				}
				continue
			}
			delay = 0
		}
	}()
}

func (s *server) runUDPEventLoop(newSession NewSessionCallback) {
	if !s.startEventLoop() {
		return
	}
	go func() {
		defer s.wg.Done()
		var (
			err  error
			conn *net.UDPConn
			ss   Session
		)

		conn = s.pktListener.(*net.UDPConn)
		ss = newUDPSession(conn, s)
		if err = s.initSession(newSession, ss); err != nil {
			if perrors.Cause(err) == errServerClosed {
				_ = conn.Close()
				return
			}
			// #125: this runs in a goroutine the library started, so a panic here
			// is unrecoverable and takes the whole process down over a caller
			// error. The TCP accept path logs and keeps serving; do the same.
			log.Errorf("server{%s}.newSession(ss{%#v}) = err {%s}", s.addr, ss, perrors.WithStack(err))
			_ = conn.Close()
			return
		}
	}()
}

type wsHandler struct {
	http.ServeMux
	server     *server
	newSession NewSessionCallback
	upgrader   websocket.Upgrader
}

func newWSHandler(server *server, newSession NewSessionCallback) *wsHandler {
	return &wsHandler{
		server:     server,
		newSession: newSession,
		upgrader: websocket.Upgrader{
			// in default, ReadBufferSize & WriteBufferSize is 4k
			// HandshakeTimeout: server.HTTPTimeout,
			CheckOrigin:       func(_ *http.Request) bool { return true }, // allow connections from any origin
			EnableCompression: true,
		},
	}
}

func (s *wsHandler) serveWSRequest(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		// w.WriteHeader(http.StatusMethodNotAllowed)
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	if s.server.IsClosed() {
		http.Error(w, "HTTP server is closed(code:500-11).", http.StatusInternalServerError)
		log.Warnf("server{%s} stop acceptting client connect request.", s.server.addr)
		return
	}

	conn, err := s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Warnf("upgrader.Upgrader(http.Request{%#v}) = error:%+v", r, err)
		return
	}
	if conn.RemoteAddr().String() == conn.LocalAddr().String() {
		_ = conn.Close()
		log.Warnf("conn.localAddr{%s} == conn.RemoteAddr{%s}", conn.LocalAddr().String(), conn.RemoteAddr().String())
		return
	}
	// conn.SetReadLimit(int64(handler.maxMsgLen))
	ss := newWSSession(conn, s.server)
	err = callNewSession(s.newSession, ss)
	if err != nil {
		_ = conn.Close()
		log.Warnf("server{%s}.newSession(ss{%#v}) = err {%s}", s.server.addr, ss, err)
		return
	}
	if ss.(*session).maxMsgLen > 0 {
		conn.SetReadLimit(int64(ss.(*session).maxMsgLen))
	}
	ss.(*session).run()
}

// runWSEventLoop serve websocket client request
// @newSession: new websocket connection callback
func (s *server) runWSEventLoop(newSession NewSessionCallback) {
	if !s.startEventLoop() {
		return
	}
	go func() {
		defer s.wg.Done()
		var (
			err     error
			handler *wsHandler
			server  *http.Server
		)
		handler = newWSHandler(s, newSession)
		handler.HandleFunc(s.path, handler.serveWSRequest)
		server = &http.Server{
			Addr:    s.addr,
			Handler: handler,
			// ReadTimeout:    server.HTTPTimeout,
			// WriteTimeout:   server.HTTPTimeout,
		}
		s.lock.Lock()
		s.server = server
		s.lock.Unlock()
		err = server.Serve(s.streamListener)
		if err != nil && !errors.Is(err, http.ErrServerClosed) && !s.IsClosed() {
			// the wss branch has always filtered this: Shutdown() makes Serve
			// return ErrServerClosed, so a normal Close() logged an error level
			// line on every shutdown of a ws server.
			log.Errorf("http.server.Serve(addr{%s}) = err:%+v", s.addr, perrors.WithStack(err))
		}
	}()
}

// serve websocket client request
// RunWSSEventLoop serve websocket client request
func (s *server) runWSSEventLoop(newSession NewSessionCallback) {
	if !s.startEventLoop() {
		return
	}
	go func() {
		var (
			err         error
			certPem     []byte
			certificate tls.Certificate
			certPool    *x509.CertPool
			config      *tls.Config
			handler     *wsHandler
			server      *http.Server
		)
		defer s.wg.Done()

		if certificate, err = tls.LoadX509KeyPair(s.cert, s.privateKey); err != nil {
			panic(fmt.Sprintf("tls.LoadX509KeyPair(certs{%s}, privateKey{%s}) = err:%+v",
				s.cert, s.privateKey, perrors.WithStack(err)))
		}
		config = &tls.Config{
			InsecureSkipVerify: true, // do not verify peer certs
			ClientAuth:         tls.NoClientCert,
			NextProtos:         []string{"http/1.1"},
			Certificates:       []tls.Certificate{certificate},
		}

		if s.caCert != "" {
			certPem, err = os.ReadFile(s.caCert)
			if err != nil {
				panic(fmt.Errorf("os.ReadFile(certFile{%s}) = err:%+v", s.caCert, perrors.WithStack(err)))
			}
			certPool = x509.NewCertPool()
			if ok := certPool.AppendCertsFromPEM(certPem); !ok {
				panic("failed to parse root certificate file")
			}
			config.ClientCAs = certPool
			config.ClientAuth = tls.RequireAndVerifyClientCert
			config.InsecureSkipVerify = false
		}

		handler = newWSHandler(s, newSession)
		handler.HandleFunc(s.path, handler.serveWSRequest)
		server = &http.Server{
			Addr:    s.addr,
			Handler: handler,
			// ReadTimeout:    server.HTTPTimeout,
			// WriteTimeout:   server.HTTPTimeout,
		}
		server.SetKeepAlivesEnabled(true)
		s.lock.Lock()
		s.server = server
		s.lock.Unlock()
		err = server.Serve(tls.NewListener(s.streamListener, config))
		if err != nil && !errors.Is(err, http.ErrServerClosed) && !s.IsClosed() {
			log.Errorf("http.server.Serve(addr{%s}) = err:%+v", s.addr, perrors.WithStack(err))
		}
	}()
}

// RunEventLoop serves client request.
// @newSession: new connection callback
func (s *server) RunEventLoop(newSession NewSessionCallback) {
	if err := s.listen(); err != nil {
		if perrors.Cause(err) == errServerClosed {
			return
		}
		panic(fmt.Errorf("server.listen() = error:%+v", perrors.WithStack(err)))
	}

	switch s.endPointType {
	case TCP_SERVER:
		s.runTCPEventLoop(newSession)
	case UDP_ENDPOINT:
		s.runUDPEventLoop(newSession)
	case WS_SERVER:
		s.runWSEventLoop(newSession)
	case WSS_SERVER:
		s.runWSSEventLoop(newSession)
	default:
		panic(fmt.Sprintf("illegal server type %s", s.endPointType.String()))
	}
}

func (s *server) Listener() net.Listener {
	s.lock.RLock()
	defer s.lock.RUnlock()
	return s.streamListener
}

func (s *server) PacketConn() net.PacketConn {
	s.lock.RLock()
	defer s.lock.RUnlock()
	return s.pktListener
}

// Close stops accepting connections. It waits for registered event loops unless
// a TCP/UDP initialization callback is active; see Server for the caller contract.
func (s *server) Close() {
	s.stop()
	s.lock.RLock()
	initializing := s.initializing > 0
	s.lock.RUnlock()
	if !initializing {
		s.wg.Wait()
	}
}