package main

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/textproto"
	"strconv"
	"time"

	"buf.build/go/protovalidate"
	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/gin-gonic/gin"
	grpclogging "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
	protovalidate_middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/protovalidate"
	grpcrecovery "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/recovery"
	gwruntime "github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
	certmanagerv1 "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1"
	asappsv1 "github.com/pingcap/advanced-statefulset/client/apis/apps/v1"
	"github.com/pingcap/log"
	opv2br "github.com/pingcap/tidb-operator/api/v2/br/v1alpha1"
	opv2core "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	operatorv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	tigrpcc "github.com/tidbcloud/tiinf/transport/tigrpc/client"
	"go.opentelemetry.io/contrib/instrumentation/github.com/gin-gonic/gin/otelgin"
	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/health"
	"google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
	"k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/infra-provider/cluster-agent/proxy"

	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-agent/v1"
	"github.com/tidbcloud/infra-provider/cluster-agent/service"
	"github.com/tidbcloud/infra-provider/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/pkg/cloud/azure"
	"github.com/tidbcloud/infra-provider/pkg/cloud/gcp"
	"github.com/tidbcloud/infra-provider/pkg/consts"
)

const (
	baseAPIPath = "/api/v1"
	// baseAPIPathWithPrefix when cluster-agent is deployed behind Cluster Gateway, a prefix like '/cluster-agent/tidb123'
	// is added.
	baseAPIPathWithPrefix = "/cluster-agent/:tcns" + baseAPIPath
	healthzEndpoint       = "/healthz"
)

func setupAndRunGRPCServer(ctx context.Context, cfg *Config) {
	var (
		provider = cfg.Provider
		addr     = cfg.InternalGRPCAddr
	)
	validator, err := protovalidate.New()
	if err != nil {
		tilog.Fatal(ctx, "Failed to create proto validator", zap.Error(err))
	}

	lis, err := net.Listen("tcp", addr)
	if err != nil {
		tilog.Fatal(ctx, "Failed to listen for gRPC server", zap.String("addr", addr), zap.Error(err))
	}
	s := grpc.NewServer(
		grpc.StatsHandler(otelgrpc.NewServerHandler()),
		grpc.ChainUnaryInterceptor(
			grpcrecovery.UnaryServerInterceptor(grpcrecovery.WithRecoveryHandlerContext(recoveryHandler)),
			grpclogging.UnaryServerInterceptor(InterceptorLogger(log.L())),
			reqAndRespInterceptor,
			protovalidate_middleware.UnaryServerInterceptor(validator),
		),
	)

	// create a k8s client with the in-cluster config
	dataplaneCli := buildDataplaneCli(ctx, cfg)

	// create agent service with the specified provider's client
	var opts []service.Option
	var nodeRestarter service.NodeRestartProvider
	switch provider {
	case "aws":
		cfg, err := awsconfig.LoadDefaultConfig(ctx)
		if err != nil {
			tilog.Fatal(ctx, "failed to load aws config: %v", zap.Error(err))
		}
		nodeRestarter = service.NewAWSNodeRestarter(aws.NewClientSet(cfg))
	case consts.ProviderTypeGCP, consts.ProviderTypeGKE:
		gcpCli, err := gcp.NewClientSet()
		if err != nil {
			tilog.Fatal(ctx, "failed to create gcp client: %v", zap.Error(err))
		}
		nodeRestarter = service.NewGCPNodeRestarter(gcpCli)
	case "kind":
		nodeRestarter = service.NewKindNodeRestarter()
	case "azure":
		cred, err := azidentity.NewDefaultAzureCredential(nil)
		if err != nil {
			tilog.Fatal(ctx, "failed to get azure default credential: %v", zap.Error(err))
		}
		azureCli, err := azure.NewClientSet(cred)
		if err != nil {
			tilog.Fatal(ctx, "failed to create azure client: %v", zap.Error(err))
		}
		nodeRestarter = service.NewAzureNodeRestarter(azureCli)
	default:
		tilog.Fatal(ctx, "unknown provider", zap.String("provider", provider))
	}
	opts = append(opts,
		service.WithProvider(provider),
		service.WithNodeRestarter(nodeRestarter),
		service.WithTiDBOperatorVersion(cfg.TiDBOperatorVersion),
		service.WithClusterNamespace(cfg.ClusterNamespace),
		service.WithGKEPolicyControlTarget(
			cfg.GKEPolicyControlProjectID,
			cfg.GKEPolicyControlLocation,
			cfg.GKEPolicyControlClusterName,
			cfg.GKEPolicyControlFleetProjectID,
		),
	)
	agentSvc := service.NewAgentService(dataplaneCli, opts...)
	certSvc := service.NewCertService(dataplaneCli)
	cdcSvc := service.NewCdcAgentService(dataplaneCli)
	operatorV2Svc := service.NewOperatorV2Service(dataplaneCli)

	apiv1.RegisterClusterAgentServiceServer(s, agentSvc)
	apiv1.RegisterCertServiceServer(s, certSvc)
	apiv1.RegisterCdcAgentServiceServer(s, cdcSvc)
	apiv1.RegisterOperatorV2ServiceServer(s, operatorV2Svc)
	// refer to https://grpc.io/docs/guides/health-checking/
	grpc_health_v1.RegisterHealthServer(s, health.NewServer())
	tilog.Info(ctx, "Starting internal gRPC server", zap.String("addr", addr))
	go func() {
		if err2 := s.Serve(lis); err2 != nil {
			tilog.Info(ctx, "Internal gRPC server returned with error", zap.String("addr", addr), zap.Error(err2))
		}
	}()

	<-ctx.Done()

	tilog.Info(ctx, "Stopping internal gRPC server...", zap.String("addr", addr))
	s.GracefulStop()
	tilog.Info(ctx, "Internal gRPC server stopped", zap.String("addr", addr))
}

