66import it .tdlight .Init ;
77import it .tdlight .ResultHandler ;
88import it .tdlight .TelegramClient ;
9+ import it .tdlight .jni .TdApi .Message ;
910import it .tdlight .jni .TdApi .Update ;
1011import it .tdlight .util .UnsupportedNativeLibraryException ;
1112import it .tdlight .jni .TdApi ;
2021import java .util .Objects ;
2122import java .util .Set ;
2223import java .util .concurrent .CompletableFuture ;
24+ import java .util .concurrent .CompletionStage ;
25+ import java .util .concurrent .ConcurrentHashMap ;
26+ import java .util .concurrent .ConcurrentMap ;
2327import java .util .concurrent .CountDownLatch ;
28+ import java .util .concurrent .ExecutionException ;
29+ import java .util .concurrent .Future ;
2430import java .util .concurrent .RejectedExecutionException ;
2531import org .slf4j .Logger ;
2632import org .slf4j .LoggerFactory ;
@@ -54,7 +60,8 @@ public final class SimpleTelegramClient implements Authenticable, MutableTelegra
5460 private final AuthorizationStateReadyLoadChats mainChatsLoader ;
5561 private final AuthorizationStateReadyLoadChats archivedChatsLoader ;
5662
57- private final CountDownLatch closed = new CountDownLatch (1 );
63+ private final CompletableFuture <Void > closed = new CompletableFuture <>();
64+ private final ConcurrentMap <TemporaryMessageURL , CompletableFuture <Message >> temporaryMessages = new ConcurrentHashMap <>();
5865
5966 SimpleTelegramClient (ClientFactory clientFactory ,
6067 TDLibSettings settings ,
@@ -109,12 +116,15 @@ public final class SimpleTelegramClient implements Authenticable, MutableTelegra
109116 this ::handleDefaultException
110117 )
111118 );
112- this .addUpdateHandler (TdApi .UpdateAuthorizationState .class , new AuthorizationStateWaitForExit (this . closed ));
119+ this .addUpdateHandler (TdApi .UpdateAuthorizationState .class , new AuthorizationStateWaitForExit (this :: onCloseUpdate ));
113120 this .mainChatsLoader = new AuthorizationStateReadyLoadChats (client , new ChatListMain ());
114121 this .archivedChatsLoader = new AuthorizationStateReadyLoadChats (client , new ChatListArchive ());
115122 this .addUpdateHandler (TdApi .UpdateAuthorizationState .class ,
116123 this .meGetter = new AuthorizationStateReadyGetMe (client , mainChatsLoader , archivedChatsLoader ));
117124 this .addUpdateHandler (TdApi .UpdateNewMessage .class , new CommandsHandler (client , this .commandHandlers , this ::getMe ));
125+ var temporaryMessageHandler = new TemporaryMessageHandler (this .temporaryMessages );
126+ this .addUpdateHandler (TdApi .UpdateMessageSendSucceeded .class , temporaryMessageHandler );
127+ this .addUpdateHandler (TdApi .UpdateMessageSendFailed .class , temporaryMessageHandler );
118128 this .authenticationData = authenticationData ;
119129 createDirectories ();
120130 client .initialize (this ::handleUpdate , this ::handleUpdateException , this ::handleDefaultException );
@@ -191,7 +201,7 @@ public <T extends TdApi.Update> void addCommandHandler(String commandName, Comma
191201 }
192202
193203 @ Override
194- public <T extends TdApi .Update > void addUpdateHandler (Class <T > updateType , GenericUpdateHandler <T > handler ) {
204+ public <T extends TdApi .Update > void addUpdateHandler (Class <T > updateType , GenericUpdateHandler <? super T > handler ) {
195205 this .updateHandlers .put (new SimpleResultHandler <>(updateType , handler ), null );
196206 }
197207
@@ -257,6 +267,44 @@ public <R extends TdApi.Object> void send(TdApi.Function<R> function, GenericRes
257267 client .send (function , result -> resultHandler .onResult (Result .of (result )), resultHandlerExceptionHandler );
258268 }
259269
270+ /**
271+ * Sends a request to TDLib and get the result.
272+ *
273+ * @param function The request to TDLib.
274+ * @throws NullPointerException if function is null.
275+ */
276+ @ SuppressWarnings ("unchecked" )
277+ public <R extends TdApi .Object > CompletableFuture <R > send (TdApi .Function <R > function ) {
278+ CompletableFuture <R > future = new CompletableFuture <>();
279+ client .send (function , result -> {
280+ if (result instanceof TdApi .Error ) {
281+ future .completeExceptionally (new TelegramError ((TdApi .Error ) result ));
282+ } else {
283+ future .complete ((R ) result );
284+ }
285+ }, future ::completeExceptionally );
286+ return future ;
287+ }
288+
289+ /**
290+ * Sends a message
291+ *
292+ * @param function The request to TDLib.
293+ * @param wait Wait until the message has been sent.
294+ * If false, the returned message will not represent the real message, but it will be a temporary message.
295+ * @throws NullPointerException if function is null.
296+ */
297+ public CompletableFuture <TdApi .Message > sendMessage (TdApi .SendMessage function , boolean wait ) {
298+ return this .send (function ).thenCompose (msg -> {
299+ CompletableFuture <TdApi .Message > future = new CompletableFuture <>();
300+ CompletableFuture <?> prev = temporaryMessages .put (new TemporaryMessageURL (msg .chatId , msg .id ), future );
301+ if (prev != null ) {
302+ prev .completeExceptionally (new IllegalStateException ("Another temporary message has the same id" ));
303+ }
304+ return future ;
305+ });
306+ }
307+
260308 /**
261309 * Execute a synchronous function.
262310 * <strong>Please note that only some functions can be executed using this method.</strong>
@@ -266,6 +314,13 @@ public <R extends TdApi.Object> Result<R> execute(TdApi.Function<R> function) {
266314 return Result .of (client .execute (function ));
267315 }
268316
317+ /**
318+ * Send the close signal but don't wait
319+ */
320+ public CompletableFuture <Void > close () {
321+ return this .send (new TdApi .Close ()).thenCompose (x -> this .waitForExitAsync ());
322+ }
323+
269324 /**
270325 * Send the close signal but don't wait
271326 */
@@ -293,7 +348,23 @@ public void closeAndWait() throws InterruptedException {
293348 * Wait until TDLight is closed
294349 */
295350 public void waitForExit () throws InterruptedException {
296- closed .await ();
351+ try {
352+ closed .get ();
353+ } catch (ExecutionException e ) {
354+ throw new RuntimeException (e );
355+ }
356+ }
357+
358+ /**
359+ * Wait until TDLight is closed
360+ */
361+ public CompletableFuture <Void > waitForExitAsync () {
362+ return closed .copy ();
363+ }
364+
365+ private void onCloseUpdate () {
366+ this .closed .complete (null );
367+ this .temporaryMessages .clear ();
297368 }
298369
299370 private final class SimpleTelegramClientInteraction implements ClientInteraction {
0 commit comments