package middlewares

import (
	"context"
	"errors"
	"testing"

	"github.com/go-logr/logr"
	"github.com/go-logr/zapr"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/testutil"
	dto "github.com/prometheus/client_model/go"
	"github.com/stretchr/testify/require"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest/observer"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"

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

const testFullMethod = "/test.Service/Check"

func newObservedMiddlewareLogger() (logr.Logger, *observer.ObservedLogs) {
	core, logs := observer.New(zapcore.Level(-10))
	return zapr.NewLogger(zap.New(core)), logs
}

func testUnaryInfo() *grpc.UnaryServerInfo {
	return &grpc.UnaryServerInfo{FullMethod: testFullMethod}
}

func testTraceContext() context.Context {
	spanCtx := trace.NewSpanContext(trace.SpanContextConfig{
		TraceID:    trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16},
		TraceFlags: trace.FlagsSampled,
	})
	return trace.ContextWithSpanContext(context.Background(), spanCtx)
}

func TestLogMetricMiddlewareLogsSuccessfulUnaryRPC(t *testing.T) {
	requestLogger, requestLogs := newObservedMiddlewareLogger()
	accessLogger, accessLogs := newObservedMiddlewareLogger()
	response := &healthpb.HealthCheckResponse{
		Status: healthpb.HealthCheckResponse_SERVING,
	}
	request := &healthpb.HealthCheckRequest{Service: "test-service"}
	builder := &LogMetricMiddlewareBuilder{
		Request: RequestLoggerOptions{
			Logger:     requestLogger,
			LogPayload: true,
		},
		Access: &AccessLoggerOptions{
			Logger: accessLogger,
			ExtraFieldsFunc: func(context.Context, any, *grpc.UnaryServerInfo) map[string]any {
				return map[string]any{"tenant_id": "tenant-1"}
			},
		},
	}

	got, err := builder.unaryServerInterceptor()(
		testTraceContext(),
		request,
		testUnaryInfo(),
		func(context.Context, any) (any, error) {
			return response, nil
		},
	)

	require.NoError(t, err)
	require.Same(t, response, got)

	requestEntries := requestLogs.AllUntimed()
	require.Len(t, requestEntries, 2)
	require.Equal(t, "request started", requestEntries[0].Message)
	require.Equal(t, zapcore.InfoLevel, requestEntries[0].Level)
	require.Equal(t, testFullMethod, requestEntries[0].ContextMap()["method"])
	require.Equal(t, request.String(), requestEntries[0].ContextMap()["request"])
	require.Equal(t, "request completed", requestEntries[1].Message)

	completedFields := requestEntries[1].ContextMap()
	require.Equal(t, zapcore.InfoLevel, requestEntries[1].Level)
	require.Equal(t, testFullMethod, completedFields["method"])
	require.Equal(t, codes.OK.String(), completedFields["status_code"])
	require.EqualValues(t, proto.Size(response), completedFields["body_size"])
	require.Equal(t, response.String(), completedFields["response"])
	require.Equal(t, "0102030405060708090a0b0c0d0e0f10", completedFields["trace_id"])

	accessEntries := accessLogs.AllUntimed()
	require.Len(t, accessEntries, 1)
	require.Equal(t, testFullMethod, accessEntries[0].Message)

	accessFields := accessEntries[0].ContextMap()
	require.Equal(t, codes.OK.String(), accessFields["status_code"])
	require.EqualValues(t, proto.Size(response), accessFields["body_size"])
	require.Equal(t, "tenant-1", accessFields["tenant_id"])
	require.Equal(t, "0102030405060708090a0b0c0d0e0f10", accessFields["trace_id"])
}

func TestLogMetricMiddlewareLogsUnaryRPCError(t *testing.T) {
	requestLogger, logs := newObservedMiddlewareLogger()
	expectedErr := status.Error(codes.Internal, "failed")
	request := &healthpb.HealthCheckRequest{Service: "test-service"}
	builder := &LogMetricMiddlewareBuilder{
		Request: RequestLoggerOptions{
			Logger:     requestLogger,
			LogPayload: true,
		},
	}

	response, err := builder.unaryServerInterceptor()(
		context.Background(),
		request,
		testUnaryInfo(),
		func(context.Context, any) (any, error) {
			return nil, expectedErr
		},
	)

	require.Nil(t, response)
	require.ErrorIs(t, err, expectedErr)

	entries := logs.AllUntimed()
	require.Len(t, entries, 3)
	require.Equal(t, "request started", entries[0].Message)
	require.Equal(t, "request completed", entries[1].Message)
	require.Equal(t, testFullMethod, entries[2].Message)
	require.Equal(t, zapcore.InfoLevel, entries[0].Level)
	require.Equal(t, zapcore.InfoLevel, entries[1].Level)
	require.Equal(t, zapcore.ErrorLevel, entries[2].Level)
	require.Equal(t, request.String(), entries[0].ContextMap()["request"])
	require.Contains(t, entries[1].ContextMap(), "response")
	require.Nil(t, entries[1].ContextMap()["response"])
	require.Equal(t, codes.Internal.String(), entries[1].ContextMap()["status_code"])
	require.EqualValues(t, 0, entries[1].ContextMap()["body_size"])
	require.Equal(t, codes.Internal.String(), entries[2].ContextMap()["status_code"])
}