func setupAndRunGRPCGateway(ctx context.Context, addr, grpcAddr string) {
	conn, err := tigrpcc.NewClientConn(grpcAddr, tigrpcc.WithDefault())
	if err != nil {
		tilog.Fatal(ctx, "Failed to dial gRPC server", zap.String("addr", grpcAddr), zap.Error(err))
	}
	defer func() {
		go func() {
			<-ctx.Done()
			if cerr := conn.Close(); cerr != nil {
				tilog.Info(ctx, "Failed to close conn to %s: %v", zap.String("addr", grpcAddr), zap.Error(cerr))
			}
		}()
	}()

	mux := gwruntime.NewServeMux(
		gwruntime.WithForwardResponseOption(httpResponseModifier),
		// refer to https://grpc-ecosystem.github.io/grpc-gateway/docs/operations/health_check/
		gwruntime.WithHealthEndpointAt(grpc_health_v1.NewHealthClient(conn), healthzEndpoint),
	)
	if err := apiv1.RegisterClusterAgentServiceHandler(ctx, mux, conn); err != nil {
		tilog.Fatal(ctx, "Failed to register gRPC gateway for cluster agent", zap.Error(err))
	}
	if err := apiv1.RegisterCertServiceHandler(ctx, mux, conn); err != nil {
		tilog.Fatal(ctx, "Failed to register gRPC gateway for cert service", zap.Error(err))
	}
	if err := apiv1.RegisterCdcAgentServiceHandler(ctx, mux, conn); err != nil {
		tilog.Fatal(ctx, "Failed to register gRPC gateway for cdc agent", zap.Error(err))
	}
	if err := apiv1.RegisterOperatorV2ServiceHandler(ctx, mux, conn); err != nil {
		tilog.Fatal(ctx, "Failed to register gRPC gateway for operator v2 services", zap.Error(err))
	}

	apiRouter := gin.New()
	// log with custom format; strip the base path (/api/v1) from the request path
	apiRouter.Use(
		gin.CustomRecovery(func(c *gin.Context, err any) {
			tilog.Error(c.Request.Context(), "panic triggered", zap.String("panic", fmt.Sprintf("%v", err)))
		}),
		otelgin.Middleware(serviceName),
		//middlewares.CommonLogging,
		gin.Logger(),
		ServePathStartFromAny,
		//tigin.RecordMetrics,
		SetResponseTraceID,
	)
	apiRouter.Any(fmt.Sprintf("%s/*any", baseAPIPath), gin.WrapH(mux))
	apiRouter.Any(fmt.Sprintf("%s/*any", baseAPIPathWithPrefix), gin.WrapH(mux))
	apiRouter.Any(healthzEndpoint, gin.WrapH(mux))
	apiRouter.Any(fmt.Sprintf("%s/*any", proxy.ProxyBasePath), proxy.ProxyEntry)
	srv := &http.Server{
		Addr:    addr,
		Handler: apiRouter,
	}

	lis, err := net.Listen("tcp", addr)
	if err != nil {
		tilog.Fatal(ctx, "Failed to listen for gRPC gateway", zap.String("addr", addr), zap.Error(err))
	}
	tilog.Info(ctx, "Starting gRPC gateway", zap.String("addr", addr))
	go func() {
		if err2 := srv.Serve(lis); err2 != nil && !errors.Is(err2, http.ErrServerClosed) {
			tilog.Info(ctx, "The gRPC gateway returned with error", zap.String("addr", addr), zap.Error(err2))
		}
	}()

	<-ctx.Done()

	// The context is used to inform the server it has 5 seconds to finish
	// the request it is currently handling
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	tilog.Info(ctx, "Stopping gRPC gatway...", zap.String("addr", addr))
	if err := srv.Shutdown(ctx); err != nil {
		tilog.Error(ctx, "gRPC gateway forced to shutdown", zap.String("addr", addr), zap.Error(err))
	} else {
		tilog.Info(ctx, "gRPC gateway stopped", zap.String("addr", addr))
	}
}

