// 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 sub

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

	"github.com/go-chi/chi/v5"
	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/recovery"
	grpc_validator "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/validator"
	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
	"github.com/pkg/errors"
	grpclib "google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/health"
	"google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/status"

	"github.com/apache/skywalking-banyandb/api/common"
	clusterv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/cluster/v1"
	databasev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/database/v1"
	fodcv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/fodc/v1"
	measurev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/measure/v1"
	streamv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/stream/v1"
	tracev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/trace/v1"
	"github.com/apache/skywalking-banyandb/banyand/liaison/grpc/route"
	"github.com/apache/skywalking-banyandb/banyand/metadata"
	"github.com/apache/skywalking-banyandb/banyand/metadata/schema/property"
	"github.com/apache/skywalking-banyandb/banyand/observability"
	"github.com/apache/skywalking-banyandb/banyand/queue"
	"github.com/apache/skywalking-banyandb/pkg/bus"
	"github.com/apache/skywalking-banyandb/pkg/healthcheck"
	"github.com/apache/skywalking-banyandb/pkg/logger"
	"github.com/apache/skywalking-banyandb/pkg/meter"
	"github.com/apache/skywalking-banyandb/pkg/panicdiag"
	banyandbpath "github.com/apache/skywalking-banyandb/pkg/path"
	"github.com/apache/skywalking-banyandb/pkg/run"
	pkgtls "github.com/apache/skywalking-banyandb/pkg/tls"
)

const defaultRecvSize = 10 << 20

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

	_ run.PreRunner = (*server)(nil)
	_ run.Service   = (*server)(nil)

	queueSubScope = observability.RootScope.SubScope("queue_sub")
)

type server struct {
	clusterv1.UnimplementedChunkedSyncServiceServer
	streamv1.UnimplementedStreamServiceServer
	databasev1.UnimplementedSnapshotServiceServer
	databasev1.UnimplementedNodeQueryServiceServer
	databasev1.UnimplementedClusterStateServiceServer
	clusterv1.UnimplementedServiceServer
	fodcv1.UnimplementedGroupLifecycleServiceServer
	omr                   observability.MetricsRegistry
	creds                 credentials.TransportCredentials
	metadataRepo          metadata.Repo
	nodeSchemaStatusRepo  metadata.Service
	clientCloser          context.CancelFunc
	ser                   *grpclib.Server
	listeners             map[bus.Topic][]bus.MessageListener
	topicMap              map[string]bus.Topic
	chunkedSyncHandlers   map[bus.Topic]queue.ChunkedSyncHandler
	log                   *logger.Logger
	httpSrv               *http.Server
	tlsReloader           *pkgtls.Reloader
	metrics               *metrics
	routeTableProvider    map[string]route.TableProvider
	curNode               *databasev1.Node
	flagNamePrefix        string
	certFile              string
	addr                  string
	httpAddr              string
	host                  string
	keyFile               string
	chunkBufferTimeout    time.Duration
	maxRecvMsgSize        run.Bytes
	listenersLock         sync.RWMutex
	routeTableProviderMu  sync.RWMutex
	port                  uint32
	httpPort              uint32
	maxChunkBufferSize    uint32
	maxChunkGapSize       uint32
	tls                   bool
	enableChunkReordering bool
}

// NewServer returns a new gRPC server.
func NewServer(omr observability.MetricsRegistry) queue.Server {
	return NewServerWithPorts(omr, "", 17912, 17913)
}

// NewServerWithPorts returns a new gRPC server with specified ports.
func NewServerWithPorts(omr observability.MetricsRegistry, flagNamePrefix string, port, httpPort uint32) queue.Server {
	srv := &server{
		listeners:           make(map[bus.Topic][]bus.MessageListener),
		topicMap:            make(map[string]bus.Topic),
		chunkedSyncHandlers: make(map[bus.Topic]queue.ChunkedSyncHandler),
		omr:                 omr,
		maxRecvMsgSize:      defaultRecvSize,
		flagNamePrefix:      flagNamePrefix,
		port:                port,
		httpPort:            httpPort,
		// Default chunk ordering configuration
		enableChunkReordering: true,
		maxChunkBufferSize:    10,
		chunkBufferTimeout:    5 * time.Second,
		maxChunkGapSize:       5,
	}
	return srv
}

