package server

import (
	"context"
	"fmt"
	"net/http"
	"strconv"

	grpcAuth "github.com/grpc-ecosystem/go-grpc-middleware/auth"
	grpcZap "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap"
	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
	grpcRecovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
	grpcValidator "github.com/grpc-ecosystem/go-grpc-middleware/validator"
	grpcPrometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
	"github.com/pingcap/log"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"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/status"
	"google.golang.org/protobuf/proto"
)

const (
	httpStatusCodeKey     = "x-http-code"
	grpcHTTPStatusCodeKey = "Grpc-Metadata-X-Http-Code"
)

func WithPostServerOption(o grpc.ServerOption) Option {
	return func(s *serverOption) {
		s.postGrpcOpts = append(s.postGrpcOpts, o)
	}
}

func WithPreServerOption(o grpc.ServerOption) Option {
	return func(s *serverOption) {
		s.preGrpcOpts = append(s.preGrpcOpts, o)
	}
}

func WithAuthFunc(authFunc grpcAuth.AuthFunc) Option {
	return func(s *serverOption) {
		s.unaryInterceptors = append(s.unaryInterceptors, grpcAuth.UnaryServerInterceptor(authFunc))
		s.streamInterceptors = append(s.streamInterceptors, grpcAuth.StreamServerInterceptor(authFunc))
	}
}

func WithValidator() Option {
	return func(s *serverOption) {
		s.unaryInterceptors = append(s.unaryInterceptors, grpcValidator.UnaryServerInterceptor())
		s.streamInterceptors = append(s.streamInterceptors, grpcValidator.StreamServerInterceptor())
	}
}

func WithPrometheus() Option {
	return func(s *serverOption) {
		grpcPrometheus.EnableHandlingTimeHistogram()
		s.unaryInterceptors = append(s.unaryInterceptors, grpcPrometheus.UnaryServerInterceptor)
		s.streamInterceptors = append(s.streamInterceptors, grpcPrometheus.StreamServerInterceptor)
	}
}

func WithTracing(opts ...otelgrpc.Option) Option {
	return func(s *serverOption) {
		s.unaryInterceptors = append(s.unaryInterceptors, otelgrpc.UnaryServerInterceptor(opts...))
		s.streamInterceptors = append(s.streamInterceptors, otelgrpc.StreamServerInterceptor(opts...))
	}
}

func WithZapLogger(logger *zap.Logger, opts ...grpcZap.Option) Option {
	return func(s *serverOption) {
		s.unaryInterceptors = append(s.unaryInterceptors, grpcZap.UnaryServerInterceptor(logger, opts...))
		s.streamInterceptors = append(s.streamInterceptors, grpcZap.StreamServerInterceptor(logger, opts...))
	}
}

func WithRecovery(opts ...grpcRecovery.Option) Option {
	return func(s *serverOption) {
		s.unaryInterceptors = append(s.unaryInterceptors, grpcRecovery.UnaryServerInterceptor(opts...))
		s.streamInterceptors = append(s.streamInterceptors, grpcRecovery.StreamServerInterceptor(opts...))
	}
}

// `WithUnaryInterceptors` can add your custom unary interceptors to the grpc server before it's initialized
func WithUnaryInterceptors(interceptors ...grpc.UnaryServerInterceptor) Option {
	return func(s *serverOption) {
		s.unaryInterceptors = append(s.unaryInterceptors, interceptors...)
	}
}

// `WithStreamInterceptors` can add your custom stream interceptors to the grpc server before it's initialized
func WithStreamInterceptors(interceptors ...grpc.StreamServerInterceptor) Option {
	return func(s *serverOption) {
		s.streamInterceptors = append(s.streamInterceptors, interceptors...)
	}
}

// `WithDefault` add tigrpc default interceptors to the server, including metrics, tracing, server log, and panic recovery.
// The default recovery behavior is log the panic stack and return `status.Internal` and the panic message to client
// The default log behavior is log the trace id
func WithDefault(opts ...InterceptorOption) Option {
	return func(s *serverOption) {
		ios := &interceptorOptions{
			recoveryOpts: []grpcRecovery.Option{
				grpcRecovery.WithRecoveryHandlerContext(DefaultPanicHandler),
			},
			zapLogger: log.L(),
			zapOpts: []grpcZap.Option{
				grpcZap.WithMessageProducer(TraceIDMessageProducer),
			},
		}
		for _, opt := range opts {
			opt(ios)
		}

		options := []Option{
			WithPrometheus(),
			WithTracing(),
			WithZapLogger(ios.zapLogger, ios.zapOpts...),
			WithRecovery(ios.recoveryOpts...),
		}
		for _, option := range options {
			option(s)
		}
	}
}

type interceptorOptions struct {
	recoveryOpts []grpcRecovery.Option
	zapOpts      []grpcZap.Option
	zapLogger    *zap.Logger
}

type InterceptorOption func(*interceptorOptions)

// FOR DEVELOPER: when you extend interceptorOptions, make sure the function's prefix is `Customize`
// it's intended to distinguish from `WithXXX`
func CustomizeRecoveryOption(recoveryOpt ...grpcRecovery.Option) InterceptorOption {
	return func(c *interceptorOptions) {
		c.recoveryOpts = recoveryOpt
	}
}

func CustomizeZapOption(zapOpt ...grpcZap.Option) InterceptorOption {
	return func(c *interceptorOptions) {
		c.zapOpts = zapOpt
	}
}

func CustomizeZapLogger(logger *zap.Logger) InterceptorOption {
	return func(c *interceptorOptions) {
		c.zapLogger = logger
	}
}

func TraceIDMessageProducer(ctx context.Context, msg string, level zapcore.Level, code codes.Code, err error, duration zapcore.Field) {
	// re-extract logger from newCtx, as it may have extra fields that changed in the holder.
	spanCtx := trace.SpanContextFromContext(ctx)
	if spanCtx.HasTraceID() {
		ctxzap.Extract(ctx).Check(level, msg).Write(
			zap.Error(err),
			zap.String("grpc.code", code.String()),
			duration,
			zap.String("trace_id", spanCtx.TraceID().String()),
		)
	} else {
		ctxzap.Extract(ctx).Check(level, msg).Write(
			zap.Error(err),
			zap.String("grpc.code", code.String()),
			duration,
		)
	}
}

// GRPCGatewayHTTPResponseModifier use this handler for grpc gateway, it can work with tierr, set HTTP status code as what tierr indicates.
func GRPCGatewayHTTPResponseModifier(ctx context.Context, w http.ResponseWriter, p proto.Message) error {
	md, ok := runtime.ServerMetadataFromContext(ctx)
	if !ok {
		return nil
	}

	// set http status code
	if vals := md.HeaderMD.Get(httpStatusCodeKey); 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, httpStatusCodeKey)
		delete(w.Header(), grpcHTTPStatusCodeKey)
		w.WriteHeader(code)
	}

	return nil
}

func DefaultPanicHandler(ctx context.Context, p interface{}) error {
	tilog.Error(ctx, "panic triggered", zap.String("panic", fmt.Sprintf("%v", p)))
	return status.Errorf(codes.Internal, "%v", p)
}