func httpResponseModifier(ctx context.Context, w http.ResponseWriter, p proto.Message) error {
	md, ok := gwruntime.ServerMetadataFromContext(ctx)
	if !ok {
		return nil
	}

	// set http status code
	if vals := md.HeaderMD.Get(ResponseKeyStatusCode); len(vals) > 0 {
		code, err := strconv.Atoi(vals[0])
		if err != nil {
			return err
		}
		// delete the headers to not expose any grpc-metadata in http response
		delete(md.HeaderMD, ResponseKeyStatusCode)
		delete(w.Header(), fmt.Sprintf("Grpc-Metadata-%s", textproto.CanonicalMIMEHeaderKey(ResponseKeyStatusCode)))
		w.WriteHeader(code)
	}

	return nil
}

func recoveryHandler(ctx context.Context, p any) (err error) {
	tilog.Error(ctx, "panic triggered", zap.Any("panic", p))
	return status.Errorf(codes.Unknown, "panic triggered: %v", p)
}

// InterceptorLogger adapts zap logger to interceptor logger.
// This code is simple enough to be copied and not imported.
func InterceptorLogger(logger *zap.Logger) grpclogging.Logger {
	return grpclogging.LoggerFunc(func(ctx context.Context, lvl grpclogging.Level, msg string, fields ...any) {
		f := make([]zap.Field, 0, len(fields)/2)

		for i := 0; i < len(fields); i += 2 {
			key := fields[i]
			value := fields[i+1]

			switch v := value.(type) {
			case string:
				f = append(f, zap.String(key.(string), v))
			case int:
				f = append(f, zap.Int(key.(string), v))
			case bool:
				f = append(f, zap.Bool(key.(string), v))
			default:
				f = append(f, zap.Any(key.(string), v))
			}
		}

		logger := logger.WithOptions(zap.AddCallerSkip(1)).With(f...).With(tilog.FieldsFromContext(ctx)...)

		levelMap := map[grpclogging.Level]zapcore.Level{
			grpclogging.LevelDebug: zap.DebugLevel,
			grpclogging.LevelInfo:  zap.InfoLevel,
			grpclogging.LevelWarn:  zap.WarnLevel,
			grpclogging.LevelError: zap.ErrorLevel,
		}
		zapLevel, ok := levelMap[lvl]
		if !ok {
			zapLevel = zap.InfoLevel
		}
		logger.Log(zapLevel, msg)
	})
}

func reqAndRespInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
	logger := tilog.WithContextL(ctx)
	logger.Debug("request", zap.Any("req", req))
	resp, err = handler(ctx, req)
	logger.Debug("response", zap.Any("resp", resp), zap.Error(err))
	return
}

// SetResponseTraceID is copied from tiinf/transport/tigin/middlewares.go to avoid import too many useless packages
func SetResponseTraceID(c *gin.Context) {
	spanCtx := trace.SpanContextFromContext(c.Request.Context())
	if spanCtx.HasTraceID() {
		c.Header("x-debug-trace-id", spanCtx.TraceID().String())
	}
	c.Next()
}

func buildDataplaneCli(ctx context.Context, cfg *Config) client.Client {
	var config *rest.Config
	var err error
	if cfg.DataplaneKubeConfig != "" {
		tilog.Info(ctx, "Using dataplane kubeconfig file", zap.String("file", cfg.DataplaneKubeConfig))
		config, err = clientcmd.BuildConfigFromFlags("", cfg.DataplaneKubeConfig)
		if err != nil {
			tilog.Fatal(ctx, "failed to build dataplane config: %v", zap.Error(err))
		}
	} else {
		// create a k8s client with the in-cluster config
		tilog.Info(ctx, "Using in-cluster config")
		config, err = rest.InClusterConfig()
		if err != nil {
			tilog.Fatal(ctx, "failed to get in-cluster config: %v", zap.Error(err))
		}
	}

	scheme := runtime.NewScheme()
	utilruntime.Must(clientgoscheme.AddToScheme(scheme))
	utilruntime.Must(operatorv1alpha1.SchemeBuilder.AddToScheme(scheme))
	utilruntime.Must(asappsv1.AddToScheme(scheme))
	utilruntime.Must(certmanagerv1.AddToScheme(scheme))
	utilruntime.Must(opv2core.AddToScheme(scheme))
	utilruntime.Must(opv2br.AddToScheme(scheme))
	dataplaneCli, err := client.New(config, client.Options{Scheme: scheme})
	if err != nil {
		tilog.Fatal(ctx, "failed to create dataplane client: %v", zap.Error(err))
	}

	return dataplaneCli
}
