From 82a59e4e609680d852b3e4e1f05d9fcf499d41a9 Mon Sep 17 00:00:00 2001
From: Ilya Laktyushin <ilya@laktyush.in>
Date: Sun, 11 Jun 2023 17:56:32 +0400
Subject: [PATCH] Camera and editor improvements

---
 .../Sources/MediaEditorComposerEntity.swift   | 71 ++++++++++++++-----
 .../Sources/MediaEditorScreen.swift           |  6 ++
 2 files changed, 60 insertions(+), 17 deletions(-)

diff --git a/submodules/TelegramUI/Components/MediaEditor/Sources/MediaEditorComposerEntity.swift b/submodules/TelegramUI/Components/MediaEditor/Sources/MediaEditorComposerEntity.swift
index 7737380486..dc069ab269 100644
--- a/submodules/TelegramUI/Components/MediaEditor/Sources/MediaEditorComposerEntity.swift
+++ b/submodules/TelegramUI/Components/MediaEditor/Sources/MediaEditorComposerEntity.swift
@@ -197,19 +197,10 @@ private class MediaEditorComposerStickerEntity: MediaEditorComposerEntity {
                 tintColor = .white
             }
             
-            self.disposables.add((self.frameSource.get()
-            |> take(1)
-            |> deliverOn(self.queue)).start(next: { [weak self] frameSource in
+            let processFrame: (Double, Int, (Int) -> AnimatedStickerFrame?) -> Void = { [weak self] duration, frameCount, takeFrame in
                 guard let strongSelf = self else {
-                    completion(nil)
                     return
                 }
-                
-                guard let frameSource, let duration = strongSelf.totalDuration, let frameCount = strongSelf.frameCount else {
-                    completion(nil)
-                    return
-                }
-                                                
                 let relativeTime = currentTime - floor(currentTime / duration) * duration
                 var t = relativeTime / duration
                 t = max(0.0, t)
@@ -233,12 +224,8 @@ private class MediaEditorComposerStickerEntity: MediaEditorComposerEntity {
                         delta = max(1, frameIndex - previousFrameIndex)
                     }
                     
-                    var frame: AnimatedStickerFrame?
-                    frameSource.syncWith { frameSource in
-                        for i in 0 ..< delta {
-                            frame = frameSource.takeFrame(draw: i == delta - 1)
-                        }
-                    }
+                    let frame = takeFrame(delta)
+                    
                     if let frame {
                         var imagePixelBuffer: CVPixelBuffer?
                         if let pixelBuffer = strongSelf.imagePixelBuffer {
@@ -273,7 +260,57 @@ private class MediaEditorComposerStickerEntity: MediaEditorComposerEntity {
                 } else {
                     completion(strongSelf.image)
                 }
-            }))
+            }
+            
+            if self.isVideo {
+                self.disposables.add((self.videoFrameSource.get()
+                |> take(1)
+                |> deliverOn(self.queue)).start(next: { [weak self] frameSource in
+                    guard let strongSelf = self else {
+                        completion(nil)
+                        return
+                    }
+                    
+                    guard let frameSource, let duration = strongSelf.totalDuration, let frameCount = strongSelf.frameCount else {
+                        completion(nil)
+                        return
+                    }
+                    
+                    processFrame(duration, frameCount, { delta in
+                        var frame: AnimatedStickerFrame?
+                        frameSource.syncWith { frameSource in
+                            for i in 0 ..< delta {
+                                frame = frameSource.takeFrame(draw: i == delta - 1)
+                            }
+                        }
+                        return frame
+                    })
+                }))
+            } else {
+                self.disposables.add((self.frameSource.get()
+                |> take(1)
+                |> deliverOn(self.queue)).start(next: { [weak self] frameSource in
+                    guard let strongSelf = self else {
+                        completion(nil)
+                        return
+                    }
+                    
+                    guard let frameSource, let duration = strongSelf.totalDuration, let frameCount = strongSelf.frameCount else {
+                        completion(nil)
+                        return
+                    }
+                    
+                    processFrame(duration, frameCount, { delta in
+                        var frame: AnimatedStickerFrame?
+                        frameSource.syncWith { frameSource in
+                            for i in 0 ..< delta {
+                                frame = frameSource.takeFrame(draw: i == delta - 1)
+                            }
+                        }
+                        return frame
+                    })
+                }))
+            }
         } else {
             var image: CIImage?
             if let cachedImage = self.image {
diff --git a/submodules/TelegramUI/Components/MediaEditorScreen/Sources/MediaEditorScreen.swift b/submodules/TelegramUI/Components/MediaEditorScreen/Sources/MediaEditorScreen.swift
index 8d41b7e1ba..9f182499d4 100644
--- a/submodules/TelegramUI/Components/MediaEditorScreen/Sources/MediaEditorScreen.swift
+++ b/submodules/TelegramUI/Components/MediaEditorScreen/Sources/MediaEditorScreen.swift
@@ -2059,6 +2059,9 @@ public final class MediaEditorScreen: ViewController {
                                         if let self, let file {
                                             let stickerEntity = DrawingStickerEntity(content: .file(file))
                                             self.interaction?.insertEntity(stickerEntity)
+                                            
+                                            self.controller?.isSavingAvailable = true
+                                            self.controller?.requestLayout(transition: .immediate)
                                         }
                                     }
                                     self.controller?.present(controller, in: .current)
@@ -2066,6 +2069,9 @@ public final class MediaEditorScreen: ViewController {
                                 case .text:
                                     let textEntity = DrawingTextEntity(text: NSAttributedString(), style: .regular, animation: .none, font: .sanFrancisco, alignment: .center, fontSize: 1.0, color: DrawingColor(color: .white))
                                     self.interaction?.insertEntity(textEntity)
+                                    
+                                    self.controller?.isSavingAvailable = true
+                                    self.controller?.requestLayout(transition: .immediate)
                                     return
                                 case .drawing:
                                     self.interaction?.deactivate()