/*
 * Copyright 2024 Google LLC
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#include "src/codec/SkPngCodecBase.h"

#include <algorithm>
#include <cstddef>
#include <tuple>
#include <utility>

#include "include/codec/SkCodec.h"
#include "include/codec/SkEncodedImageFormat.h"
#include "include/core/SkAlphaType.h"
#include "include/core/SkColor.h"
#include "include/core/SkColorType.h"
#include "include/core/SkImageInfo.h"
#include "include/core/SkRect.h"
#include "include/core/SkSpan.h"
#include "include/core/SkStream.h"
#include "include/private/SkAssert.h"
#include "include/private/SkEncodedInfo.h"
#include "modules/skcms/skcms.h"
#include "src/codec/SkCodecPriv.h"
#include "src/codec/SkColorPalette.h"
#include "src/codec/SkSwizzler.h"
#include "src/core/SkMemset.h"
#include "src/core/SkSwizzlePriv.h"

namespace {

constexpr SkColorType kXformSrcColorType = kRGBA_8888_SkColorType;

inline bool needs_premul(SkAlphaType dstAT, SkEncodedInfo::Alpha encodedAlpha) {
    return kPremul_SkAlphaType == dstAT && SkEncodedInfo::kUnpremul_Alpha == encodedAlpha;
}

skcms_PixelFormat ToPixelFormat(const SkEncodedInfo& info) {
    // We use kRGB and kRGBA formats because color PNGs are always RGB or RGBA.
    if (16 == info.bitsPerComponent()) {
        if (SkEncodedInfo::kRGBA_Color == info.color()) {
            return skcms_PixelFormat_RGBA_16161616BE;
        } else if (SkEncodedInfo::kRGB_Color == info.color()) {
            return skcms_PixelFormat_RGB_161616BE;
        }
    } else if (SkEncodedInfo::kGray_Color == info.color()) {
        return skcms_PixelFormat_G_8;
    }

    return skcms_PixelFormat_RGBA_8888;
}

}  // namespace

SkPngCodecBase::~SkPngCodecBase() = default;

// static
bool SkPngCodecBase::IsPng(const void* buf, size_t len) {
    static constexpr uint8_t kPngSignature[] = {0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A};
    if (len < sizeof(kPngSignature)) {
        return false;
    }

    return memcmp(buf, kPngSignature, sizeof(kPngSignature)) == 0;
}

// static
bool SkPngCodecBase::isCompatibleColorProfileAndType(const SkCodecs::ColorProfile* profile,
                                                     SkEncodedInfo::Color color) {
    if (profile) {
        switch (profile->dataSpace()) {
            case SkCodecs::ColorProfile::DataSpace::kCMYK:
                return false;
            case SkCodecs::ColorProfile::DataSpace::kGray:
                if (SkEncodedInfo::kGray_Color != color &&
                    SkEncodedInfo::kGrayAlpha_Color != color) {
                    return false;
                }
                break;
            default:
                break;
        }
    }

    return true;
}

SkPngCodecBase::SkPngCodecBase(SkEncodedInfo&& encodedInfo,
                               std::unique_ptr<SkStream> stream,
                               SkEncodedOrigin origin,
                               sk_sp<SkPngCompositeChunkReader> chunkReader,
                               std::unique_ptr<SkStream> gainmapStream,
                               std::optional<SkGainmapInfo> gainmapInfo)
        : SkCodec(std::move(encodedInfo), ToPixelFormat(encodedInfo), std::move(stream), origin)
        , fPngChunkReader(std::move(chunkReader))
        , fGainmapStream(std::move(gainmapStream))
        , fGainmapInfo(gainmapInfo) {}

SkEncodedImageFormat SkPngCodecBase::onGetEncodedFormat() const {
    return SkEncodedImageFormat::kPNG;
}

SkCodec::Result SkPngCodecBase::initializeXforms(const SkImageInfo& dstInfo,
                                                 const Options& options,
                                                 int frameWidth) {
    if (frameWidth != dstInfo.width() && options.fSubset) {
        return kInvalidParameters;
    }
    fXformWidth = frameWidth;

    {
        uint32_t encodedBitsPerPixel = getEncodedInfo().bitsPerPixel();

        // We assume that `frameWidth` and the encoded format have already been
        // sanitized earlier (preventing overflow in the row bytes calculation).
        SkASSERT_RELEASE(0 < frameWidth);
        SkASSERT_RELEASE(frameWidth < 0xFFFFFF);
        SkASSERT_RELEASE(encodedBitsPerPixel < 128);

        fEncodedRowBytes = SkCodecPriv::ComputeRowBytes(frameWidth, encodedBitsPerPixel);
    }

    // Reset fSwizzler and this->colorXform().  We can't do this in onRewind() because the
    // interlaced scanline decoder may need to rewind.
    fSwizzler.reset(nullptr);

    // If skcms directly supports the encoded PNG format, we should skip format
    // conversion in the swizzler (or skip swizzling altogether).
    bool skipFormatConversion = false;
    switch (this->getEncodedInfo().color()) {
        case SkEncodedInfo::kRGB_Color:
            if (this->getEncodedInfo().bitsPerComponent() != 16) {
                break;
            }
            [[fallthrough]];
        case SkEncodedInfo::kRGBA_Color:
        case SkEncodedInfo::kGray_Color:
            skipFormatConversion = this->colorXform();
            break;
        default:
            break;
    }

    if (skipFormatConversion && !options.fSubset) {
        fXformMode = kColorOnly_XformMode;
    } else {
        if (SkEncodedInfo::kPalette_Color == this->getEncodedInfo().color()) {
            if (!this->createColorTable(dstInfo)) {
                return kInvalidInput;
            }
        }

        Result result =
                this->initializeSwizzler(dstInfo, options, skipFormatConversion, frameWidth);
        if (result != kSuccess) {
            return result;
        }
    }

    if (Result result = this->allocateStorage(dstInfo); result != kSuccess) {
        return result;
    }

    // We can't call `initializeXformParams` here, because `swizzleWidth` may
    // change *after* `onStartIncrementalDecode`
    // (`SkSampledCodec::sampledDecode` first [transitively] calls
    // `onStartIncrementalDecode` and *then* `SkSwizzler::onSetSampleX`).

    return kSuccess;
}

void SkPngCodecBase::initializeXformParams() {
    switch (fXformMode) {
        case kSwizzleOnly_XformMode:
        case kSwizzleColor_XformMode:
            fXformWidth = this->swizzler()->swizzleWidth();
            break;
        case kColorOnly_XformMode:
            SkASSERT_RELEASE(!fSwizzler);
            break;
    }

    {
        // We assume that `fXformWidth` and the destination format have already been
        // sanitized earlier (preventing overflow in the destination row bytes calculation).
        SkASSERT_RELEASE(0 < fXformWidth);
        SkASSERT_RELEASE(fXformWidth < 0xFFFFFF);

        fDstRowSize = SkCodecPriv::ComputeRowBytesBytesPerPixel(fXformWidth,
                                                                this->dstInfo().bytesPerPixel());
    }
}

SkCodec::Result SkPngCodecBase::allocateStorage(const SkImageInfo& dstInfo) {
    fStorage.reset();
    switch (fXformMode) {
        case kSwizzleOnly_XformMode:
            break;
        case kColorOnly_XformMode:
            // Intentional fall through.  A swizzler hasn't been created yet, but one will
            // be created later if we are sampling.  We'll go ahead and allocate
            // enough memory to swizzle if necessary.
        case kSwizzleColor_XformMode: {
            const int bitsPerPixel = this->getEncodedInfo().bitsPerPixel();

            // If we have more than 8-bits (per component) of precision, we will keep that
            // extra precision.  Otherwise, we will swizzle to RGBA_8888 before transforming.
            const size_t bytesPerPixel = (bitsPerPixel > 32) ? bitsPerPixel / 8 : 4;
            const size_t colorXformBytes = dstInfo.width() * bytesPerPixel;
            return fStorage.allocateFromBudget(this, colorXformBytes, /*zeroInit=*/false);
        }
    }
    return kSuccess;
}

