/*
 * Copyright (C) 2010 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <android/bitmap.h>
#include <android/graphics/bitmap.h>
#include <android/gui/DisplayCaptureArgs.h>
#include <binder/ProcessState.h>
#include <errno.h>
#include <fcntl.h>
#include <ftl/concat.h>
#include <ftl/optional.h>
#include <getopt.h>
#include <gui/ISurfaceComposer.h>
#include <gui/SurfaceComposerClient.h>
#include <gui/SyncScreenCaptureListener.h>
#include <linux/fb.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/wait.h>
#include <system/graphics.h>
#include <ui/GraphicTypes.h>
#include <ui/PixelFormat.h>

#include "aidl/android/hardware/graphics/common/PixelFormat.h"
#include "screencap_utils.h"
#include "utils/Errors.h"

using namespace android;

#define COLORSPACE_UNKNOWN 0
#define COLORSPACE_SRGB 1
#define COLORSPACE_DISPLAY_P3 2

void usage(const char* pname, ftl::Optional<DisplayId> displayIdOpt) {
    fprintf(stderr, R"(
usage: %s [-ahp] [-d display-id] [FILENAME]
   -h: this message
   -a: captures all the active displays. This appends an integer postfix to the FILENAME.
       e.g., FILENAME_0.png, FILENAME_1.png. If both -a and -d are given, it ignores -d.
   -d: specify the display ID to capture%s
       see "dumpsys SurfaceFlinger --display-id" for valid display IDs.
   -p: outputs in png format.
   --preserve-display-colors: Set to true to preserves the native display colorspace. Useful
       for mixed HDR + SDR content, using identical processing as the display's
   --require-dpu: Set to true to force the device to use the DPU to composite the screenshot.
       Does nothing if the device is not able to allocate a display resource to use the DPU.
       Setting --require-dpu also implies --preserve-display-colors.

If FILENAME ends with .png it will be saved as a png.
If FILENAME is not given, the results will be printed to stdout.
)",
            pname,
            displayIdOpt
                    .transform([](DisplayId id) {
                        return std::string(ftl::Concat(" (If the id is not given, it defaults to ",
                                                       id.value, ')')
                                                   .str());
                    })
                    .value_or(std::string())
                    .c_str());
}

// For options that only exist in long-form. Anything in the
// 0-255 range is reserved for short options (which just use their ASCII value)
namespace LongOpts {
enum {
    Reserved = 255,
    PreservedDisplayColors,
    RequireDpu,
};
}

static const struct option LONG_OPTIONS[] = {{"png", no_argument, nullptr, 'p'},
                                             {"jpeg", no_argument, nullptr, 'j'},
                                             {"help", no_argument, nullptr, 'h'},
                                             {"preserve-display-colors", no_argument, nullptr,
                                              LongOpts::PreservedDisplayColors},
                                             {"require-dpu", no_argument, nullptr,
                                              LongOpts::RequireDpu},
                                             {0, 0, 0, 0}};

static int32_t flinger2bitmapFormat(aidl::android::hardware::graphics::common::PixelFormat f) {
    switch (f) {
        case aidl::android::hardware::graphics::common::PixelFormat::RGB_565:
            return ANDROID_BITMAP_FORMAT_RGB_565;
        case aidl::android::hardware::graphics::common::PixelFormat::RGBA_1010102:
            return ANDROID_BITMAP_FORMAT_RGBA_1010102;
        case aidl::android::hardware::graphics::common::PixelFormat::BGRA_1010102:
            return ANDROID_BITMAP_FORMAT_BGRA_1010102;
        case aidl::android::hardware::graphics::common::PixelFormat::BGRX_1010102:
            return ANDROID_BITMAP_FORMAT_BGRX_1010102;
        default:
            return ANDROID_BITMAP_FORMAT_RGBA_8888;
    }
}

static uint32_t dataSpaceToInt(ui::Dataspace d) {
    switch (d) {
        case ui::Dataspace::V0_SRGB:
            return COLORSPACE_SRGB;
        case ui::Dataspace::DISPLAY_P3:
            return COLORSPACE_DISPLAY_P3;
        default:
            return COLORSPACE_UNKNOWN;
    }
}

static status_t notifyMediaScanner(const char* fileName) {
    std::string filePath("file://");
    filePath.append(fileName);
    char* cmd[] = {(char*)"am", (char*)"broadcast",
                   (char*)"-a", (char*)"android.intent.action.MEDIA_SCANNER_SCAN_FILE",
                   (char*)"-d", &filePath[0],
                   nullptr};

    int status;
    int pid = fork();
    if (pid < 0) {
        fprintf(stderr, "Unable to fork in order to send intent for media scanner.\n");
        return UNKNOWN_ERROR;
    }
    if (pid == 0) {
        int fd = open("/dev/null", O_WRONLY);
        if (fd < 0) {
            fprintf(stderr, "Unable to open /dev/null for media scanner stdout redirection.\n");
            exit(1);
        }
        dup2(fd, 1);
        int result = execvp(cmd[0], cmd);
        close(fd);
        exit(result);
    }
    wait(&status);

    if (status < 0) {
        fprintf(stderr, "Unable to broadcast intent for media scanner.\n");
        return UNKNOWN_ERROR;
    }
    return NO_ERROR;
}

status_t saveImage(const char* fn, std::optional<AndroidBitmapCompressFormat> format,
                   const ScreenCaptureResults& captureResults) {
    void* base = nullptr;
    ui::Dataspace dataspace = captureResults.capturedDataspace;
    const sp<GraphicBuffer>& buffer = captureResults.buffer;

    status_t result = buffer->lock(GraphicBuffer::USAGE_SW_READ_OFTEN, &base);

    if (base == nullptr || result != NO_ERROR) {
        String8 reason;
        if (result != NO_ERROR) {
            reason.appendFormat(" Error Code: %d", result);
        } else {
            reason = "Failed to write to buffer";
        }
        fprintf(stderr, "Failed to take screenshot (%s)\n", reason.c_str());
        return 1;
    }

    void* gainmapBase = nullptr;
    sp<GraphicBuffer> gainmap = captureResults.optionalGainMap;

    if (gainmap) {
        result = gainmap->lock(GraphicBuffer::USAGE_SW_READ_OFTEN, &gainmapBase);
        if (gainmapBase == nullptr || result != NO_ERROR) {
            fprintf(stderr, "Failed to capture gainmap with error code (%d)\n", result);
            gainmapBase = nullptr;
            // Fall-through: just don't attempt to write the gainmap
        }
    }

    int fd = -1;
    if (fn == nullptr) {
        fd = dup(STDOUT_FILENO);
        if (fd == -1) {
            fprintf(stderr, "Error writing to stdout. (%s)\n", strerror(errno));
            if (gainmapBase) {
                gainmap->unlock();
            }

            if (base) {
                buffer->unlock();
            }
            return 1;
        }
    } else {
        fd = open(fn, O_WRONLY | O_CREAT | O_TRUNC, 0664);
        if (fd == -1) {
            fprintf(stderr, "Error opening file: %s (%s)\n", fn, strerror(errno));
            if (gainmapBase) {
                gainmap->unlock();
            }

            if (base) {
                buffer->unlock();
            }
            return 1;
        }
    }

    if (format) {
        AndroidBitmapInfo info;
        info.format = flinger2bitmapFormat(
                static_cast<aidl::android::hardware::graphics::common::PixelFormat>(
                        buffer->getPixelFormat()));
        info.flags = ANDROID_BITMAP_FLAGS_ALPHA_PREMUL;
        info.width = buffer->getWidth();
        info.height = buffer->getHeight();
        info.stride = buffer->getStride() * bytesPerPixel(buffer->getPixelFormat());

        int bitmapResult =
                ABitmap_compressWithGainmap(&info, static_cast<ADataSpace>(dataspace), base,
                                            gainmapBase, captureResults.hdrSdrRatio, *format, 100,
                                            &fd,
                                            [](void* fdPtr, const void* data, size_t size) -> bool {
                                                int bytesWritten = write(*static_cast<int*>(fdPtr),
                                                                         data, size);
                                                return bytesWritten == size;
                                            });

        if (bitmapResult != ANDROID_BITMAP_RESULT_SUCCESS) {
            fprintf(stderr, "Failed to compress (error code: %d)\n", bitmapResult);
        }

        if (fn != NULL) {
            notifyMediaScanner(fn);
        }
    } else {
        uint32_t w = buffer->getWidth();
        uint32_t h = buffer->getHeight();
        uint32_t s = buffer->getStride();
        uint32_t f = buffer->getPixelFormat();
        uint32_t c = dataSpaceToInt(dataspace);

        write(fd, &w, 4);
        write(fd, &h, 4);
        write(fd, &f, 4);
        write(fd, &c, 4);
        size_t Bpp = bytesPerPixel(f);
        for (size_t y = 0; y < h; y++) {
            write(fd, base, w * Bpp);
            base = (void*)((char*)base + s * Bpp);
        }
    }
    close(fd);

    if (gainmapBase) {
        gainmap->unlock();
    }

    if (base) {
        buffer->unlock();
    }

    return 0;
}

int main(int argc, char** argv) {
    const std::vector<PhysicalDisplayId> physicalDisplays =
            SurfaceComposerClient::getPhysicalDisplayIds();

    if (physicalDisplays.empty()) {
        fprintf(stderr, "Failed to get ID for any displays.\n");
        return 1;
    }
    std::optional<DisplayId> displayIdOpt;
    std::vector<DisplayId> displaysToCapture;
    gui::CaptureArgs captureArgs;
    const char* pname = argv[0];
    bool png = false;
    bool jpeg = false;
    bool all = false;
    int c;
    while ((c = getopt_long(argc, argv, "apjhd:", LONG_OPTIONS, nullptr)) != -1) {
        switch (c) {
            case 'p':
                png = true;
                break;
            case 'j':
                jpeg = true;
                break;
            case 'd': {
                errno = 0;
                char* end = nullptr;
                const uint64_t id = strtoull(optarg, &end, 10);
                if (!end || *end != '\0' || errno == ERANGE) {
                    fprintf(stderr, "Invalid display ID: Out of range [0, 2^64).\n");
                    return 1;
                }

                displayIdOpt = DisplayId::fromValue(id);
                if (!displayIdOpt) {
                    fprintf(stderr, "Invalid display ID: Incorrect encoding.\n");
                    return 1;
                }
                displaysToCapture.push_back(displayIdOpt.value());
                break;
            }
            case 'a': {
                all = true;
                break;
            }
            case '?':
            case 'h':
                if (physicalDisplays.size() >= 1) {
                    displayIdOpt = physicalDisplays.front();
                }
                usage(pname, displayIdOpt);
                return 1;
            case LongOpts::RequireDpu:
                // require-dpu implies --preserve-display-colors
                captureArgs.captureMode = android::gui::CaptureMode::RequireOptimized;
                captureArgs.pixelFormat = 0; // unknown
                [[fallthrough]];
            // require-dpu implies --preserve-display-colors
            case LongOpts::PreservedDisplayColors:
                captureArgs.preserveDisplayColors = true;
                break;
        }
    }

    argc -= optind;
    argv += optind;

    // We don't expect more than 2 arguments.
    if (argc >= 2) {
        if (physicalDisplays.size() >= 1) {
            usage(pname, physicalDisplays.front());
        } else {
            usage(pname, std::nullopt);
        }
        return 1;
    }

    std::string baseName;
    std::string suffix;

    if (argc == 1) {
        std::string_view filename = {argv[0]};
        if (filename.ends_with(".png")) {
            baseName = filename.substr(0, filename.size() - 4);
            suffix = ".png";
            png = true;
        } else if (filename.ends_with(".jpeg")) {
            baseName = filename.substr(0, filename.size() - 5);
            suffix = ".jpeg";
            jpeg = true;
        } else if (filename.ends_with(".jpg")) {
            baseName = filename.substr(0, filename.size() - 4);
            suffix = ".jpg";
            jpeg = true;
        } else {
            baseName = filename;
        }
    }

    if (all) {
        // Ignores -d if -a is given.
        displaysToCapture.clear();
        for (int i = 0; i < physicalDisplays.size(); i++) {
            displaysToCapture.push_back(physicalDisplays[i]);
        }
    }

    if (displaysToCapture.empty()) {
        displaysToCapture.push_back(physicalDisplays.front());
        if (physicalDisplays.size() > 1) {
            fprintf(stderr,
                    "[Warning] Multiple displays were found, but no display id was specified! "
                    "Defaulting to the first display found, however this default is not guaranteed "
                    "to be consistent across captures. A display id should be specified.\n");
            fprintf(stderr, "A display ID can be specified with the [-d display-id] option.\n");
            fprintf(stderr, "See \"dumpsys SurfaceFlinger --display-id\" for valid display IDs.\n");
        }
    }

    if (png && jpeg) {
        fprintf(stderr, "Ambiguous file type");
        return 1;
    }

    std::optional<AndroidBitmapCompressFormat> format = std::nullopt;

    if (png) {
        format = ANDROID_BITMAP_COMPRESS_FORMAT_PNG;
    } else if (jpeg) {
        format = ANDROID_BITMAP_COMPRESS_FORMAT_JPEG;
    }

    // setThreadPoolMaxThreadCount(0) actually tells the kernel it's
    // not allowed to spawn any additional threads, but we still spawn
    // a binder thread from userspace when we call startThreadPool().
    // See b/36066697 for rationale
    ProcessState::self()->setThreadPoolMaxThreadCount(0);
    ProcessState::self()->startThreadPool();

    std::vector<ScreenCaptureResults> results;
    const size_t numDisplays = displaysToCapture.size();
    for (int i = 0; i < numDisplays; i++) {
        // 1. Capture the screen
        auto captureResult = screencap::capture(displaysToCapture[i], captureArgs);
        if (!captureResult.ok()) {
            fprintf(stderr, "%sCapturing failed.\n", captureResult.error().message().c_str());
            return 1;
        }

        // 2. Save the capture result as an image.
        // When there's more than one file to capture, add the index as postfix.
        std::string filename;
        if (!baseName.empty()) {
            filename = baseName;
            if (numDisplays > 1) {
                filename += "_";
                filename += std::to_string(i);
            }
            filename += suffix;
        }
        const char* fn = nullptr;
        if (!filename.empty()) {
            fn = filename.c_str();
        }
        if (const status_t saveImageStatus = saveImage(fn, format, captureResult.value()) != 0) {
            fprintf(stderr, "Saving image failed.\n");
            return saveImageStatus;
        }
    }

    return 0;
}