From 25d601920ec8c96049624b2ae9997ed2885f5550 Mon Sep 17 00:00:00 2001
From: Mikhail Filimonov <overtakeful@gmail.com>
Date: Tue, 17 Mar 2026 08:25:16 +0100
Subject: [PATCH] composeMessageWithAI

---
 .../Messages/TelegramEngineMessages.swift     |  4 ++
 .../TelegramEngine/Messages/Translate.swift   | 37 +++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/submodules/TelegramCore/Sources/TelegramEngine/Messages/TelegramEngineMessages.swift b/submodules/TelegramCore/Sources/TelegramEngine/Messages/TelegramEngineMessages.swift
index 38d7db955d..619b24158d 100644
--- a/submodules/TelegramCore/Sources/TelegramEngine/Messages/TelegramEngineMessages.swift
+++ b/submodules/TelegramCore/Sources/TelegramEngine/Messages/TelegramEngineMessages.swift
@@ -635,6 +635,10 @@ public extension TelegramEngine {
             }
         }
         
+        public func composeMessageWithAI(text: String, entities: [MessageTextEntity], proofread: Bool = false, translateToLang: String? = nil, changeTone: String? = nil, emojify: Bool = false) -> Signal<(String, [MessageTextEntity]), TranslationError> {
+            return _internal_composeMessageWithAI(network: self.account.network, text: text, entities: entities, proofread: proofread, translateToLang: translateToLang, changeTone: changeTone, emojify: emojify)
+        }
+
         public func translate(texts: [(String, [MessageTextEntity])], toLang: String, tone: TranslationTone = .neutral) -> Signal<[(String, [MessageTextEntity])], TranslationError> {
             return _internal_translateTexts(network: self.account.network, texts: texts, toLang: toLang, tone: tone)
         }
diff --git a/submodules/TelegramCore/Sources/TelegramEngine/Messages/Translate.swift b/submodules/TelegramCore/Sources/TelegramEngine/Messages/Translate.swift
index 04f1a29f62..7485f7db2c 100644
--- a/submodules/TelegramCore/Sources/TelegramEngine/Messages/Translate.swift
+++ b/submodules/TelegramCore/Sources/TelegramEngine/Messages/Translate.swift
@@ -68,6 +68,43 @@ func _internal_translate(network: Network, text: String, toLang: String, entitie
     }
 }
 
+func _internal_composeMessageWithAI(network: Network, text: String, entities: [MessageTextEntity], proofread: Bool = false, translateToLang: String? = nil, changeTone: String? = nil, emojify: Bool = false) -> Signal<(String, [MessageTextEntity]), TranslationError> {
+    var flags: Int32 = 0
+    if proofread {
+        flags |= (1 << 0)
+    }
+    if translateToLang != nil {
+        flags |= (1 << 1)
+    }
+    if changeTone != nil {
+        flags |= (1 << 2)
+    }
+    if emojify {
+        flags |= (1 << 3)
+    }
+
+    let apiText: Api.TextWithEntities = .textWithEntities(.init(text: text, entities: apiEntitiesFromMessageTextEntities(entities, associatedPeers: SimpleDictionary())))
+
+    return network.request(Api.functions.messages.composeMessageWithAI(flags: flags, text: apiText, translateToLang: translateToLang, changeTone: changeTone))
+    |> mapError { error -> TranslationError in
+        if error.errorDescription.hasPrefix("FLOOD_WAIT") {
+            return .limitExceeded
+        } else if error.errorDescription == "INPUT_TEXT_EMPTY" {
+            return .textIsEmpty
+        } else if error.errorDescription == "INPUT_TEXT_TOO_LONG" {
+            return .textTooLong
+        } else {
+            return .generic
+        }
+    }
+    |> map { result -> (String, [MessageTextEntity]) in
+        switch result {
+        case let .textWithEntities(data):
+            return (data.text, messageTextEntitiesFromApiEntities(data.entities))
+        }
+    }
+}
+
 func _internal_translateTexts(network: Network, texts: [(String, [MessageTextEntity])], toLang: String, tone: TranslationTone = .neutral) -> Signal<[(String, [MessageTextEntity])], TranslationError> {
     var flags: Int32 = 0
     flags |= (1 << 1)