6 files changed, 82 insertions, 17 deletions

diff --git a/src/core/hle/service/am/am.cpp b/src/core/hle/service/am/am.cpp
index 98765b81a..0886531b2 100644
--- a/src/core/hle/service/am/am.cpp
+++ b/src/core/hle/service/am/am.cpp
@@ -31,6 +31,7 @@
 #include "core/hle/service/apm/apm_controller.h"
 #include "core/hle/service/apm/apm_interface.h"
 #include "core/hle/service/bcat/backend/backend.h"
+#include "core/hle/service/caps/caps_su.h"
 #include "core/hle/service/caps/caps_types.h"
 #include "core/hle/service/filesystem/filesystem.h"
 #include "core/hle/service/ipc_helpers.h"
@@ -702,9 +703,17 @@ void ISelfController::SetAlbumImageTakenNotificationEnabled(HLERequestContext& c
 void ISelfController::SaveCurrentScreenshot(HLERequestContext& ctx) {
     IPC::RequestParser rp{ctx};
 
-    const auto album_report_option = rp.PopEnum<Capture::AlbumReportOption>();
+    const auto report_option = rp.PopEnum<Capture::AlbumReportOption>();
 
-    LOG_WARNING(Service_AM, "(STUBBED) called. album_report_option={}", album_report_option);
+    LOG_INFO(Service_AM, "called, report_option={}", report_option);
+
+    const auto screenshot_service =
+        system.ServiceManager().GetService<Service::Capture::IScreenShotApplicationService>(
+            "caps:su");
+
+    if (screenshot_service) {
+        screenshot_service->CaptureAndSaveScreenshot(report_option);
+    }
 
     IPC::ResponseBuilder rb{ctx, 2};
     rb.Push(ResultSuccess);
diff --git a/src/core/hle/service/caps/caps_manager.cpp b/src/core/hle/service/caps/caps_manager.cpp
index 7d733eb54..96b225d5f 100644
--- a/src/core/hle/service/caps/caps_manager.cpp
+++ b/src/core/hle/service/caps/caps_manager.cpp
@@ -228,12 +228,14 @@ Result AlbumManager::LoadAlbumScreenShotThumbnail(
 
 Result AlbumManager::SaveScreenShot(ApplicationAlbumEntry& out_entry,
                                     const ScreenShotAttribute& attribute,
-                                    std::span<const u8> image_data, u64 aruid) {
-    return SaveScreenShot(out_entry, attribute, {}, image_data, aruid);
+                                    AlbumReportOption report_option, std::span<const u8> image_data,
+                                    u64 aruid) {
+    return SaveScreenShot(out_entry, attribute, report_option, {}, image_data, aruid);
 }
 
 Result AlbumManager::SaveScreenShot(ApplicationAlbumEntry& out_entry,
                                     const ScreenShotAttribute& attribute,
+                                    AlbumReportOption report_option,
                                     const ApplicationData& app_data, std::span<const u8> image_data,
                                     u64 aruid) {
     const u64 title_id = system.GetApplicationProcessProgramID();
@@ -407,10 +409,14 @@ Result AlbumManager::LoadImage(std::span<u8> out_image, const std::filesystem::p
     return ResultSuccess;
 }
 
-static void PNGToMemory(void* context, void* png, int len) {
+void AlbumManager::FlipVerticallyOnWrite(bool flip) {
+    stbi_flip_vertically_on_write(flip);
+}
+
+static void PNGToMemory(void* context, void* data, int len) {
     std::vector<u8>* png_image = static_cast<std::vector<u8>*>(context);
-    png_image->reserve(len);
-    std::memcpy(png_image->data(), png, len);
+    unsigned char* png = static_cast<unsigned char*>(data);
+    png_image->insert(png_image->end(), png, png + len);
 }
 
 Result AlbumManager::SaveImage(ApplicationAlbumEntry& out_entry, std::span<const u8> image,
diff --git a/src/core/hle/service/caps/caps_manager.h b/src/core/hle/service/caps/caps_manager.h
index 44d85117f..e20c70c7b 100644
--- a/src/core/hle/service/caps/caps_manager.h
+++ b/src/core/hle/service/caps/caps_manager.h
@@ -59,14 +59,17 @@ public:
                                         const ScreenShotDecodeOption& decoder_options) const;
 
     Result SaveScreenShot(ApplicationAlbumEntry& out_entry, const ScreenShotAttribute& attribute,
-                          std::span<const u8> image_data, u64 aruid);
-    Result SaveScreenShot(ApplicationAlbumEntry& out_entry, const ScreenShotAttribute& attribute,
-                          const ApplicationData& app_data, std::span<const u8> image_data,
+                          AlbumReportOption report_option, std::span<const u8> image_data,
                           u64 aruid);
+    Result SaveScreenShot(ApplicationAlbumEntry& out_entry, const ScreenShotAttribute& attribute,
+                          AlbumReportOption report_option, const ApplicationData& app_data,
+                          std::span<const u8> image_data, u64 aruid);
     Result SaveEditedScreenShot(ApplicationAlbumEntry& out_entry,
                                 const ScreenShotAttribute& attribute, const AlbumFileId& file_id,
                                 std::span<const u8> image_data);
 
+    void FlipVerticallyOnWrite(bool flip);
+
 private:
     static constexpr std::size_t NandAlbumFileLimit = 1000;
     static constexpr std::size_t SdAlbumFileLimit = 10000;
diff --git a/src/core/hle/service/caps/caps_ss.cpp b/src/core/hle/service/caps/caps_ss.cpp
index 1ba2b7972..eab023568 100644
--- a/src/core/hle/service/caps/caps_ss.cpp
+++ b/src/core/hle/service/caps/caps_ss.cpp
@@ -34,7 +34,7 @@ void IScreenShotService::SaveScreenShotEx0(HLERequestContext& ctx) {
     IPC::RequestParser rp{ctx};
     struct Parameters {
         ScreenShotAttribute attribute{};
-        u32 report_option{};
+        AlbumReportOption report_option{};
         INSERT_PADDING_BYTES(0x4);
         u64 applet_resource_user_id{};
     };
@@ -49,13 +49,16 @@ void IScreenShotService::SaveScreenShotEx0(HLERequestContext& ctx) {
              parameters.applet_resource_user_id);
 
     ApplicationAlbumEntry entry{};
-    const auto result = manager->SaveScreenShot(entry, parameters.attribute, image_data_buffer,
-                                                parameters.applet_resource_user_id);
+    manager->FlipVerticallyOnWrite(false);
+    const auto result =
+        manager->SaveScreenShot(entry, parameters.attribute, parameters.report_option,
+                                image_data_buffer, parameters.applet_resource_user_id);
 
     IPC::ResponseBuilder rb{ctx, 10};
     rb.Push(result);
     rb.PushRaw(entry);
 }
+
 void IScreenShotService::SaveEditedScreenShotEx1(HLERequestContext& ctx) {
     IPC::RequestParser rp{ctx};
     struct Parameters {
@@ -83,6 +86,7 @@ void IScreenShotService::SaveEditedScreenShotEx1(HLERequestContext& ctx) {
              image_data_buffer.size(), thumbnail_image_data_buffer.size());
 
     ApplicationAlbumEntry entry{};
+    manager->FlipVerticallyOnWrite(false);
     const auto result = manager->SaveEditedScreenShot(entry, parameters.attribute,
                                                       parameters.file_id, image_data_buffer);
 
diff --git a/src/core/hle/service/caps/caps_su.cpp b/src/core/hle/service/caps/caps_su.cpp
index e85625ee4..296b07b00 100644
--- a/src/core/hle/service/caps/caps_su.cpp
+++ b/src/core/hle/service/caps/caps_su.cpp
@@ -2,10 +2,12 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 
 #include "common/logging/log.h"
+#include "core/core.h"
 #include "core/hle/service/caps/caps_manager.h"
 #include "core/hle/service/caps/caps_su.h"
 #include "core/hle/service/caps/caps_types.h"
 #include "core/hle/service/ipc_helpers.h"
+#include "video_core/renderer_base.h"
 
 namespace Service::Capture {
 
@@ -58,8 +60,10 @@ void IScreenShotApplicationService::SaveScreenShotEx0(HLERequestContext& ctx) {
              parameters.applet_resource_user_id);
 
     ApplicationAlbumEntry entry{};
-    const auto result = manager->SaveScreenShot(entry, parameters.attribute, image_data_buffer,
-                                                parameters.applet_resource_user_id);
+    manager->FlipVerticallyOnWrite(false);
+    const auto result =
+        manager->SaveScreenShot(entry, parameters.attribute, parameters.report_option,
+                                image_data_buffer, parameters.applet_resource_user_id);
 
     IPC::ResponseBuilder rb{ctx, 10};
     rb.Push(result);
@@ -88,13 +92,43 @@ void IScreenShotApplicationService::SaveScreenShotEx1(HLERequestContext& ctx) {
     ApplicationAlbumEntry entry{};
     ApplicationData app_data{};
     std::memcpy(&app_data, app_data_buffer.data(), sizeof(ApplicationData));
+    manager->FlipVerticallyOnWrite(false);
     const auto result =
-        manager->SaveScreenShot(entry, parameters.attribute, app_data, image_data_buffer,
-                                parameters.applet_resource_user_id);
+        manager->SaveScreenShot(entry, parameters.attribute, parameters.report_option, app_data,
+                                image_data_buffer, parameters.applet_resource_user_id);
 
     IPC::ResponseBuilder rb{ctx, 10};
     rb.Push(result);
     rb.PushRaw(entry);
 }
 
+void IScreenShotApplicationService::CaptureAndSaveScreenshot(AlbumReportOption report_option) {
+    auto& renderer = system.Renderer();
+    Layout::FramebufferLayout layout =
+        Layout::DefaultFrameLayout(screenshot_width, screenshot_height);
+
+    const Capture::ScreenShotAttribute attribute{
+        .unknown_0{},
+        .orientation = Capture::AlbumImageOrientation::None,
+        .unknown_1{},
+        .unknown_2{},
+    };
+
+    renderer.RequestScreenshot(
+        image_data.data(),
+        [attribute, report_option, this](bool invert_y) {
+            // Convert from BGRA to RGBA
+            for (std::size_t i = 0; i < image_data.size(); i += bytes_per_pixel) {
+                const u8 temp = image_data[i];
+                image_data[i] = image_data[i + 2];
+                image_data[i + 2] = temp;
+            }
+
+            Capture::ApplicationAlbumEntry entry{};
+            manager->FlipVerticallyOnWrite(invert_y);
+            manager->SaveScreenShot(entry, attribute, report_option, image_data, {});
+        },
+        layout);
+}
+
 } // namespace Service::Capture
diff --git a/src/core/hle/service/caps/caps_su.h b/src/core/hle/service/caps/caps_su.h
index 89e71f506..21912e95f 100644
--- a/src/core/hle/service/caps/caps_su.h
+++ b/src/core/hle/service/caps/caps_su.h
@@ -10,6 +10,7 @@ class System;
 }
 
 namespace Service::Capture {
+enum class AlbumReportOption : s32;
 class AlbumManager;
 
 class IScreenShotApplicationService final : public ServiceFramework<IScreenShotApplicationService> {
@@ -18,11 +19,19 @@ public:
                                            std::shared_ptr<AlbumManager> album_manager);
     ~IScreenShotApplicationService() override;
 
+    void CaptureAndSaveScreenshot(AlbumReportOption report_option);
+
 private:
+    static constexpr std::size_t screenshot_width = 1280;
+    static constexpr std::size_t screenshot_height = 720;
+    static constexpr std::size_t bytes_per_pixel = 4;
+
     void SetShimLibraryVersion(HLERequestContext& ctx);
     void SaveScreenShotEx0(HLERequestContext& ctx);
     void SaveScreenShotEx1(HLERequestContext& ctx);
 
+    std::array<u8, screenshot_width * screenshot_height * bytes_per_pixel> image_data;
+
     std::shared_ptr<AlbumManager> manager;
 };