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

#include <memory>
#include <variant>
#include <vector>

#include "include/core/SkSpan.h"
#include "rust/png/FFI.rs.h"
#include "src/codec/SkCodecPriv.h"
#include "src/codec/SkFrameHolder.h"
#include "src/codec/SkPngCodecBase.h"
#include "third_party/rust/cxx/v1/cxx.h"

class SkFrame;
class SkStream;
struct SkEncodedInfo;

// This class provides the Skia image decoding API (`SkCodec`) on top of:
// * The third-party `png` crate (PNG decompression and decoding implemented in
//   Rust)
// * Skia's `SkSwizzler` and `skcms_Transform` (pixel format and color space
//   transformations implemented in C++).
class SkPngRustCodec final : public SkPngCodecBase {
public:
    static std::unique_ptr<SkPngRustCodec> MakeFromStream(std::unique_ptr<SkStream>,
                                                          Result*,
                                                          SkPngChunkReader* = nullptr);

    // `public` to support `std::make_unique<SkPngRustCodec>(...)`.
    SkPngRustCodec(SkEncodedInfo&&,
                   std::unique_ptr<SkStream>,
                   rust::Box<rust_png::Reader>,
                   sk_sp<SkPngCompositeChunkReader>,
                   std::unique_ptr<SkStream> gainmapStream,
                   std::optional<SkGainmapInfo> gainmapInfo);

    ~SkPngRustCodec() override;

private:
    // Holds raw destination parameters captured in `startDecoding` before
    // sampler/swizzler options are finalized.
    struct UninitializedDstInfo {
        // Pointer to the start of the caller's destination pixel buffer.
        uint8_t* fDstPtr;

        // Size of a row (in bytes) in the full image. Based on `rowBytes`
        // passed to `onGetPixels` or `onStartIncrementalDecode`.
        size_t fDstRowStride;
    };

    // Holds validated and calculated destination spans after `initializeDecodingDst`
    // has run.
    struct DecodingDstInfo {
        // Size of a row (in bytes) in the full image.
        size_t fDstRowStride;

        // `fDst` is based on `fDstPtr` from `UninitializedDstInfo`. For interlaced
        // and non-interlaced images, `initializeDecodingDst` initializes `fDst`
        // to start at the (0,0) (top-left) pixel of the current frame (which may
        // be offset from `fDstPtr` if the current frame is a sub-rect of the full
        // image). After decoding a non-interlaced row this moves (by
        // `fDstRowStride`) to the next row.
        SkSpan<uint8_t> fDst;

        // Size of a row (in bytes) in the current frame (or subset/sampled output).
        size_t fDstRowSize;

        // Height (in rows) of the current frame (or subset/sampled output).
        int fDstRowCount;

        // Bytes per pixel of fDst.
        uint8_t fDstBytesPerPixel;
    };

    struct DecodingState {
        // Transitions from `UninitializedDstInfo` (set in `startDecoding`) ->
        // `DecodingDstInfo` (populated by `initializeDecodingDst` once sampling
        // and swizzler parameters are finalized).
        std::variant<UninitializedDstInfo, DecodingDstInfo> fDst;

        DecodingDstInfo& dst() { return std::get<DecodingDstInfo>(fDst); }

        // Intermediate buffer that holds color-transformed pixels that are
        // ready to be blended with the destination. Used only when this frame
        // uses `SkCodecAnimation::Blend::kSrcOver`. For interlaced images this
        // buffer holds the whole frame; otherwise it holds only a single row.
        //
        // This is also used in the case of subsets or sampling for interlaced
        // images. We use this buffer as a full sized encoded image, which we then
        // take the subset / sampled rows from.
        // TODO: Subsets of APNG not supported, but if we need to, we would need
        // a separate fInterlacedBuffer along with fPreblendBuffer.
        SkCodecPriv::BudgetedBuffer fPreblendBuffer;