func (s *server) PreRun(ctx context.Context) error {
	s.log = logger.GetLogger("server-queue-sub")
	s.metrics = newMetrics(s.omr.With(queueSubScope))

	// Initialize TLS reloader if TLS is enabled
	if s.tls {
		var err error
		if s.certFile, err = banyandbpath.Get(s.certFile); err != nil {
			return errors.Wrap(err, "failed to get absolute path for certFile")
		}
		if s.keyFile, err = banyandbpath.Get(s.keyFile); err != nil {
			return errors.Wrap(err, "failed to get absolute path for keyFile")
		}
		s.tlsReloader, err = pkgtls.NewReloader(s.certFile, s.keyFile, s.log)
		if err != nil {
			return errors.Wrap(err, "failed to initialize TLS reloader for queue server")
		}
		s.log.Info().Str("certFile", s.certFile).Str("keyFile", s.keyFile).Msg("Initialized TLS reloader for queue server")
	}

	nodeVal := ctx.Value(common.ContextNodeKey)
	roleVal := ctx.Value(common.ContextNodeRolesKey)
	if nodeVal == nil || roleVal == nil {
		s.log.Warn().Msg("node or role value not found in context")
		return nil
	}
	nodeRoles := roleVal.([]databasev1.Role)
	node := nodeVal.(common.Node)
	s.curNode = node.ToProtoNode(nodeRoles)

	return nil
}

func (s *server) Name() string {
	return "server-queue"
}

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

func (s *server) GetPort() *uint32 {
	return &s.port
}

func (s *server) FlagSet() *run.FlagSet {
	fs := run.NewFlagSet("grpc")
	prefixFlag := func(name string) string {
		if s.flagNamePrefix == "" {
			return name
		}
		return s.flagNamePrefix + "-" + name
	}
	fs.VarP(&s.maxRecvMsgSize, prefixFlag("max-recv-msg-size"), "", "the size of max receiving message")
	fs.BoolVar(&s.tls, prefixFlag("tls"), false, "connection uses TLS if true, else plain TCP")
	fs.StringVar(&s.certFile, prefixFlag("cert-file"), "", "the TLS cert file")
	fs.StringVar(&s.keyFile, prefixFlag("key-file"), "", "the TLS key file")
	fs.StringVar(&s.host, prefixFlag("grpc-host"), "", "the host of banyand listens")
	fs.Uint32Var(&s.port, prefixFlag("grpc-port"), s.port, "the port of banyand listens")
	fs.Uint32Var(&s.httpPort, prefixFlag("http-port"), s.httpPort, "the port of banyand http api listens")
	// Chunk ordering configuration flags
	fs.BoolVar(&s.enableChunkReordering, prefixFlag("enable-chunk-reordering"), s.enableChunkReordering, "enable out-of-order chunk handling")
	fs.Uint32Var(&s.maxChunkBufferSize, prefixFlag("max-chunk-buffer-size"), s.maxChunkBufferSize, "maximum chunks to buffer for reordering")
	fs.DurationVar(&s.chunkBufferTimeout, prefixFlag("chunk-buffer-timeout"), s.chunkBufferTimeout, "maximum time to wait for missing chunks")
	fs.Uint32Var(&s.maxChunkGapSize, prefixFlag("max-chunk-gap-size"), s.maxChunkGapSize, "maximum gap in chunk sequence")
	return fs
}

func (s *server) Validate() error {
	if s.port == 0 {
		s.port = 17912
	}
	if s.httpPort == 0 {
		s.httpPort = 17913
	}
	s.addr = net.JoinHostPort(s.host, strconv.FormatUint(uint64(s.port), 10))
	if s.addr == ":" {
		return errNoAddr
	}
	s.httpAddr = net.JoinHostPort(s.host, strconv.FormatUint(uint64(s.httpPort), 10))
	if s.httpAddr == ":" {
		return errNoAddr
	}
	if !s.tls {
		return nil
	}
	if s.certFile == "" {
		return errServerCert
	}
	if s.keyFile == "" {
		return errServerKey
	}
	// TLS reloader will be initialized in PreRun, so we don't need to load credentials here
	// The credentials will be loaded dynamically when the server starts
	return nil
}