SkCodec::Result SkPngCodecBase::initializeSwizzler(const SkImageInfo& dstInfo,
                                                   const Options& options,
                                                   bool skipFormatConversion,
                                                   int frameWidth) {
    SkImageInfo swizzlerInfo = dstInfo;
    Options swizzlerOptions = options;
    fXformMode = kSwizzleOnly_XformMode;
    if (this->colorXform() && this->xformOnDecode()) {
        if (SkEncodedInfo::kGray_Color == this->getEncodedInfo().color()) {
            swizzlerInfo = swizzlerInfo.makeColorType(kGray_8_SkColorType);
        } else {
            swizzlerInfo = swizzlerInfo.makeColorType(kXformSrcColorType);
        }
        if (kPremul_SkAlphaType == dstInfo.alphaType()) {
            swizzlerInfo = swizzlerInfo.makeAlphaType(kUnpremul_SkAlphaType);
        }

        fXformMode = kSwizzleColor_XformMode;

        // Here, we swizzle into temporary memory, which is not zero initialized.
        // FIXME (msarett):
        // Is this a problem?
        swizzlerOptions.fZeroInitialized = kNo_ZeroInitialized;
    }

    SkIRect frameRect = SkIRect::MakeWH(frameWidth, 1);
    const SkIRect* frameRectPtr = nullptr;
    if (options.fSubset) {
        SkASSERT_RELEASE(frameWidth == dstInfo.width());
    } else {
        frameRectPtr = &frameRect;
    }

    if (skipFormatConversion) {
        // We cannot skip format conversion when there is a color table.
        SkASSERT_RELEASE(!fColorTable);
        int srcBPP = 0;
        switch (this->getEncodedInfo().color()) {
            case SkEncodedInfo::kRGB_Color:
                SkASSERT_RELEASE(this->getEncodedInfo().bitsPerComponent() == 16);
                srcBPP = 6;
                break;
            case SkEncodedInfo::kRGBA_Color:
                srcBPP = this->getEncodedInfo().bitsPerComponent() / 2;
                break;
            case SkEncodedInfo::kGray_Color:
                srcBPP = 1;
                break;
            default:
                SkASSERT_RELEASE(false);
                break;
        }
        fSwizzler = SkSwizzler::MakeSimple(srcBPP, swizzlerInfo, swizzlerOptions, frameRectPtr);
    } else {
        const SkPMColor* colors = SkCodecPriv::GetColorPtr(fColorTable.get());
        fSwizzler = SkSwizzler::Make(
                this->getEncodedInfo(), colors, swizzlerInfo, swizzlerOptions, frameRectPtr);
    }

    return !!fSwizzler ? kSuccess : kUnimplemented;
}

