// Licensed to 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. Apache Software Foundation (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 gossip

import (
	"context"
	"fmt"
	"net"
	"runtime/debug"
	"strconv"
	"sync"
	"time"

	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/recovery"
	grpc_validator "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/validator"
	"github.com/pkg/errors"
	grpclib "google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/status"

	"github.com/apache/skywalking-banyandb/api/common"
	databasev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/database/v1"
	propertyv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/property/v1"
	"github.com/apache/skywalking-banyandb/banyand/metadata"
	"github.com/apache/skywalking-banyandb/banyand/metadata/schema"
	"github.com/apache/skywalking-banyandb/banyand/observability"
	"github.com/apache/skywalking-banyandb/banyand/queue"
	"github.com/apache/skywalking-banyandb/pkg/logger"
	"github.com/apache/skywalking-banyandb/pkg/node"
	"github.com/apache/skywalking-banyandb/pkg/partition"
	banyandbpath "github.com/apache/skywalking-banyandb/pkg/path"
	"github.com/apache/skywalking-banyandb/pkg/run"
)

// AddressExtractor extracts the gossip gRPC address from a node.
type AddressExtractor func(*databasev1.Node) string

var (
	errServerCert = errors.New("invalid server cert file")
	errServerKey  = errors.New("invalid server key file")
	errNoAddr     = errors.New("no address")
)

const (
	defaultRecvSize = 10 << 20

	defaultTotalTimeout = 20 * time.Minute
)

type service struct {
	schema.UnimplementedOnInitHandler
	addressExtractor    AddressExtractor
	pipeline            queue.Client
	metadata            metadata.Repo
	creds               credentials.TransportCredentials
	omr                 observability.MetricsRegistry
	sel                 node.Selector
	traceStreamSelector node.Selector
	ser                 *grpclib.Server
	log                 *logger.Logger
	closer              *run.Closer
	serverMetrics       *serverMetrics
	protocolHandler     *protocolHandler
	registered          map[string]*databasev1.Node
	traceSpanNotified   *int32
	prefix              string
	caCertPath          string
	host                string
	addr                string
	nodeID              string
	certFile            string
	keyFile             string
	traceSpanSending    []*recordTraceSpan
	listeners           []MessageListener
	serviceRegister     []func(s *grpclib.Server)
	traceEntityLocator  partition.Locator
	maxRecvMsgSize      run.Bytes
	totalTimeout        time.Duration
	scheduleInterval    time.Duration
	serviceRegisterLock sync.RWMutex
	traceSpanLock       sync.RWMutex
	listenersLock       sync.RWMutex
	mu                  sync.RWMutex
	port                uint32
	traceLogEnabled     bool
	tls                 bool
}

// NewMessenger creates a new instance of Messenger for gossip propagation communication between nodes.
func NewMessenger(prefix string, addressExtractor AddressExtractor,
	omr observability.MetricsRegistry, metadata metadata.Repo, pipeline queue.Client, defaultPort uint32,
) Messenger {
	serverScope := observability.RootScope.SubScope(prefix + "_gossip_server")
	return &service{
		prefix:           prefix,
		addressExtractor: addressExtractor,
		metadata:         metadata,
		omr:              omr,
		closer:           run.NewCloser(1),
		pipeline:         pipeline,
		serverMetrics:    newServerMetrics(omr.With(serverScope)),
		maxRecvMsgSize:   defaultRecvSize,
		totalTimeout:     defaultTotalTimeout,
		listeners:        make([]MessageListener, 0),
		registered:       make(map[string]*databasev1.Node),
		scheduleInterval: time.Hour * 2,
		sel:              node.NewRoundRobinSelector("", metadata),
		port:             defaultPort,
	}
}

// NewMessengerWithoutMetadata creates a new instance of Messenger without metadata.
func NewMessengerWithoutMetadata(prefix string, addressExtractor AddressExtractor, omr observability.MetricsRegistry, port int) Messenger {
	return NewMessenger(prefix, addressExtractor, omr, nil, nil, uint32(port))
}

func (s *service) PreRun(ctx context.Context) error {
	val := ctx.Value(common.ContextNodeKey)
	if val == nil {
		return errors.New("node id is empty")
	}
	node := val.(common.Node)
	s.nodeID = node.NodeID
	s.log = logger.GetLogger(s.prefix + "-gossip-messenger")
	s.listeners = make([]MessageListener, 0)
	metricsScope := observability.RootScope.SubScope(s.prefix + "_gossip_server")
	s.serverMetrics = newServerMetrics(s.omr.With(metricsScope))
	if s.metadata != nil {
		s.sel.OnInit([]schema.Kind{schema.KindGroup})
		s.metadata.RegisterHandler(s.prefix+"-nodes", schema.KindNode, s)
		s.metadata.RegisterHandler(s.prefix+"-groups", schema.KindGroup, s)
		if err := s.initTracing(ctx); err != nil {
			s.log.Warn().Err(err).Msg("failed to init internal trace stream")
		}
	}
	s.protocolHandler = newProtocolHandler(s)
	go s.protocolHandler.processPropagation()
	return nil
}

func (s *service) GetServerPort() *uint32 {
	return &s.port
}

func (s *service) Name() string {
	return s.prefix + "-gossip-messenger"
}

func (s *service) Role() databasev1.Role {
	return databasev1.Role_ROLE_UNSPECIFIED
}

