1/*
2 * Copyright (c) 2023 Huawei Device Co., Ltd.
3 * Licensed under the Apache License, Version 2.0 (the "License");
4 * you may not use this file except in compliance with the License.
5 * You may obtain a copy of the License at
6 *
7 *     http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software
10 * distributed under the License is distributed on an "AS IS" BASIS,
11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 * See the License for the specific language governing permissions and
13 * limitations under the License.
14 */
15
16#include <condition_variable>
17#include <display_manager.h>
18#include <jpeglib.h>
19#include <memory>
20#include <mutex>
21#include <pixel_map.h>
22#include <refbase.h>
23#include <screen_manager.h>
24#include <securec.h>
25#include <csetjmp>
26#include "common_utilities_hpp.h"
27#include "screen_copy.h"
28
29namespace OHOS::uitest {
30using namespace std;
31using namespace OHOS;
32using namespace OHOS::Media;
33using namespace OHOS::Rosen;
34
35using OnScreenChangeHandler = function<void()>;
36class ScreenChangeListener : public ScreenManager::IScreenListener {
37public:
38    explicit ScreenChangeListener(OnScreenChangeHandler hdl, ScreenId id): handler_(hdl), targetId_(id) {}
39    void OnConnect(ScreenId id) override {};
40    void OnDisconnect(ScreenId id) override {};
41    void OnChange(ScreenId id) override
42    {
43        if (handler_ != nullptr && id == targetId_) {
44            handler_();
45        }
46    };
47    void Destroy() { handler_ = nullptr; }
48private:
49    OnScreenChangeHandler handler_ = nullptr;
50    ScreenId targetId_ = SCREEN_ID_INVALID;
51};
52
53class ScreenCopy {
54public:
55    explicit ScreenCopy(float scale): scale_(scale) {};
56    virtual ~ScreenCopy();
57    bool Run();
58    void Destroy();
59    const char* pendingError_ = nullptr;
60    const float scale_ = 0.5f;
61private:
62    void PollAndNotifyFrames();
63    void WaitAndConsumeFrames();
64    void UpdateFrameLocked(shared_ptr<PixelMap> frame, bool &changed, bool &muted);
65    shared_ptr<PixelMap> ScaleNewsetFrameLocked();
66    sptr<Screen> sourceScreen_;
67    shared_ptr<PixelMap> lastFrame_ = nullptr;
68    shared_ptr<PixelMap> newestFrame_ = nullptr;
69    mutex frameLock_;
70    condition_variable frameCond_;
71    unique_ptr<thread> snapshotThread = nullptr;
72    unique_ptr<thread> encodeThread = nullptr;
73    atomic_bool stopped_ = false;
74    static sptr<ScreenChangeListener> screenChangeListener_;
75};
76sptr<ScreenChangeListener> ScreenCopy::screenChangeListener_ = nullptr;
77static unique_ptr<ScreenCopy> g_screenCopy = nullptr;
78static ScreenCopyHandler g_screenCopyHandler = nullptr;
79
80
81static void AdapteScreenChange()
82{
83    if (g_screenCopy == nullptr) {
84        return;
85    }
86    // destrory current one and create a new one
87    LOG_D("Screen changed, auto restart ScreenCopy");
88    const auto scale = g_screenCopy->scale_;
89    g_screenCopy->Destroy();
90    g_screenCopy = make_unique<ScreenCopy>(scale);
91    g_screenCopy->Run();
92}
93
94ScreenCopy::~ScreenCopy()
95{
96    if (!stopped_.load()) {
97        Destroy();
98    }
99}
100
101bool ScreenCopy::Run()
102{
103    if (scale_ <= 0 || scale_ > 1.0) {
104        pendingError_ = "Error: Illegal scale value!";
105        return false;
106    }
107    // get source screen
108    auto id = static_cast<ScreenId>(DisplayManager::GetInstance().GetDefaultDisplayId());
109    sourceScreen_ = ScreenManager::GetInstance().GetScreenById(id);
110    if (id == SCREEN_ID_INVALID || sourceScreen_ == nullptr) {
111        pendingError_ = "Error: Get main screen failed!";
112        return false;
113    }
114    // listen screen changes for auto-adapting
115    if (screenChangeListener_ == nullptr) {
116        screenChangeListener_ = new ScreenChangeListener([]() { AdapteScreenChange(); }, id);
117        auto ret = ScreenManager::GetInstance().RegisterScreenListener(screenChangeListener_);
118        LOG_D("Register ScreenListener, ret=%{public}d", ret);
119    }
120    // run snapshot thread and encode thread
121    snapshotThread = make_unique<thread>([this]() { this->PollAndNotifyFrames(); });
122    encodeThread = make_unique<thread>([this]() { this->WaitAndConsumeFrames(); });
123    return true;
124}
125
126void ScreenCopy::Destroy()
127{
128    if (stopped_.load()) {
129        return;
130    }
131    unique_lock<mutex> lock(frameLock_);
132    stopped_.store(true);
133    frameCond_.notify_all(); // mark stopped and wakeup the waiting thread
134    lock.unlock();
135    LOG_D("Begin to wait for threads exit");
136    if (snapshotThread != nullptr && snapshotThread->joinable()) {
137        snapshotThread->join();
138        snapshotThread = nullptr;
139    }
140    sourceScreen_ = nullptr;
141    lastFrame_ = nullptr;
142    newestFrame_ = nullptr;
143    if (encodeThread != nullptr && encodeThread->joinable()) {
144        encodeThread->join();
145        encodeThread = nullptr;
146    }
147    LOG_D("All threads exited");
148}
149
150void ScreenCopy::PollAndNotifyFrames()
151{
152    constexpr int32_t screenCheckIntervalUs = 50 * 1000;
153    auto &dm = DisplayManager::GetInstance();
154    const auto displayId = dm.GetDefaultDisplayId();
155    LOG_I("Start PollAndNotifyFrames");
156    bool changed = false;
157    bool screenOff = false;
158    while (!stopped_.load()) {
159        if (screenOff) {
160            usleep(screenCheckIntervalUs);
161            if (dm.GetDisplayState(displayId) != DisplayState::OFF) {
162                screenOff = false;
163                LOG_I("Screen turned on! resume screenCopy");
164            }
165            continue;
166        }
167        shared_ptr<PixelMap> frame = dm.GetScreenshot(displayId);
168        if (frame == nullptr) {
169            continue;
170        }
171        frameLock_.lock();
172        UpdateFrameLocked(frame, changed, screenOff);
173        frameLock_.unlock();
174        LOG_D("GetOneFrameDone, Changed=%{public}d", changed);
175        if (changed) {
176            frameCond_.notify_all();
177        }
178        if (screenOff) {
179            LOG_I("Screen turned off! mute screenCopy");
180        }
181    }
182}
183
184void ScreenCopy::UpdateFrameLocked(shared_ptr<PixelMap> frame, bool &changed, bool &screenOff)
185{
186    lastFrame_ = newestFrame_;
187    newestFrame_ = frame;
188    changed = true;
189    const size_t newestFrameSize = newestFrame_->GetHeight() * newestFrame_->GetRowStride();
190    // if this is the first frame
191    if (lastFrame_ == nullptr) {
192        // if screen copy starts with screen-off, given a black image
193        if (DisplayManager::GetInstance().GetDisplayState(sourceScreen_->GetId()) == DisplayState::OFF) {
194            memset_s(frame->GetWritablePixels(), newestFrameSize, 0, newestFrameSize);
195        }
196        return;
197    }
198    // compare this frame and last frame
199    const size_t lastFrameSize = lastFrame_->GetHeight() * lastFrame_->GetRowStride();
200    if (lastFrameSize == newestFrameSize) {
201        changed = memcmp(lastFrame_->GetPixels(), newestFrame_->GetPixels(), newestFrameSize) != 0;
202    }
203    // detect screen of only when not changed
204    if (!changed && !screenOff) {
205        screenOff = DisplayManager::GetInstance().GetDisplayState(sourceScreen_->GetId()) == DisplayState::OFF;
206        if (screenOff) {
207            // mark changed and reset pixels to black so we provide a black image
208            changed = true;
209            memset_s(frame->GetWritablePixels(), newestFrameSize, 0, newestFrameSize);
210        }
211    }
212}
213
214struct MissionErrorMgr : public jpeg_error_mgr {
215    jmp_buf setjmp_buffer;
216};
217
218static void AdaptJpegSize(jpeg_compress_struct &jpeg, uint32_t width, uint32_t height)
219{
220    constexpr int32_t alignment = 32;
221    if (width % alignment == 0) {
222        jpeg.image_width = width;
223    } else {
224        LOG_D("The width need to be adapted!");
225        jpeg.image_width = ceil((double)width / (double)alignment) * alignment;
226    }
227    jpeg.image_height = height;
228}
229
230shared_ptr<PixelMap> ScreenCopy::ScaleNewsetFrameLocked()
231{
232    if (newestFrame_ == nullptr) {
233        return newestFrame_;
234    }
235    // resize the pixelmap to fit scale
236    auto origWidth = newestFrame_->GetWidth();
237    auto origHeight = newestFrame_->GetHeight();
238    Media::Rect rect = {.left = 0, .top = 0, .width = origWidth, .height = origHeight};
239    Media::InitializationOptions opt;
240    opt.size.width = ceil(origWidth * scale_);
241    opt.size.height = ceil(origHeight * scale_);
242    opt.scaleMode = Media::ScaleMode::FIT_TARGET_SIZE;
243    opt.editable = false;
244    return Media::PixelMap::Create(*newestFrame_, rect, opt);
245}
246
247void ScreenCopy::WaitAndConsumeFrames()
248{
249    LOG_I("Start WaitAndConsumeFrames");
250    while (!stopped_.load()) {
251        unique_lock<mutex> lock(frameLock_);
252        frameCond_.wait(lock);
253        if (stopped_.load()) {
254            break;
255        }
256        LOG_D("ConsumeFrame_Begin");
257        // resize the pixelmap to fit scale
258        auto scaledPixels = ScaleNewsetFrameLocked();
259        lock.unlock();
260        if (scaledPixels == nullptr) {
261            continue;
262        }
263        constexpr int32_t rgbaPixelBytes = 4;
264        jpeg_compress_struct jpeg;
265        MissionErrorMgr jerr;
266        jpeg.err = jpeg_std_error(&jerr);
267        jpeg_create_compress(&jpeg);
268        AdaptJpegSize(jpeg, scaledPixels->GetWidth(), scaledPixels->GetHeight());
269        jpeg.input_components = rgbaPixelBytes;
270        jpeg.in_color_space = JCS_EXT_RGBX;
271        jpeg_set_defaults(&jpeg);
272        constexpr int32_t compressQuality = 75;
273        jpeg_set_quality(&jpeg, compressQuality, 1);
274        uint8_t *imgBuf = nullptr;
275        unsigned long imgSize = 0;
276        jpeg_mem_dest(&jpeg, &imgBuf, &imgSize);
277        jpeg_start_compress(&jpeg, 1);
278        JSAMPROW rowPointer[1024 * 4];
279        const auto stride = scaledPixels->GetRowStride();
280        auto memAddr = (uint8_t *)(scaledPixels->GetPixels());
281        for (int32_t rowIndex = 0; rowIndex < scaledPixels->GetHeight(); rowIndex++) {
282            rowPointer[rowIndex] = memAddr;
283            memAddr += stride;
284        }
285        (void)jpeg_write_scanlines(&jpeg, rowPointer, jpeg.image_height);
286        jpeg_finish_compress(&jpeg);
287        jpeg_destroy_compress(&jpeg);
288        LOG_D("ConsumeFrame_End, size=%{public}lu", imgSize);
289        if (g_screenCopyHandler != nullptr) {
290            g_screenCopyHandler(imgBuf, imgSize);
291        } else {
292            free(imgBuf);
293        }
294    }
295    LOG_I("Stop WaitAndConsumeFrames");
296}
297
298bool StartScreenCopy(float scale, ScreenCopyHandler handler)
299{
300    if (scale <= 0 || scale > 1 || handler == nullptr) {
301        LOG_E("Illegal arguments");
302        return false;
303    }
304    StopScreenCopy();
305    g_screenCopyHandler = handler;
306    g_screenCopy = make_unique<ScreenCopy>(scale);
307    bool success = g_screenCopy != nullptr && g_screenCopy->Run();
308    if (!success) {
309        constexpr size_t BUF_SIZE = 128;
310        auto buf = (uint8_t *)malloc(BUF_SIZE);
311        memset_s(buf, BUF_SIZE, 0, BUF_SIZE);
312        if (g_screenCopy->pendingError_ != nullptr) {
313            g_screenCopy->pendingError_ = "Failed to run ScreenCopy, unknown error";
314        }
315        memcpy_s(buf, BUF_SIZE, g_screenCopy->pendingError_, strlen(g_screenCopy->pendingError_));
316        LOG_E("The error message is %{public}s", buf);
317        handler(buf, strlen(g_screenCopy->pendingError_));
318        g_screenCopy->pendingError_ = nullptr;
319        StopScreenCopy();
320    }
321    return success;
322}
323
324void StopScreenCopy()
325{
326    if (g_screenCopy != nullptr) {
327        g_screenCopy->Destroy();
328        g_screenCopy = nullptr;
329    }
330}
331}
332