func TestLogMetricMiddlewarePayloadLoggingDisabledByDefault(t *testing.T) {
	requestLogger, logs := newObservedMiddlewareLogger()
	builder := &LogMetricMiddlewareBuilder{
		Request: RequestLoggerOptions{Logger: requestLogger},
	}

	_, err := builder.unaryServerInterceptor()(
		context.Background(),
		&healthpb.HealthCheckRequest{},
		testUnaryInfo(),
		func(context.Context, any) (any, error) {
			return &healthpb.HealthCheckResponse{}, nil
		},
	)

	require.NoError(t, err)
	entries := logs.AllUntimed()
	require.Len(t, entries, 2)
	require.Equal(t, int8(-4), int8(entries[0].Level))
	require.Equal(t, int8(-4), int8(entries[1].Level))
	require.NotContains(t, entries[0].ContextMap(), "request")
	require.NotContains(t, entries[1].ContextMap(), "response")
}

func TestLogMetricMiddlewareAccessFilter(t *testing.T) {
	requestLogger, requestLogs := newObservedMiddlewareLogger()
	accessLogger, accessLogs := newObservedMiddlewareLogger()
	builder := &LogMetricMiddlewareBuilder{
		Request: RequestLoggerOptions{Logger: requestLogger},
		Access: &AccessLoggerOptions{
			Logger: accessLogger,
			FilterFunc: func(context.Context, any, *grpc.UnaryServerInfo) bool {
				return true
			},
		},
	}

	_, err := builder.unaryServerInterceptor()(
		context.Background(),
		&healthpb.HealthCheckRequest{},
		testUnaryInfo(),
		func(context.Context, any) (any, error) {
			return &healthpb.HealthCheckResponse{}, nil
		},
	)

	require.NoError(t, err)
	require.Len(t, requestLogs.AllUntimed(), 2)
	require.Empty(t, accessLogs.AllUntimed())
}

func TestLogMetricMiddlewareRecordsMetrics(t *testing.T) {
	registerTestRPCMetrics(t, "rpc_middleware_success_test")

	response := &healthpb.HealthCheckResponse{
		Status: healthpb.HealthCheckResponse_SERVING,
	}
	builder := &LogMetricMiddlewareBuilder{}

	_, err := builder.unaryServerInterceptor()(
		context.Background(),
		&healthpb.HealthCheckRequest{},
		testUnaryInfo(),
		func(context.Context, any) (any, error) {
			return response, nil
		},
	)

	require.NoError(t, err)
	require.EqualValues(t, 1, histogramSampleCount(t,
		metrics.RPC.ReqDur.WithLabelValues(codes.OK.String(), testFullMethod),
	))
	count, sum := summaryValue(t, metrics.RPC.RespSize)
	require.EqualValues(t, 1, count)
	require.EqualValues(t, proto.Size(response), sum)
}

func TestLogMetricMiddlewareRequestFilterSkipsLogsAndMetrics(t *testing.T) {
	registerTestRPCMetrics(t, "rpc_middleware_filter_test")

	requestLogger, logs := newObservedMiddlewareLogger()
	handlerCalled := false
	builder := &LogMetricMiddlewareBuilder{
		Request: RequestLoggerOptions{
			Logger: requestLogger,
			FilterFunc: func(context.Context, any, *grpc.UnaryServerInfo) bool {
				return true
			},
		},
	}

	_, err := builder.unaryServerInterceptor()(
		context.Background(),
		&healthpb.HealthCheckRequest{},
		testUnaryInfo(),
		func(context.Context, any) (any, error) {
			handlerCalled = true
			return nil, errors.New("failed")
		},
	)

	require.Error(t, err)
	require.True(t, handlerCalled)
	require.Empty(t, logs.AllUntimed())
	require.Equal(t, 0, testutil.CollectAndCount(metrics.RPC.ReqDur))
	count, sum := summaryValue(t, metrics.RPC.RespSize)
	require.Zero(t, count)
	require.Zero(t, sum)
}

func registerTestRPCMetrics(t *testing.T, namespace string) {
	t.Helper()

	metrics.RegisterRPCMetrics(namespace)
	t.Cleanup(func() {
		prometheus.Unregister(metrics.RPC.ReqDur)
		prometheus.Unregister(metrics.RPC.RespSize)
		metrics.RPC = nil
	})
}

func histogramSampleCount(t *testing.T, observer prometheus.Observer) uint64 {
	t.Helper()

	metric := observer.(prometheus.Metric)
	var value dto.Metric
	require.NoError(t, metric.Write(&value))
	return value.GetHistogram().GetSampleCount()
}

func summaryValue(t *testing.T, metric prometheus.Metric) (uint64, float64) {
	t.Helper()

	var value dto.Metric
	require.NoError(t, metric.Write(&value))
	return value.GetSummary().GetSampleCount(), value.GetSummary().GetSampleSum()
}