        // Single-row scratch buffers used when decoding interlaced images (unless
        // decoding a subset or sampling, where the whole untransformed image is
        // decoded into `fPreblendBuffer` first):
        // * `fDecodedInterlacedFullWidthRow` (size `getEncodedRowBytes()`): holds
        //   the decoded interlaced row padded to full width (used by `read_row` in
        //   `incrementalDecode`, or to pad `srcRow` for `applyXformRow` in
        //   `expandDecodedInterlacedRow`).
        // * `fXformedInterlacedRow` (size `fDstRowSize`): holds the color-transformed
        //   interlaced row in `expandDecodedInterlacedRow` when `!canReadRow()`.
        SkCodecPriv::BudgetedBuffer fDecodedInterlacedFullWidthRow;
        SkCodecPriv::BudgetedBuffer fXformedInterlacedRow;

        int fFirstRow = 0;
        int fLastRow = 0;

        // Tracks the current row index of the source image being processed.
        // Needs to be persisted in DecodingState to support incremental decoding
        // across multiple calls without resetting progress.
        int fCurrentSourceRow = 0;
        // Number of destination rows written to `fDst` so far:
        // * Non-interlaced: incremented after each decoded row.
        // * Interlaced full-image: updated in `expandDecodedInterlacedRow` from
        //   `fInterlacedDecodedHeight`.
        // * Interlaced subset/sampled: updated when flushing `fPreblendBuffer` in
        //   `writeSubsetOrSampleToOutput`.
        int fRowsWrittenToOutput = 0;

        // The y offset for a subset in the encoded color type, not the dst color type.
        // Only used for interlaced images.
        size_t fYByteOffset = 0;

        // Height (in rows) initialized so far by `splat_last_interlaced_row` (+8
        // rows per Pass 1 scanline).
        int fInterlacedDecodedHeight = 0;
    };

    // Helper for validating parameters of `onGetPixels` and/or
    // `onStartIncrementalDecode`. If `kSuccess` is returned then
    // `decodingState` output parameter got populated.
    Result startDecoding(const SkImageInfo& dstInfo,
                         void* pixels,
                         size_t rowBytes,
                         const Options& options,
                         std::optional<DecodingState>& decodingState);

    // Helper for initializing the destination span (`fDst`), row size, and
    // preblend buffer once sampling/subsetting parameters are finalized.
    Result initializeDecodingDst(DecodingState& decodingState);

    // Helper for taking a decoded interlaced `srcRow`, applying color
    // transformations, and then expanding it into the `frame`.
    void expandDecodedInterlacedRow(DecodingState& decodingState,
                                    SkSpan<uint8_t> dstFrame,
                                    SkSpan<const uint8_t> srcRow,
                                    const DecodingDstInfo& decodingDst,
                                    bool xFormNeeded);

    // Helper for row-by-row decoding which is used from `onGetPixels` and/or
    // `onIncrementalDecode`.
    Result incrementalDecode(DecodingState& decodingState);
    // The same as incrementalDecode but uses `applyXFormRow()`. Should only
    // be used if this->canReadRows() is false.
    Result incrementalDecodeXForm(DecodingState& decodingState);

    // Extracts and/or samples the rows decoded so far in `fPreblendBuffer` (in
    // the encoded color type) into `fDst` using `applyXformRow()`, and records
    // the number of destination rows written in `fRowsWrittenToOutput`. Used
    // for interlaced subset/sampled decodes both on halt and at end-of-frame.
    void writeSubsetOrSampleToOutput(DecodingState& decodingState);

    // Populates `rowsDecoded` after a failed decode (`kIncompleteInput` or
    // `kErrorInInput`) with the number of destination rows written so far,
    // matching `SkPngCodec` so callers only fill the remaining rows.
    void reportRowsDecoded(Result result,
                           const DecodingState& decodingState,
                           int* rowsDecoded) const;

    // Helper for reading until the start of the next `fdAT` sequence.
    Result readToStartOfNextFrame();

    // Helper for seeking to the start of image data for the given frame.
    Result seekToStartOfFrame(int index);

    // The number of frames calculated based on 1) the presence, and 2) the
    // contents of an `acTL` chunk.  "raw" in the sense that it reports all the
    // frames, while `SkCodec::getFrameCount` and
    // `SkPngRustCodec::onGetFrameCount` only report frames for which we have
    // successfully populated `fFrameHolder` with frame info parsed from `IHDR`
    // and/or `fcTL` chunks.
    int getRawFrameCount() const;

