package agentsidecar

import (
	"context"
	"net"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"google.golang.org/grpc"
	"google.golang.org/grpc/test/bufconn"

	apiproxyapi "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/agent-sidecar/apiproxy/v1"
	projectapi "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/agent-sidecar/project/v1"
	secretstoreapi "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/api/agent-sidecar/secretstore/v1"
	apiproxyentity "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/entity/agent-sidecar/apiproxy/v1"
	projectentity "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/entity/agent-sidecar/project/v1"
	secretstoreentity "github.com/tidbcloud/aisupport/packages/sdk-go/contracts/entity/agent-sidecar/secretstore/v1"
	"github.com/tidbcloud/aisupport/services/agent-gateway/config"
)

func TestRegistryRoutesAgentsToIndependentSidecarsAndClosesConnections(t *testing.T) {
	listenerA := startFakeSidecar(t, "from-a")
	listenerB := startFakeSidecar(t, "from-b")
	listeners := map[string]*bufconn.Listener{"agent-a": listenerA, "agent-b": listenerB}
	originalNewClient := newClientConnection
	newClientConnection = func(target string, options ...grpc.DialOption) (*grpc.ClientConn, error) {
		listener := listeners[target]
		options = append(options, grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
			return listener.Dial()
		}))
		return grpc.NewClient("passthrough:///"+target, options...)
	}
	t.Cleanup(func() { newClientConnection = originalNewClient })
	conf := &config.RootConfig{AgentRuntime: config.AgentRuntimeConfig{Agents: map[string]config.AgentProviderConfig{
		"agent-b":   {Sidecar: &config.AgentSidecarConfig{Endpoint: "agent-b", Timeout: 2 * time.Second}},
		"agent-a":   {Sidecar: &config.AgentSidecarConfig{Endpoint: "agent-a", Timeout: time.Second}},
		"unmanaged": {},
	}}}

	registry, err := NewRegistry(conf)
	require.NoError(t, err)

	clientA, ok := registry.Get("agent-a")
	require.True(t, ok)
	clientB, ok := registry.Get("agent-b")
	require.True(t, ok)
	_, ok = registry.Get("unmanaged")
	require.False(t, ok)
	require.Equal(t, time.Second, clientA.Timeout)
	require.Equal(t, 2*time.Second, clientB.Timeout)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	responseA, err := clientA.Project.ListProjects(ctx, &projectapi.ListProjectsRequest{})
	require.NoError(t, err)
	responseB, err := clientB.Project.ListProjects(ctx, &projectapi.ListProjectsRequest{})
	require.NoError(t, err)
	require.Equal(t, "from-a", responseA.Projects[0].Name)
	require.Equal(t, "from-b", responseB.Projects[0].Name)

	for expected, client := range map[string]*Client{"from-a": clientA, "from-b": clientB} {
		routes, err := client.APIProxy.ListRoutes(ctx, &apiproxyapi.ListRoutesRequest{})
		require.NoError(t, err)
		require.Equal(t, expected, routes.GetRoutes()[0].GetRouteId())
		secrets, err := client.SecretStore.ListSecrets(ctx, &secretstoreapi.ListSecretsRequest{})
		require.NoError(t, err)
		require.Equal(t, expected, secrets.GetSecrets()[0].GetName())
	}
	require.NoError(t, registry.Close())
	_, err = clientA.APIProxy.ListRoutes(ctx, &apiproxyapi.ListRoutesRequest{})
	require.Error(t, err)
	_, err = clientA.SecretStore.ListSecrets(ctx, &secretstoreapi.ListSecretsRequest{})
	require.Error(t, err)
	_, err = clientA.Project.ListProjects(context.Background(), &projectapi.ListProjectsRequest{})
	require.Error(t, err)
}

func TestRegistryDoesNotRequireSidecarToBeOnlineAtStartup(t *testing.T) {
	registry, err := NewRegistry(&config.RootConfig{AgentRuntime: config.AgentRuntimeConfig{Agents: map[string]config.AgentProviderConfig{
		"offline": {Sidecar: &config.AgentSidecarConfig{Endpoint: "127.0.0.1:1", Timeout: time.Millisecond}},
	}}})
	require.NoError(t, err)
	require.NoError(t, registry.Close())
}

type fakeProjectServer struct {
	projectapi.UnimplementedProjectServiceServer
	projectName string
	apiproxyapi.UnimplementedAPIProxyServiceServer
	secretstoreapi.UnimplementedSecretStoreServiceServer
}

func (s fakeProjectServer) ListProjects(context.Context, *projectapi.ListProjectsRequest) (*projectapi.ListProjectsResponse, error) {
	return &projectapi.ListProjectsResponse{Projects: []*projectentity.Project{{Name: s.projectName}}}, nil
}

func startFakeSidecar(t *testing.T, projectName string) *bufconn.Listener {
	t.Helper()
	listener := bufconn.Listen(1024 * 1024)
	server := grpc.NewServer()
	fake := fakeProjectServer{projectName: projectName}
	projectapi.RegisterProjectServiceServer(server, fake)
	apiproxyapi.RegisterAPIProxyServiceServer(server, fake)
	secretstoreapi.RegisterSecretStoreServiceServer(server, fake)
	go func() {
		_ = server.Serve(listener)
	}()
	t.Cleanup(func() {
		server.Stop()
		_ = listener.Close()
	})
	return listener
}

func (s fakeProjectServer) ListRoutes(context.Context, *apiproxyapi.ListRoutesRequest) (*apiproxyapi.ListRoutesResponse, error) {
	return &apiproxyapi.ListRoutesResponse{Routes: []*apiproxyentity.Route{{RouteId: s.projectName}}}, nil
}

func (s fakeProjectServer) ListSecrets(context.Context, *secretstoreapi.ListSecretsRequest) (*secretstoreapi.ListSecretsResponse, error) {
	return &secretstoreapi.ListSecretsResponse{Secrets: []*secretstoreentity.Secret{{Name: s.projectName}}}, nil
}