SkSampler* SkPngCodecBase::getSampler(bool createIfNecessary) {
    if (fSwizzler || !createIfNecessary) {
        return fSwizzler.get();
    }

    // Ok to ignore `initializeSwizzler`'s result, because if it fails, then
    // `fSwizzler` will be `nullptr` and we want to return `nullptr` upon
    // failure.
    std::ignore = this->initializeSwizzler(
            this->dstInfo(), this->options(), true, this->dstInfo().width());

    return fSwizzler.get();
}

void SkPngCodecBase::applyXformRow(SkSpan<uint8_t> dstRow, SkSpan<const uint8_t> srcRow) {
    SkASSERT_RELEASE(dstRow.size() >= fDstRowSize);
    SkASSERT_RELEASE(srcRow.size() >= fEncodedRowBytes);
    applyXformRow(dstRow.data(), srcRow.data());
}

void SkPngCodecBase::applyXformRow(void* dstRow, const uint8_t* srcRow) {
    switch (fXformMode) {
        case kSwizzleOnly_XformMode:
            fSwizzler->swizzle(dstRow, srcRow);
            break;
        case kColorOnly_XformMode:
            this->applyColorXform(dstRow, srcRow, fXformWidth);
            break;
        case kSwizzleColor_XformMode:
            SkASSERT_RELEASE(fSwizzler->swizzleWidth() == fXformWidth);
            fSwizzler->swizzle(fStorage.data(), srcRow);
            this->applyColorXform(dstRow, fStorage.data(), fXformWidth);
            break;
    }
}

