package middlewares

import (
	"context"
	"time"

	"github.com/go-logr/logr"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"

	"github.com/tidbcloud/aisupport/packages/sdk-go/infra/logger"
	"github.com/tidbcloud/aisupport/packages/sdk-go/infra/metrics"
)

type RequestLoggerOptions struct {
	// Logger is used for logging RPC errors. It should not be nil.
	Logger logr.Logger
	// LogPayload logs RPC requests and responses at info level.
	LogPayload bool
	// FilterFunc can be set to filter out some requests from logging and metrics. Return true to skip logging and metrics for the request.
	FilterFunc func(ctx context.Context, req any, info *grpc.UnaryServerInfo) bool
}

type AccessLoggerOptions struct {
	// Logger is used for access log. If it's zero value, access log will be disabled.
	Logger logr.Logger
	// FilterFunc can be set to filter out some requests from access logging. Return true to skip access logging for the request.
	FilterFunc func(ctx context.Context, req any, info *grpc.UnaryServerInfo) bool
	// ExtraFieldsFunc can be set to add extra fields to access log.
	ExtraFieldsFunc func(ctx context.Context, req any, info *grpc.UnaryServerInfo) map[string]any
}

type LogMetricMiddlewareBuilder struct {
	Request RequestLoggerOptions
	Access  *AccessLoggerOptions
}

func (b *LogMetricMiddlewareBuilder) Build() grpc.ServerOption {
	return grpc.ChainUnaryInterceptor(b.unaryServerInterceptor())
}

func (b *LogMetricMiddlewareBuilder) unaryServerInterceptor() grpc.UnaryServerInterceptor {
	return func(
		ctx context.Context,
		req any,
		info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler,
	) (any, error) {
		start := time.Now()
		filtered := b.Request.FilterFunc != nil && b.Request.FilterFunc(ctx, req, info)

		if !filtered {
			log := logger.FromContextOrDef(ctx, b.Request.Logger)
			if b.Request.LogPayload {
				log.Info("request started",
					"method", info.FullMethod,
					"request", req,
				)
			} else {
				log.V(4).Info("request started",
					"method", info.FullMethod,
				)
			}
		}

		resp, err := handler(ctx, req)
		if filtered {
			return resp, err
		}

		code := status.Code(err)
		codeStr := code.String()
		latency := float64(time.Since(start)) / float64(time.Second)
		respSize := responseSize(resp)

		if metrics.RPC != nil {
			metrics.RPC.ReqDur.WithLabelValues(codeStr, info.FullMethod).Observe(latency)
			metrics.RPC.RespSize.Observe(float64(respSize))
		}

		if b.Access != nil && (b.Access.FilterFunc == nil || !b.Access.FilterFunc(ctx, req, info)) {
			fields := []any{
				"latency", latency,
				"status_code", codeStr,
				"body_size", respSize,
			}

			if b.Access.ExtraFieldsFunc != nil {
				for key, value := range b.Access.ExtraFieldsFunc(ctx, req, info) {
					fields = append(fields, key, value)
				}
			}

			logger.WithTraceID(ctx, b.Access.Logger).Info(info.FullMethod, fields...)
		}

		log := logger.FromContextOrDef(ctx, b.Request.Logger)
		if b.Request.LogPayload {
			log.Info("request completed",
				"latency", latency,
				"method", info.FullMethod,
				"status_code", codeStr,
				"body_size", respSize,
				"response", resp,
			)
		} else {
			log.V(4).Info("request completed",
				"latency", latency,
				"method", info.FullMethod,
				"status_code", codeStr,
				"body_size", respSize,
			)
		}

		if code != codes.OK {
			log.Error(err, info.FullMethod,
				"latency", latency,
				"status_code", codeStr,
			)
		}

		return resp, err
	}
}

func responseSize(resp any) int {
	if resp == nil {
		return 0
	}
	return proto.Size(resp.(proto.Message))
}