func (s *server) Serve() run.StopNotify {
	var opts []grpclib.ServerOption
	if s.tls {
		// Start TLS reloader if enabled
		if s.tlsReloader != nil {
			if err := s.tlsReloader.Start(context.Background()); err != nil {
				s.log.Error().Err(err).Msg("Failed to start TLS reloader for queue server")
				stopCh := make(chan struct{})
				close(stopCh)
				return stopCh
			}
			s.log.Info().Str("certFile", s.certFile).Str("keyFile", s.keyFile).Msg("Started TLS file monitoring for queue server")
			tlsConfig := s.tlsReloader.GetTLSConfig()
			creds := credentials.NewTLS(tlsConfig)
			opts = []grpclib.ServerOption{grpclib.Creds(creds)}
		} else {
			// Fallback to static loading if reloader is not available
			creds, errTLS := credentials.NewServerTLSFromFile(s.certFile, s.keyFile)
			if errTLS != nil {
				s.log.Error().Err(errTLS).Msg("Failed to load TLS credentials")
				stopCh := make(chan struct{})
				close(stopCh)
				return stopCh
			}
			opts = []grpclib.ServerOption{grpclib.Creds(creds)}
		}
	}
	grpcPanicRecoveryHandler := func(ctx context.Context, p any) (err error) {
		breadcrumbs := panicdiag.BreadcrumbsFromContext(ctx)
		stages := make([]string, len(breadcrumbs))
		for idx, bc := range breadcrumbs {
			stages[idx] = bc.Stage
		}
		s.log.Error().Interface("panic", p).Str("stack", string(debug.Stack())).Strs("breadcrumbs", stages).Msg("recovered from panic")
		return status.Errorf(codes.Internal, "%s", p)
	}

	streamChain := []grpclib.StreamServerInterceptor{
		panicdiag.BreadcrumbStreamInterceptor(),
		recovery.StreamServerInterceptor(recovery.WithRecoveryHandlerContext(grpcPanicRecoveryHandler)),
	}
	unaryChain := []grpclib.UnaryServerInterceptor{
		panicdiag.BreadcrumbUnaryInterceptor(),
		grpc_validator.UnaryServerInterceptor(),
		recovery.UnaryServerInterceptor(recovery.WithRecoveryHandlerContext(grpcPanicRecoveryHandler)),
	}

	opts = append(opts, grpclib.MaxRecvMsgSize(int(s.maxRecvMsgSize)),
		grpclib.ChainUnaryInterceptor(unaryChain...),
		grpclib.ChainStreamInterceptor(streamChain...),
	)
	s.ser = grpclib.NewServer(opts...)
	clusterv1.RegisterServiceServer(s.ser, s)
	clusterv1.RegisterChunkedSyncServiceServer(s.ser, s)
	grpc_health_v1.RegisterHealthServer(s.ser, health.NewServer())
	databasev1.RegisterSnapshotServiceServer(s.ser, s)
	databasev1.RegisterNodeQueryServiceServer(s.ser, s)
	databasev1.RegisterClusterStateServiceServer(s.ser, s)
	streamv1.RegisterStreamServiceServer(s.ser, &streamService{ser: s})
	measurev1.RegisterMeasureServiceServer(s.ser, &measureService{ser: s})
	tracev1.RegisterTraceServiceServer(s.ser, &traceService{ser: s})
	if s.metadataRepo != nil {
		fodcv1.RegisterGroupLifecycleServiceServer(s.ser, s)
	}
	if s.nodeSchemaStatusRepo != nil {
		// The registry is resolved per request (closure) so the service
		// registers even when SchemaRegistry isn't ready at Serve time;
		// the fail-closed nil-cache contract handles in-flight probes
		// during initialization.
		svc := s.nodeSchemaStatusRepo
		clusterv1.RegisterNodeSchemaStatusServiceServer(s.ser, property.NewNodeSchemaStatusServerForRegistryWithNodeRepo(
			func() *property.SchemaRegistry {
				reg, regOk := svc.SchemaRegistry().(*property.SchemaRegistry)
				if !regOk {
					return nil
				}
				return reg
			},
			svc.NodeRepoRegistry,
		))
	}

	var ctx context.Context
	ctx, s.clientCloser = context.WithCancel(context.Background())
	clientOpts := make([]grpclib.DialOption, 0, 1)
	if s.creds == nil {
		clientOpts = append(clientOpts, grpclib.WithTransportCredentials(insecure.NewCredentials()))
	} else {
		clientOpts = append(clientOpts, grpclib.WithTransportCredentials(s.creds))
	}
	stopCh := make(chan struct{})
	client, err := healthcheck.NewClient(ctx, s.log, s.addr, clientOpts)
	if err != nil {
		s.log.Error().Err(err).Msg("Failed to health check client")
		close(stopCh)
		return stopCh
	}
	gwMux := runtime.NewServeMux(runtime.WithHealthzEndpoint(client))
	if err := databasev1.RegisterSnapshotServiceHandlerFromEndpoint(ctx, gwMux, s.addr, clientOpts); err != nil {
		s.log.Error().Err(err).Msg("Failed to register snapshot service")
		close(stopCh)
		return stopCh
	}
	if err := databasev1.RegisterClusterStateServiceHandlerFromEndpoint(ctx, gwMux, s.addr, clientOpts); err != nil {
		s.log.Error().Err(err).Msg("Failed to register cluster state service")
		close(stopCh)
		return stopCh
	}
	mux := chi.NewRouter()
	mux.Mount("/api", http.StripPrefix("/api", gwMux))
	// Mount /metrics when the metrics registry exposes a Prometheus handler.
	// The data node's HTTP server doubles as a Prometheus scrape target, so
	// monitoring tooling (and integration tests) can read banyandb_queue_sub_*
	// directly. The handler is only registered when omr is wired to a real
	// registry; BypassRegistry / nil registries leave /metrics absent (404).
	if reg, ok := s.omr.(observability.PrometheusHandlerProvider); ok {
		mux.Handle("/metrics", reg.PrometheusHandler())
	}
	s.httpSrv = &http.Server{
		Addr:              s.httpAddr,
		Handler:           mux,
		ReadHeaderTimeout: 3 * time.Second,
	}
	var wg sync.WaitGroup
	wg.Add(2)
	run.Go(context.Background(), "queue.sub.grpc-server", s.log, func(_ context.Context) {
		lis, err := net.Listen("tcp", s.addr)
		if err != nil {
			s.log.Error().Err(err).Msg("Failed to listen")
			close(stopCh)
			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()
	})
	run.Go(context.Background(), "queue.sub.healthz-http", s.log, func(_ context.Context) {
		s.log.Info().Str("listenAddr", s.httpAddr).Msg("Start healthz http server")
		err := s.httpSrv.ListenAndServe()
		if err != http.ErrServerClosed {
			s.log.Error().Err(err)
		}
		wg.Done()
	})
	run.Go(context.Background(), "queue.sub.shutdown-watcher", s.log, func(_ context.Context) {
		wg.Wait()
		s.log.Info().Msg("All servers are stopped")
		close(stopCh)
	})
	return stopCh
}

func (s *server) GracefulStop() {
	s.log.Info().Msg("stopping")

	// Stop TLS reloader if enabled
	if s.tlsReloader != nil {
		s.tlsReloader.Stop()
	}

	stopped := make(chan struct{})
	s.clientCloser()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_ = s.httpSrv.Shutdown(ctx)
	run.Go(context.Background(), "queue.sub.graceful-stop", s.log, func(_ context.Context) {
		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")
	}
}

// RegisterChunkedSyncHandler implements queue.Server.
func (s *server) RegisterChunkedSyncHandler(topic bus.Topic, handler queue.ChunkedSyncHandler) {
	s.chunkedSyncHandlers[topic] = handler
}

func (s *server) SetRouteProviders(providers map[string]route.TableProvider) {
	s.routeTableProviderMu.Lock()
	s.routeTableProvider = providers
	s.routeTableProviderMu.Unlock()
}

// SetMetadataRepo sets the metadata repository for the internal gRPC server.
func (s *server) SetMetadataRepo(repo metadata.Repo) {
	s.metadataRepo = repo
}

// SetNodeSchemaStatusRepo wires the metadata.Service whose SchemaRegistry +
// NodeRepoRegistry back the per-node NodeSchemaStatusService.
func (s *server) SetNodeSchemaStatusRepo(svc metadata.Service) {
	s.nodeSchemaStatusRepo = svc
}

type metrics struct {
	totalStarted         meter.Counter
	totalFinished        meter.Counter
	totalLatency         meter.Histogram
	totalErr             meter.Counter
	receivedBytes        meter.Counter
	totalBatchStarted    meter.Counter
	totalBatchFinished   meter.Counter
	totalBatchLatency    meter.Histogram
	totalMessageStarted  meter.Counter
	totalMessageFinished meter.Counter
}

func newMetrics(factory observability.Factory) *metrics {
	labels := []string{"operation", "group", "remote_node", "remote_role", "remote_tier"}
	return &metrics{
		totalStarted:         factory.NewCounter("total_started", labels...),
		totalFinished:        factory.NewCounter("total_finished", labels...),
		totalLatency:         factory.NewHistogram("total_latency", meter.DefBuckets, labels...),
		totalErr:             factory.NewCounter("total_err", append(labels, "error_type")...),
		receivedBytes:        factory.NewCounter("received_bytes", labels...),
		totalBatchStarted:    factory.NewCounter("total_batch_started", labels...),
		totalBatchFinished:   factory.NewCounter("total_batch_finished", labels...),
		totalBatchLatency:    factory.NewHistogram("total_batch_latency", meter.BatchBuckets, labels...),
		totalMessageStarted:  factory.NewCounter("total_message_started", labels...),
		totalMessageFinished: factory.NewCounter("total_message_finished", labels...),
	}
}