    // Attempts to read through the input stream to parse the additional `fcTL`
    // chunks.
    Result parseAdditionalFrameInfos();

    // Determines whether or not we can read from the rust decoder directly into dst.
    bool canReadRow();
    void processUnknownChunks();
    bool isLastFrame();
    bool isSampling() const;
    // Helper to query whether the image supports sampling or subsetting
    // (only static full-canvas images are supported; APNG animations and
    // offset frames are unsupported).
    bool supportsSamplingOrSubsetting();
    // Leaves the stream just past `IEND` so that callers can keep reading it (see `Codec_end`).
    // Like libpng, partial decodes skip this and stop early. Errors are ignored: the scanlines and
    // the `IDAT` CRC are verified already, and libpng also tolerates a damaged tail (it checks the
    // `IEND` CRC as if it were ancillary). No-op unless the codec is built `for_android`.
    void finishStreamAfterLastFrame();

    // SkCodec overrides:
    Result onGetPixels(const SkImageInfo& dstInfo,
                       void* pixels,
                       size_t rowBytes,
                       const Options&,
                       int* rowsDecoded) override;
    bool onSupportsIncrementalDecode(const SkImageInfo&) override { return true; }
    Result onStartIncrementalDecode(const SkImageInfo& dstInfo,
                                    void* pixels,
                                    size_t rowBytes,
                                    const Options&) override;
    Result onIncrementalDecode(int* rowsDecoded) override;
    int onGetFrameCount() override;
    bool onGetFrameInfo(int, FrameInfo*) const override;
    int onGetRepetitionCount() override;
    IsAnimated onIsAnimated() override;
    const SkFrameHolder* getFrameHolder() const override;
    sk_sp<const SkData> getEncodedData() const override;
    SkSampler* getSampler(bool createIfNecessary) final;

    // SkPngCodecBase overrides:
    std::unique_ptr<SkCodec> onDecodeGainmap(std::unique_ptr<SkStream> stream,
                                             SkCodec::Result* result) override;
    std::optional<SkSpan<const PaletteColorEntry>> onTryGetPlteChunk() override;
    std::optional<SkSpan<const uint8_t>> onTryGetTrnsChunk() override;

    rust::Box<rust_png::Reader> fReader;

    // `-1` means that `IDAT` is not part of animation and wasn't skipped yet.
    int fFrameAtCurrentStreamPosition = -1;
    bool fStreamIsPositionedAtStartOfFrameData = false;
    const std::unique_ptr<SkStream> fPrivStream;
    // TODO(https://crbug.com/371060427): Once fast seeking is available, we can
    // remove the field that tracks the stream length.
    std::optional<size_t> fMaxStreamLengthSeenWhenParsingAdditionalFrameInfos;

    std::optional<DecodingState> fIncrementalDecodingState;

    class FrameHolder final : public SkFrameHolder {
    public:
        FrameHolder(int width, int height);
        ~FrameHolder() override;

        FrameHolder(const FrameHolder&) = delete;
        FrameHolder(FrameHolder&&) = delete;
        FrameHolder& operator=(const FrameHolder&) = delete;
        FrameHolder& operator=(FrameHolder&&) = delete;

        // Returning an `int` (rather than `size_t`) for easier interop with
        // other parts of the SkCodec API.
        int size() const;

        Result appendNewFrame(const rust_png::Reader& reader, const SkEncodedInfo& info);
        void markFrameAsFullyReceived(size_t index);
        bool getFrameInfo(int index, FrameInfo* info) const;

    private:
        class PngFrame;

        const SkFrame* onGetFrame(int unverifiedIndex) const override;
        Result setFrameInfoFromCurrentFctlChunk(const rust_png::Reader& reader,
                                                PngFrame* out_frame);

        std::vector<PngFrame> fFrames;
    };
    FrameHolder fFrameHolder;

    // Whether there may still be additional `fcTL` chunks to discover and parse.
    //
    // `true` if the stream hasn't been fully received (i.e. only
    // `kIncompleteInput` errors so far, no hard errors) and `fFrameHolder`
    // doesn't yet contain frame info for all `num_frames` declared in an `acTL`
    // chunk.
    bool fCanParseAdditionalFrameInfos = true;
};

#endif  // SkPngRustCodec_DEFINED