package grpc

import (
	"context"
	"net"
	"strconv"

	"github.com/gogo/grpc-example/insecure"
	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
	grpc_zap "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap"
	grpc_recovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
	grpc_validator "github.com/grpc-ecosystem/go-grpc-middleware/validator"
	"github.com/pingcap/log"
	"github.com/pingcap/tidb/pkg/lightning/mydump"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/agent/constants"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/agent/tidbcloud"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	grpc_errcode "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/middleware/errcode"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/crypto"
	grpchttp "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/grpc/http"
	"github.com/tidbcloud/dataflow-service/dataflow-agent-ng/pkg/loader"
	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

type AgentServer struct {
	agent.UnimplementedAgentServiceServer
	// While enabling clusterTLS, the agent will use the local k8s cluster's tls context:
	// /cluster-cert/tls.crt, /cluster-cert/tls.key, /cluster-cert/ca.crt
	enableLocalClusterCert bool
	clusterProvider        tidbcloud.ClusterProvider
	encryptor              *crypto.RSAEncryptor
	importTaskManager      *importAsyncTaskManager
	// scanDataFiles is the function used to scan files from external storage.
	// It can be replaced in tests to inject mock behavior.
	scanDataFiles func(ctx context.Context, url string, maxScanFiles int) ([]*mydump.SourceFileMeta, int64, error)
	// createSchemasFn allows tests to stub schema creation behavior.
	createSchemasFn func(ctx context.Context, req *agent.CreateSchemasReq) (*agent.CreateSchemasResp, error)
	// execSQLOnClusterFn allows tests to stub cluster SQL execution behavior.
	execSQLOnClusterFn func(ctx context.Context, req *agent.ExecSQLReq) (*agent.ExecSQLResp, error)
}

func MustNewAgentGrpcServer(enableLocalClusterCert bool, clusterModel string) *grpc.Server {
	unaryOpts := grpc.UnaryInterceptor(
		grpc_middleware.ChainUnaryServer(
			otelgrpc.UnaryServerInterceptor(),
			grpc_zap.UnaryServerInterceptor(log.L().Named("grpc")),
			grpc_errcode.UnaryServerInterceptor(),
			grpchttp.LogUnaryServerInterceptor(),
			grpc_validator.UnaryServerInterceptor(),
			grpc_recovery.UnaryServerInterceptor(),
		),
	)
	streamOpts := grpc.StreamInterceptor(
		grpc_middleware.ChainStreamServer(
			otelgrpc.StreamServerInterceptor(),
			grpc_errcode.StreamServerInterceptor(),
			grpc_validator.StreamServerInterceptor(),
			grpc_recovery.StreamServerInterceptor(),
		),
	)
	s := grpc.NewServer(grpc.Creds(credentials.NewServerTLSFromCert(&insecure.Cert)), unaryOpts, streamOpts)
	encryptor, err := crypto.NewRSAEncryptor(2048)
	if err != nil {
		panic(err)
	}
	agentServer := &AgentServer{
		enableLocalClusterCert: enableLocalClusterCert,
		encryptor:              encryptor,
		importTaskManager:      newImportAsyncTaskManager(),
		scanDataFiles:          loader.ScanDataFiles,
	}
	kubeconfig, err := rest.InClusterConfig()
	if err != nil {
		panic(err)
	}
	kubeCli, err := kubernetes.NewForConfig(kubeconfig)
	if err != nil {
		panic(err)
	}
	switch clusterModel {
	case "dedicated":
		agentServer.clusterProvider = tidbcloud.NewDedicatedClusterProvider(kubeCli)
	case "next-gen":
		agentServer.clusterProvider = tidbcloud.NewNextGenClusterProvider(kubeCli)
	default:
		agentServer.clusterProvider = tidbcloud.NewDedicatedClusterProvider(kubeCli)
	}
	agent.RegisterAgentServiceServer(s, agentServer)
	return s
}

func StartGRPCServer(enableLocalClusterCert bool, clusterModel string) error {
	s := MustNewAgentGrpcServer(enableLocalClusterCert, clusterModel)
	addr := net.JoinHostPort(constants.ListenHost, strconv.Itoa(constants.GrpcPort))
	lis, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	log.L().Info("starting agent gRPC server", zap.String("addr", addr))

	// Start serving with the tigrpc server
	serveErr := s.Serve(lis)
	if serveErr != nil {
		log.Fatal("failed to serve gRPC server", zap.Error(serveErr))
	}
	log.L().Info("agent gRPC server exit now")
	return nil
}