// Note: SkColorPalette claims to store SkPMColors, which is not necessarily the case here.
bool SkPngCodecBase::createColorTable(const SkImageInfo& dstInfo) {
    if (fDstInfoOfPreviousColorTableCreation.has_value() &&
        *fDstInfoOfPreviousColorTableCreation == dstInfo) {
        return !!fColorTable;
    }
    fColorTable.reset();
    fDstInfoOfPreviousColorTableCreation = dstInfo;

    std::optional<SkSpan<const PaletteColorEntry>> maybePlteChunk = this->onTryGetPlteChunk();
    if (!maybePlteChunk.has_value()) {
        return false;
    }
    const PaletteColorEntry* palette = maybePlteChunk->data();
    constexpr size_t kMaxCountOfPaletteEntries = 256;
    size_t numColors = maybePlteChunk->size();
    numColors = std::min(numColors, kMaxCountOfPaletteEntries);

    // Contents depend on tableColorType and our choice of if/when to premultiply:
    // { kPremul, kUnpremul, kOpaque } x { RGBA, BGRA }
    SkPMColor colorTable[kMaxCountOfPaletteEntries];
    SkColorType tableColorType = this->colorXform() ? kXformSrcColorType : dstInfo.colorType();

    std::optional<SkSpan<const uint8_t>> maybeTrnsChunk = this->onTryGetTrnsChunk();
    const uint8_t* alphas = nullptr;
    size_t numColorsWithAlpha = 0;
    if (maybeTrnsChunk.has_value()) {
        alphas = maybeTrnsChunk->data();
        numColorsWithAlpha = maybeTrnsChunk->size();
        numColorsWithAlpha = std::min(numColorsWithAlpha, numColors);
    }

    bool shouldApplyColorXformToColorTable = this->colorXform() && !this->xformOnDecode();
    if (alphas) {
        bool premultiply = !shouldApplyColorXformToColorTable &&
                           needs_premul(dstInfo.alphaType(), this->getEncodedInfo().alpha());

        // Choose which function to use to create the color table. If the final destination's
        // colortype is unpremultiplied, the color table will store unpremultiplied colors.
        SkCodecPriv::PackColorProc proc =
                SkCodecPriv::ChoosePackColorProc(premultiply, tableColorType);

        for (size_t i = 0; i < numColorsWithAlpha; i++) {
            // We don't have a function in SkOpts that combines a set of alphas with a set
            // of RGBs.  We could write one, but it's hardly worth it, given that this
            // is such a small fraction of the total decode time.
            colorTable[i] = proc(alphas[i], palette->red, palette->green, palette->blue);
            palette++;
        }
    }

    if (numColorsWithAlpha < numColors) {
        // The optimized code depends on a 3-byte png_color struct with the colors
        // in RGB order.  These checks make sure it is safe to use.
        static_assert(3 == sizeof(PaletteColorEntry));
        static_assert(offsetof(PaletteColorEntry, red) == 0);
        static_assert(offsetof(PaletteColorEntry, green) == 1);
        static_assert(offsetof(PaletteColorEntry, blue) == 2);

        if (SkCodecPriv::IsRGBA(tableColorType)) {
            SkOpts::RGB_to_RGB1(colorTable + numColorsWithAlpha,
                                (const uint8_t*)palette,
                                numColors - numColorsWithAlpha);
        } else {
            SkOpts::RGB_to_BGR1(colorTable + numColorsWithAlpha,
                                (const uint8_t*)palette,
                                numColors - numColorsWithAlpha);
        }
    }

    if (shouldApplyColorXformToColorTable) {
        this->applyColorXform(colorTable, colorTable, numColors);
    }

    // Pad the color table with the last color in the table (or black) in the case that
    // invalid pixel indices exceed the number of colors in the table.
    const size_t maxColors = static_cast<size_t>(1) << this->getEncodedInfo().bitsPerComponent();
    if (numColors < maxColors) {
        SkPMColor lastColor = numColors > 0 ? colorTable[numColors - 1] : SK_ColorBLACK;
        SkOpts::memset32(colorTable + numColors, lastColor, maxColors - numColors);
    }

    fColorTable.reset(new SkColorPalette(colorTable, maxColors));
    return true;
}

bool SkPngCodecBase::onGetGainmapCodec(SkGainmapInfo* info,
                                       std::unique_ptr<SkCodec>* gainmapCodec) {
    if (!fGainmapStream) {
        return false;
    }

    sk_sp<const SkData> data = fGainmapStream->getData();
    if (!data) {
        return false;
    }

    if (!SkPngCodecBase::IsPng(data->bytes(), data->size())) {
        return false;
    }

    // The gainmap information lives on the gainmap image itself, so we need to
    // create the gainmap codec first, then check if it has a metadata chunk.
    SkCodec::Result result;
    std::unique_ptr<SkCodec> codec = this->onDecodeGainmap(fGainmapStream->duplicate(), &result);

    if (!codec || result != SkCodec::Result::kSuccess) {
        return false;
    }

    bool hasInfo = codec->onGetGainmapInfo(info);

    if (hasInfo && gainmapCodec) {
        // The ISO gainmap payload does not contain the actual alterative image
        // primaries, so we need to query the ICC profile stored on the gainmap.
        if (info->fGainmapMathColorSpace) {
            const auto* colorProfile = codec->getEncodedInfo().colorProfile();
            if (colorProfile) {
                auto colorSpace = colorProfile->getExactColorSpace();
                if (colorSpace) {
                    info->fGainmapMathColorSpace = std::move(colorSpace);
                }
            }
        }

        *gainmapCodec = std::move(codec);
    }

    return hasInfo;
}

bool SkPngCodecBase::onGetGainmapInfo(SkGainmapInfo* info) {
    if (fGainmapInfo) {
        if (info) {
            *info = *fGainmapInfo;
        }
        return true;
    }

    return false;
}