package tunserver

import (
	"context"
	"log/slog"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	"google.golang.org/grpc"
)

type FindHandle interface {
	// Get finds a tunnel to an agentk.
	// It waits for a matching tunnel to proxy a connection through. When a matching tunnel is found, it is returned.
	// It returns gRPC status errors only, ready to return from RPC handler.
	Get(ctx context.Context) (Tunnel, error)
	// Done must be called to free resources of this FindHandle instance.
	// ctx is used for tracing only.
	Done(ctx context.Context)
}

type RouterPlugin[T Target] interface {
	// FindReadyGateway finds a ready gateway tunserver to route the request through.
	// It returns an error, compatible with gRPC status package.
	FindReadyGateway(ctx context.Context, log *slog.Logger, method string) (ReadyGateway[T], *slog.Logger, api.AgentKey, error)
	// FindTunnel starts searching for a tunnel to a matching tunclient.
	// Tunnel found boolean indicates whether a suitable tunnel is immediately available from the
	// returned FindHandle object.
	FindTunnel(grpc.ServerStream, modshared.RPCAPI) (bool, *slog.Logger, FindHandle, error)
	// PrepareStreamForForwarding is used to perform any actions, like wrapping, on the stream just before
	// forwarding it to tunclient.
	// Returned error is compatible with gRPC-status error.
	PrepareStreamForForwarding(stream grpc.ServerStream) (grpc.ServerStream, error)
}

// Router routes traffic from kas to another kas to agentk.
// routing kas -> gateway kas -> agentk
type Router[T Target] struct {
	Plugin RouterPlugin[T]
	// PrivateAPIServer is the gRPC server that other kas instances can talk to.
	// Request handlers can obtain the per-request logger using modshared.RPCAPIFromContext(requestContext).
	PrivateAPIServer grpc.ServiceRegistrar
}

func (r *Router[T]) RegisterTunclientAPI(desc *grpc.ServiceDesc) {
	// 1. Munge the descriptor into the right shape:
	//    - turn all unary calls into streaming calls
	//    - all streaming calls, including the ones from above, are handled by routing handlers
	privateAPIServerDesc := mungeDescriptor(desc, r.routeToTunclient)

	// 2. Register on PrivateApiServer gRPC server so that this tunserver instance can act as the gateway tunserver instance
	//    and route to one of the matching connected tunclient instances.
	r.PrivateAPIServer.RegisterService(privateAPIServerDesc, nil)
}

func mungeDescriptor(in *grpc.ServiceDesc, handler grpc.StreamHandler) *grpc.ServiceDesc {
	streams := make([]grpc.StreamDesc, 0, len(in.Streams)+len(in.Methods))
	for _, stream := range in.Streams {
		streams = append(streams, grpc.StreamDesc{
			StreamName:    stream.StreamName,
			Handler:       handler,
			ServerStreams: true,
			ClientStreams: true,
		})
	}
	// Turn all methods into streams
	for _, method := range in.Methods {
		streams = append(streams, grpc.StreamDesc{
			StreamName:    method.MethodName,
			Handler:       handler,
			ServerStreams: true,
			ClientStreams: true,
		})
	}
	return &grpc.ServiceDesc{
		ServiceName: in.ServiceName,
		Streams:     streams,
		Metadata:    in.Metadata,
	}
}
