package rpcserver

import (
	"context"
	"encoding/json"
	"errors"
	"net"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
	"github.com/stretchr/testify/require"
	"google.golang.org/grpc"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"

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

type testGatewayHandler struct {
	name            string
	registerGateway func(context.Context, *runtime.ServeMux, *grpc.ClientConn) error
}

func (h testGatewayHandler) Name() string {
	return h.name
}

func (h testGatewayHandler) RegisterGateway(ctx context.Context, mux *runtime.ServeMux, conn *grpc.ClientConn) error {
	return h.registerGateway(ctx, mux, conn)
}

func TestNewRPCGatewayServerRegistersHandlers(t *testing.T) {
	t.Parallel()

	var registeredConnections []*grpc.ClientConn
	var registeredHandlers []string
	newHandler := func(name string) GatewayHandler {
		return testGatewayHandler{
			name: name,
			registerGateway: func(_ context.Context, _ *runtime.ServeMux, conn *grpc.ClientConn) error {
				registeredConnections = append(registeredConnections, conn)
				registeredHandlers = append(registeredHandlers, name)
				return nil
			},
		}
	}

	server, err := NewRPCGatewayServer(RPCGatewayServerParam{
		Address:    freeAddress(t),
		RPCAddress: freeAddress(t),
		Handlers: []GatewayHandler{
			newHandler("first"),
			newHandler("second"),
		},
	})
	require.NoError(t, err)
	t.Cleanup(func() {
		require.NoError(t, server.closeClient())
	})

	require.Equal(t, []string{"first", "second"}, registeredHandlers)
	require.Len(t, registeredConnections, 2)
	require.Same(t, server.conn, registeredConnections[0])
	require.Same(t, registeredConnections[0], registeredConnections[1])
}

func TestNewRPCGatewayServerReturnsGatewayRegistrationError(t *testing.T) {
	t.Parallel()

	expectedErr := errors.New("gateway registration failed")
	server, err := NewRPCGatewayServer(RPCGatewayServerParam{
		Address:    freeAddress(t),
		RPCAddress: freeAddress(t),
		Handlers: []GatewayHandler{
			testGatewayHandler{
				name: "failed-gateway-handler",
				registerGateway: func(context.Context, *runtime.ServeMux, *grpc.ClientConn) error {
					return expectedErr
				},
			},
		},
	})

	require.Nil(t, server)
	require.ErrorContains(t, err, "register gateway handler failed-gateway-handler failed")
	require.ErrorContains(t, err, expectedErr.Error())
}

func TestNewRPCGatewayServerAppliesMiddlewares(t *testing.T) {
	t.Parallel()

	server, err := NewRPCGatewayServer(RPCGatewayServerParam{
		RPCAddress: "127.0.0.1:1",
		Middlewares: []HTTPMiddleware{
			func(next http.Handler) http.Handler {
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.Header().Set("X-Test-Middleware", "first")
					next.ServeHTTP(w, r)
				})
			},
			func(next http.Handler) http.Handler {
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.Header().Add("X-Test-Middleware", "second")
					next.ServeHTTP(w, r)
				})
			},
		},
	})
	require.NoError(t, err)
	t.Cleanup(func() {
		require.NoError(t, server.closeClient())
	})

	response := httptest.NewRecorder()
	server.server.Handler.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/", nil))

	require.Equal(t, []string{"first", "second"}, response.Header().Values("X-Test-Middleware"))
}

func TestRPCGatewayServerRunAndShutdown(t *testing.T) {
	t.Parallel()

	requestLogger, requestLogs := newObservedRPCServerLogger()
	rpcAddress := freeAddress(t)
	gatewayAddress := freeAddress(t)
	rpcServer, err := NewRPCServer(RPCServerParam{
		Address: rpcAddress,
		ServerOptions: []grpc.ServerOption{
			(&middlewares.LogMetricMiddlewareBuilder{
				Request: middlewares.RequestLoggerOptions{Logger: requestLogger},
			}).Build(),
		},
	})
	require.NoError(t, err)

	rpcRunErrCh := make(chan error, 1)
	go func() {
		rpcRunErrCh <- rpcServer.Run()
	}()
	t.Cleanup(func() {
		shutdownCtx, cancelShutdown := context.WithTimeout(context.Background(), time.Second)
		defer cancelShutdown()
		require.NoError(t, rpcServer.Shutdown(shutdownCtx))
		require.NoError(t, <-rpcRunErrCh)
	})

	gatewayServer, err := NewRPCGatewayServer(RPCGatewayServerParam{
		Address:    gatewayAddress,
		RPCAddress: rpcAddress,
		Handlers: []GatewayHandler{
			testGatewayHandler{
				name:            "health-gateway",
				registerGateway: registerHealthGateway,
			},
		},
	})
	require.NoError(t, err)

	gatewayRunErrCh := make(chan error, 1)
	go func() {
		gatewayRunErrCh <- gatewayServer.Run()
	}()

	response := waitForGatewayResponse(t, "http://"+gatewayAddress+"/health")
	require.Equal(t, healthpb.HealthCheckResponse_SERVING, response.Status)
	require.Len(t, requestLogs.AllUntimed(), 2)
	require.Equal(t,
		[]string{"request started", "request completed"},
		[]string{requestLogs.AllUntimed()[0].Message, requestLogs.AllUntimed()[1].Message},
	)

	shutdownCtx, cancelShutdown := context.WithTimeout(context.Background(), time.Second)
	defer cancelShutdown()
	require.NoError(t, gatewayServer.Shutdown(shutdownCtx))
	require.NoError(t, <-gatewayRunErrCh)
	requireListenerClosed(t, gatewayAddress)
}

func registerHealthGateway(_ context.Context, mux *runtime.ServeMux, conn *grpc.ClientConn) error {
	return mux.HandlePath(http.MethodGet, "/health", func(w http.ResponseWriter, r *http.Request, _ map[string]string) {
		response := &healthpb.HealthCheckResponse{}
		err := conn.Invoke(
			r.Context(),
			"/grpc.health.v1.Health/Check",
			&healthpb.HealthCheckRequest{},
			response,
		)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		if err := json.NewEncoder(w).Encode(response); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}
	})
}

func waitForGatewayResponse(t *testing.T, url string) healthpb.HealthCheckResponse {
	t.Helper()

	var response healthpb.HealthCheckResponse
	require.Eventually(t, func() bool {
		httpResponse, err := http.Get(url)
		if err != nil {
			return false
		}
		defer httpResponse.Body.Close()

		if httpResponse.StatusCode != http.StatusOK {
			return false
		}
		return json.NewDecoder(httpResponse.Body).Decode(&response) == nil
	}, 5*time.Second, 10*time.Millisecond)
	return response
}

func requireListenerClosed(t *testing.T, address string) {
	t.Helper()

	conn, err := net.DialTimeout("tcp", address, 100*time.Millisecond)
	if err == nil {
		require.NoError(t, conn.Close())
	}
	require.Error(t, err)
}