func (s *service) FlagSet() *run.FlagSet {
	fs := run.NewFlagSet(s.prefix + "-gossip-messenger")

	fs.VarP(&s.maxRecvMsgSize, s.prefix+"-gossip-grpc-max-recv-msg-size", "", "the size of max receiving message")
	fs.StringVar(&s.host, s.prefix+"-gossip-grpc-host", "", "the host of banyand listens")
	fs.Uint32Var(&s.port, s.prefix+"-gossip-grpc-port", s.port, "the port of banyand listens")

	fs.BoolVar(&s.tls, s.prefix+"-gossip-grpc-tls", false, "connection uses TLS if true, else plain TCP")
	fs.StringVar(&s.certFile, s.prefix+"-gossip-grpc-server-cert-file", "", "the TLS cert file")
	fs.StringVar(&s.keyFile, s.prefix+"-gossip-grpc-server-key-file", "", "the TLS key file")
	fs.StringVar(&s.caCertPath, s.prefix+"-gossip-client-ca-cert", "", "Path to the CA certificate for gossip client TLS communication.")
	fs.DurationVar(&s.totalTimeout, s.prefix+"-gossip-total-timeout", defaultTotalTimeout, "the total timeout for gossip propagation")
	fs.BoolVar(&s.traceLogEnabled, s.prefix+"-gossip-trace-log", true, "enable trace log")
	return fs
}

func (s *service) Validate() error {
	var err error
	if s.caCertPath != "" {
		if s.caCertPath, err = banyandbpath.Get(s.caCertPath); err != nil {
			return err
		}
	}
	if s.certFile != "" {
		if s.certFile, err = banyandbpath.Get(s.certFile); err != nil {
			return err
		}
	}
	if s.keyFile != "" {
		if s.keyFile, err = banyandbpath.Get(s.keyFile); err != nil {
			return err
		}
	}

	// client side tls validation
	if s.tls && s.caCertPath == "" {
		return fmt.Errorf("CA certificate path must be provided when TLS is enabled")
	}

	// server side validation
	s.addr = net.JoinHostPort(s.host, strconv.FormatUint(uint64(s.port), 10))
	if s.addr == ":" {
		return errNoAddr
	}
	if !s.tls {
		return nil
	}
	if s.certFile == "" {
		return errServerCert
	}
	if s.keyFile == "" {
		return errServerKey
	}
	creds, errTLS := credentials.NewServerTLSFromFile(s.certFile, s.keyFile)
	if errTLS != nil {
		return errors.Wrap(errTLS, "failed to load cert and key")
	}
	s.creds = creds
	return nil
}

func (s *service) Serve(closer *run.Closer) {
	var opts []grpclib.ServerOption
	if s.tls {
		opts = []grpclib.ServerOption{grpclib.Creds(s.creds)}
	}
	grpcPanicRecoveryHandler := func(p any) (err error) {
		s.log.Error().Interface("panic", p).Str("stack", string(debug.Stack())).Msg("recovered from panic")

		return status.Errorf(codes.Internal, "%s", p)
	}
	streamChain := []grpclib.StreamServerInterceptor{
		recovery.StreamServerInterceptor(recovery.WithRecoveryHandler(grpcPanicRecoveryHandler)),
	}
	unaryChain := []grpclib.UnaryServerInterceptor{
		grpc_validator.UnaryServerInterceptor(),
		recovery.UnaryServerInterceptor(recovery.WithRecoveryHandler(grpcPanicRecoveryHandler)),
	}
	opts = append(opts, grpclib.MaxRecvMsgSize(int(s.maxRecvMsgSize)),
		grpclib.ChainUnaryInterceptor(unaryChain...),
		grpclib.ChainStreamInterceptor(streamChain...),
	)
	s.ser = grpclib.NewServer(opts...)

	propertyv1.RegisterGossipServiceServer(s.ser, s.protocolHandler)
	for _, register := range s.getServiceRegisters() {
		register(s.ser)
	}
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		lis, err := net.Listen("tcp", s.addr)
		if err != nil {
			s.log.Error().Err(err).Msg("Failed to listen")
			closer.CloseThenWait()
			return
		}
		s.log.Info().Str("addr", s.addr).Msg("Listening to")
		err = s.ser.Serve(lis)
		if err != nil {
			s.log.Error().Err(err).Msg("server is interrupted")
		}
		wg.Done()
	}()
	go func() {
		wg.Wait()
		s.log.Info().Msg("GRPC server is stopped")
	}()
}

// GetRouteTable implements RouteTableProvider interface.
// For gossip messenger, all registered nodes are considered active since gossip
// protocol doesn't track separate health states.
func (s *service) GetRouteTable() *databasev1.RouteTable {
	s.mu.RLock()
	defer s.mu.RUnlock()

	registered := make([]*databasev1.Node, 0, len(s.registered))

	for _, node := range s.registered {
		if node != nil {
			registered = append(registered, node)
		}
	}

	return &databasev1.RouteTable{
		Registered: registered,
		Active:     []string{},
		Evictable:  []string{},
	}
}

func (s *service) GracefulStop() {
	if s.ser == nil {
		return
	}
	s.closer.Done()
	s.closer.CloseThenWait()

	stopped := make(chan struct{})
	go func() {
		s.ser.GracefulStop()
		close(stopped)
	}()
	t := time.NewTimer(10 * time.Second)
	select {
	case <-t.C:
		s.ser.Stop()
		s.log.Info().Msg("force stopped")
	case <-stopped:
		t.Stop()
		s.log.Info().Msg("stopped gracefully")
	}
}
