/*
 * 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 <vector>

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

struct SkEncodedInfo;
class SkFrame;
class SkStream;
template <typename T> class SkSpan;

// 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:
    struct DecodingDstInfo {
        // `fDst` is based on `pixels` passed to `onGetPixels` or
        // `onStartIncrementalDecode`.  For interlaced and non-interlaced
        // images, `startDecoding` initializes `fDst` to start at the (0,0)
        // (top-left) pixel of the current frame (which may be offset from
        // `pixels` 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 full image.  Based on `rowBytes`
        // passed to `onGetPixels` or `onStartIncrementalDecode`.
        size_t fDstRowStride = 0;

        // Size of a row (in bytes) in the current frame.
        size_t fDstRowSize = 0;

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

    struct DecodingState {
        // The info and pixels we will be decoding into.
        DecodingDstInfo fDecodingDstInfo;

        // 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 for interlaced images. We use
        // this buffer as a full sized encoded image, which we then take the subset
        // from.
        // TODO: Subsets of APNG not supported, but if we need to, we would need
        // a separate fInterlacedBuffer along with fPreblendBuffer.
        std::vector<uint8_t> fPreblendBuffer;

        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;
        // Tracks the number of rows actually written to the destination buffer.
        // Used to correctly report progress via rowsDecoded in incrementalDecode.
        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;
    };

    // 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,
                         DecodingState* decodingState);

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

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

    // 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;
    Result initializeSamplerParams(DecodingState& decodingState);

    // 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;

    // When decoding interlaced subsets or sampling, decode the full image into a
    // buffer in the encoded colortype and extract/sample for the final dst using
    // this function. This function will use applyXFormRow() for needed rows.
    void getSubsetOrSampleFromFullImage(SkSpan<const uint8_t> fullImageBuffer,
                                        SkSpan<uint8_t> dst,
                                        size_t dstRowStride,
                                        size_t offset,
                                        int srcHeight,
                                        int dstHeight,
                                        size_t dstRowSize);
};

#endif  // SkPngRustCodec_DEFINED