/*
 * Copyright (c) 2026 Huawei Device Co., Ltd.
 * 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 "ApngDecoder.h"
#include <memory>
#include <cmath>
#include <cstring>
#include <iostream>
#include "png.h"
#ifndef USE_LIBPNG_APNG
#include "apng.h" // 仅在使用预编译 libpng 时需 stub 声明
#endif
#include "Error.h"
#include "StreamSourceBase.h"
#include "StreamSourceHarmony.h"
#include "Log.h"

namespace apng_drawable {
const uint8_t ALPHA_TRANSPARENT = 0U;
const uint8_t ALPHA_OPAQUE = 0xFFU;
constexpr size_t CHUNK_NAME_SIZE = 4;
constexpr uint8_t CHANNEL_COUNT_RGBA = 4;
// APNG chunk 数量 (acTL, fcTL, fdAT)
constexpr int APNG_CHUNK_COUNT = 3;

namespace {

struct PngContext {
    png_structp pngPtr = nullptr;
    png_infop infoPtr = nullptr;
    std::unique_ptr<StreamSourceBase> source;
};

struct FrameBuffers {
    std::unique_ptr<uint8_t[]> current;      // p_frame (最终累积的帧)
    std::unique_ptr<uint8_t[]> scratch;      // p_buffer (当前解码出的原始帧)
    std::unique_ptr<uint8_t[]> previous;     // p_previous_frame (用于 dispose previous)
    std::unique_ptr<png_bytep[]> rows_current;
    std::unique_ptr<png_bytep[]> rows_scratch;
    size_t size = 0;          // 当前缓冲区总字节数
    size_t rowBytes = 0;     // 每行字节数
    uint32_t height = 0;

    bool Allocate(uint32_t h, size_t rb)
    {
        height = h;
        rowBytes = rb;
        if (rowBytes == 0 || height > SIZE_MAX / rowBytes) {
            return false;
        }
        size = height * rowBytes;
        constexpr size_t maxAllocBytes = 256 * 1024 * 104;
        if (size > maxAllocBytes) {
            return false;
        }
        current.reset(new (std::nothrow) uint8_t[size]);
        scratch.reset(new (std::nothrow) uint8_t[size]);
        previous.reset(new (std::nothrow) uint8_t[size]);
        rows_current.reset(new (std::nothrow) png_bytep[height]);
        rows_scratch.reset(new (std::nothrow) png_bytep[height]);
        if (!current || !scratch || !previous || !rows_current || !rows_scratch) {
            return false;
        }
        for (uint32_t j = 0; j < height; ++j) {
            rows_current[j] = current.get() + j * rowBytes;
            rows_scratch[j] = scratch.get() + j * rowBytes;
        }
        return true;
    }

    void UpdateForFrame(uint32_t frameH, size_t frameRb)
    {
        // 如果尺寸变化，重新分配 scratch 和 rows_scratch
        if (frameH != height || frameRb != rowBytes) {
            size_t newSize = frameH * frameRb;
            if (newSize > size) {
                scratch.reset(new uint8_t[newSize]());
                size = newSize;
            }
            if (frameH > 0) {
                rows_scratch.reset(new png_bytep[frameH]);
            }
            for (uint32_t j = 0; j < frameH; ++j) {
                rows_scratch[j] = scratch.get() + j * frameRb;
            }
        }
    }

    void ResetRowsForCurrent()
    {
        for (uint32_t j = 0; j < height; ++j) {
            rows_current[j] = current.get() + j * rowBytes;
        }
    }
};

void SaveFrame(uint32_t *destination, uint32_t ** const source, uint32_t const width, uint32_t const height)
{
    if (!destination) {
        return;
    }
    uint_fast8_t alpha;
    uint32_t srcColor;
    for (uint32_t j = 0; j < height; ++j, destination += width) {
        std::copy_n(source[j], width * CHANNEL_4_BYTE_SIZE, destination);
        for (uint32_t i = 0; i < width; ++i) {
            srcColor = destination[i];
            alpha = static_cast<uint_fast8_t>((srcColor >> 24U) & 0xFFU);
            if (alpha == ALPHA_TRANSPARENT) {
                destination[i] = 0;
            }
        }
    }
}

static void BlendPixel(uint8_t* dp, const uint8_t* sp)
{
    uint8_t sourceAlpha = sp[CHANNEL_INDEX_A];
    if (sourceAlpha == ALPHA_OPAQUE) {
        // 源完全不透明，直接覆盖目标
        std::copy_n(sp, CHANNEL_4_BYTE_SIZE, dp);
    } else if (sourceAlpha != ALPHA_TRANSPARENT) {
        // 源半透明，需要与目标混合
        if (dp[CHANNEL_INDEX_A] != ALPHA_TRANSPARENT) {
            int32_t u = sourceAlpha * ALPHA_OPAQUE;
            int32_t v = (ALPHA_OPAQUE - sourceAlpha) * dp[CHANNEL_INDEX_A];
            int32_t al = ALPHA_OPAQUE * ALPHA_OPAQUE -
                         (ALPHA_OPAQUE - sourceAlpha) * (ALPHA_OPAQUE - dp[CHANNEL_INDEX_A]);
            dp[CHANNEL_INDEX_R] = static_cast<uint8_t>(
                                            (sp[CHANNEL_INDEX_R] * u + dp[CHANNEL_INDEX_R] * v) / al);
            dp[CHANNEL_INDEX_G] = static_cast<uint8_t>(
                                            (sp[CHANNEL_INDEX_G] * u + dp[CHANNEL_INDEX_G] * v) / al);
            dp[CHANNEL_INDEX_B] = static_cast<uint8_t>(
                                            (sp[CHANNEL_INDEX_B] * u + dp[CHANNEL_INDEX_B] * v) / al);
            dp[CHANNEL_INDEX_A] = static_cast<uint8_t>(al / ALPHA_OPAQUE);
        } else {
            // 目标完全透明，直接覆盖
            std::copy_n(sp, CHANNEL_4_BYTE_SIZE, dp);
        }
    }
}

struct BlendRowPlanes {
    uint8_t** destination;
    uint8_t** source;
};

struct BlendRegion {
    png_uint_32 xOffset;
    png_uint_32 yOffset;
    png_uint_32 width;
    png_uint_32 height;
};

void BlendOver(const BlendRowPlanes& planes, const BlendRegion& region)
{
    for (uint32_t j = 0; j < region.height; ++j) {
        uint8_t* sp = planes.source[j];
        uint8_t* dp = planes.destination[j + region.yOffset] + region.xOffset * CHANNEL_4_BYTE_SIZE;
        for (uint32_t i = 0; i < region.width; ++i) {
            BlendPixel(dp, sp);
            sp += CHANNEL_4_BYTE_SIZE;
            dp += CHANNEL_4_BYTE_SIZE;
        }
    }
}

inline void BlendSource(const BlendRowPlanes& planes, const BlendRegion& region)
{
    for (uint32_t j = 0; j < region.height; j++) {
        std::copy_n(planes.source[j], region.width * CHANNEL_4_BYTE_SIZE,
            planes.destination[j + region.yOffset] + region.xOffset * CHANNEL_4_BYTE_SIZE);
    }
}

struct PngInitializedInfo {
    uint32_t width = 0;
    uint32_t height = 0;
    size_t rowBytes = 0;
    unsigned int channels = 0;
};

bool InitPngRead(PngContext& ctx, PngInitializedInfo& info, int32_t& result)
{
    ctx.pngPtr = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    ctx.infoPtr = png_create_info_struct(ctx.pngPtr);
    if (!ctx.pngPtr || !ctx.infoPtr) {
        png_destroy_read_struct(&ctx.pngPtr, &ctx.infoPtr, nullptr);
        result = ERR_OUT_OF_MEMORY;
        return false;
    }
    if (setjmp(png_jmpbuf(ctx.pngPtr)) != 0) {
        result = ctx.source->GetError();
        if (!result) {
            result = ERR_INVALID_FILE_FORMAT;
        }
        OH_LOG_Print(LOG_APP, LOG_ERROR, APNG_LOG_DOMAIN, "ApngDecoder",
                     "initPngRead: PNG read error, code=%d", result);
        png_destroy_read_struct(&ctx.pngPtr, &ctx.infoPtr, nullptr);
        return false;
    }
    ctx.source->Init(ctx.pngPtr);
    static_cast<StreamSourceHarmony*>(ctx.source.get())->SetOffset(PNG_SIG_SIZE);
    png_set_sig_bytes(ctx.pngPtr, PNG_SIG_SIZE);
#ifndef USE_LIBPNG_APNG
    png_byte apngChunks[] = {'a', 'c', 'T', 'L', 0, 'f', 'c', 'T', 'L', 0, 'f', 'd', 'A', 'T', 0};
    png_set_keep_unknown_chunks(ctx.pngPtr, PNG_HANDLE_CHUNK_ALWAYS, apngChunks, APNG_CHUNK_COUNT);
#endif
    png_read_info(ctx.pngPtr, ctx.infoPtr);
    png_set_expand(ctx.pngPtr);
    png_set_strip_16(ctx.pngPtr);
    png_set_gray_to_rgb(ctx.pngPtr);
    png_set_add_alpha(ctx.pngPtr, ALPHA_OPAQUE, PNG_FILLER_AFTER);
    png_set_interlace_handling(ctx.pngPtr);
    png_read_update_info(ctx.pngPtr, ctx.infoPtr);
    info.width = static_cast<uint32_t>(png_get_image_width(ctx.pngPtr, ctx.infoPtr));
    info.height = static_cast<uint32_t>(png_get_image_height(ctx.pngPtr, ctx.infoPtr));
    info.channels = png_get_channels(ctx.pngPtr, ctx.infoPtr);
    info.rowBytes = png_get_rowbytes(ctx.pngPtr, ctx.infoPtr);
    if (info.width == 0 || info.height == 0 || info.channels != CHANNEL_COUNT_RGBA) {
        result = ERR_INVALID_FILE_FORMAT;
        png_destroy_read_struct(&ctx.pngPtr, &ctx.infoPtr, nullptr);
        return false;
    }
    return true;
}

std::unique_ptr<ApngImage> CreateApngImage(uint32_t width,
                                           uint32_t height,
                                           uint32_t frames,
                                           uint32_t plays,
                                           int32_t& result)
{
    try {
        return std::make_unique<ApngImage>(width, height, frames, plays);
    } catch (const std::bad_alloc&) {
        result = ERR_OUT_OF_MEMORY;
        return nullptr;
    }
}

bool DecodeOneFrame(PngContext& ctx, FrameBuffers& bufs, uint32_t frameIdx,
                    uint32_t firstHidden, ApngImage* png, int32_t& result)
{
    png_read_frame_head(ctx.pngPtr, ctx.infoPtr);
    png_uint_32 frameW = 0;
    png_uint_32 frameH = 0;
    png_uint_32 xOff = 0;
    png_uint_32 yOff = 0;
    png_uint_16 delayNum = 1;
    png_uint_16 delayDen = 100;
    png_byte disposeOp = 0;
    png_byte blendOp = 0;
    png_get_next_frame_fcTL(ctx.pngPtr, ctx.infoPtr, &frameW, &frameH,
                            &xOff, &yOff, &delayNum, &delayDen,
                            &disposeOp, &blendOp);
    size_t frameRowBytes = png_get_rowbytes(ctx.pngPtr, ctx.infoPtr);
    bufs.UpdateForFrame(frameH, frameRowBytes);
    auto duration = static_cast<size_t>(std::lround(
        static_cast<float>(delayNum) / delayDen * 1000.0f));
    size_t pixelCount = static_cast<size_t>(png->GetWidth()) * png->GetHeight();
    auto frame = std::make_unique<ApngFrame>(pixelCount, duration);
    // 首帧特殊处理
    if (frameIdx == firstHidden) {
        blendOp = PNG_BLEND_OP_SOURCE;
        if (disposeOp == PNG_DISPOSE_OP_PREVIOUS) {
            disposeOp = PNG_DISPOSE_OP_BACKGROUND;
        }
    }
    png_read_image(ctx.pngPtr, bufs.rows_scratch.get());
    // 处理 dispose previous
    if (disposeOp == PNG_DISPOSE_OP_PREVIOUS) {
        std::copy_n(bufs.previous.get(), bufs.size, bufs.current.get());
    }
    // 混合
    if (blendOp == PNG_BLEND_OP_OVER) {
        BlendOver({ bufs.rows_current.get(), bufs.rows_scratch.get() }, { xOff, yOff, frameW, frameH });
    } else { // PNG_BLEND_OP_SOURCE
        BlendSource({ bufs.rows_current.get(), bufs.rows_scratch.get() }, { xOff, yOff, frameW, frameH });
    }
    // 保存到 ApngFrame
    SaveFrame(frame->GetRawPixels(),
              reinterpret_cast<uint32_t**>(bufs.rows_current.get()),
              png->GetWidth(), png->GetHeight());
    // 更新 previous 或清除背景
    if (disposeOp == PNG_DISPOSE_OP_PREVIOUS) {
        std::copy_n(bufs.current.get(), bufs.size, bufs.previous.get());
    } else if (disposeOp == PNG_DISPOSE_OP_BACKGROUND) {
        for (uint32_t j = 0; j < frameH; ++j) {
            std::fill_n(bufs.rows_current[j + yOff] + xOff * CHANNEL_COUNT_RGBA, frameW * CHANNEL_COUNT_RGBA, 0);
        }
    }
    png->SetFrame(frameIdx, std::move(frame));
    return true;
}

bool ReadActlChunk(PngContext& ctx, png_uint_32& frames, png_uint_32& plays, int32_t& result)
{
    if (!png_get_acTL(ctx.pngPtr, ctx.infoPtr, &frames, &plays)) {
        result = ERR_INVALID_FILE_FORMAT;
        png_destroy_read_struct(&ctx.pngPtr, &ctx.infoPtr, nullptr);
        return false;
    }
    return true;
}

bool AllocateFrameBuffers(PngContext& ctx, FrameBuffers& bufs, uint32_t height, size_t rowBytes, int32_t& result)
{
    if (!bufs.Allocate(height, rowBytes)) {
        result = ERR_OUT_OF_MEMORY;
        png_destroy_read_struct(&ctx.pngPtr, &ctx.infoPtr, nullptr);
        return false;
    }
    return true;
}

struct ApngImageCreateParams {
    uint32_t width;
    uint32_t height;
    uint32_t frames;
    uint32_t plays;
};

std::unique_ptr<ApngImage> CreateApngImageOrFail(PngContext& ctx, const ApngImageCreateParams& spec, int32_t& result)
{
    auto png = CreateApngImage(spec.width, spec.height, spec.frames, spec.plays, result);
    if (!png) {
        png_destroy_read_struct(&ctx.pngPtr, &ctx.infoPtr, nullptr);
    }
    return png;
}

bool BeginFrameDecoding(PngContext& ctx, int32_t& result)
{
    if (setjmp(png_jmpbuf(ctx.pngPtr)) != 0) {
        result = ctx.source->GetError();
        if (!result) {
            result = ERR_INVALID_FILE_FORMAT;
        }
        OH_LOG_Print(LOG_APP, LOG_ERROR, APNG_LOG_DOMAIN, "ApngDecoder",
                     "Decode: Frame read error, code=%d", result);
        png_destroy_read_struct(&ctx.pngPtr, &ctx.infoPtr, nullptr);
        return false;
    }
    return true;
}

bool DecodeFrames(PngContext& ctx, FrameBuffers& bufs, ApngImage* png, uint32_t decodeFrameCount, int32_t& result)
{
    uint32_t firstHidden = (png_get_first_frame_is_hidden(ctx.pngPtr, ctx.infoPtr) != 0) ? 1 : 0;
    for (uint32_t i = 0; i < decodeFrameCount; ++i) {
        if (!DecodeOneFrame(ctx, bufs, i, firstHidden, png, result)) {
            png_destroy_read_struct(&ctx.pngPtr, &ctx.infoPtr, nullptr);
            return false;
        }
    }
    return true;
}
} // anonymous namespace

// 主解码函数
std::unique_ptr<ApngImage> ApngDecoder::Decode(std::unique_ptr<StreamSourceBase> source,
                                               int32_t& result)
{
    if (source->CheckPngSignature() < SUCCESS) {
        result = source->CheckPngSignature();
        return nullptr;
    }
    PngContext ctx{nullptr, nullptr, std::move(source)};
    PngInitializedInfo pngInfo;
    if (!InitPngRead(ctx, pngInfo, result)) {
        return nullptr;
    }
    png_uint_32 frames = 1;
    png_uint_32 plays = 0;
    if (!ReadActlChunk(ctx, frames, plays, result)) {
        return nullptr;
    }
#ifdef USE_LIBPNG_APNG
    uint32_t decodeFrameCount = frames;
#else
    uint32_t decodeFrameCount = 1;
#endif
    FrameBuffers bufs;
    if (!AllocateFrameBuffers(ctx, bufs, pngInfo.height, pngInfo.rowBytes, result)) {
        return nullptr;
    }
    auto png = CreateApngImageOrFail(ctx,
        { pngInfo.width, pngInfo.height, static_cast<uint32_t>(frames), static_cast<uint32_t>(plays) }, result);
    if (!png) {
        return nullptr;
    }
    if (!BeginFrameDecoding(ctx, result)) {
        return nullptr;
    }
    if (!DecodeFrames(ctx, bufs, png.get(), decodeFrameCount, result)) {
        return nullptr;
    }
    png_read_end(ctx.pngPtr, ctx.infoPtr);
    png_destroy_read_struct(&ctx.pngPtr, &ctx.infoPtr, nullptr);
    result = SUCCESS;
    return png;
}

bool ApngDecoder::IsApng(std::unique_ptr<StreamSourceBase> source)
{
    // Checks PNG signature
    int result = source->CheckPngSignature();
    if (result != SUCCESS) {
        OH_LOG_Print(LOG_APP, LOG_DEBUG, APNG_LOG_DOMAIN, "ApngDecoder",
            "IsApng: PNG signature check failed, result=%{public}d", result);
        return false;
    }
    OH_LOG_Print(LOG_APP, LOG_DEBUG, APNG_LOG_DOMAIN, "ApngDecoder", "IsApng: PNG signature check passed");

    // Checks APNG acTL chunk
    png_structp pngPtr = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    png_infop infoPtr = png_create_info_struct(pngPtr);
    if (!pngPtr || !infoPtr) {
        OH_LOG_Print(LOG_APP, LOG_ERROR, APNG_LOG_DOMAIN, "ApngDecoder",
            "IsApng: Failed to create png_struct or png_info");
        png_destroy_read_struct(&pngPtr, &infoPtr, nullptr);
        return false;
    }

    // Point to handle error (Read header and acTL)
    if (setjmp(png_jmpbuf(pngPtr)) != 0) { // NOLINT(cert-err52-cpp)
        OH_LOG_Print(LOG_APP, LOG_ERROR, APNG_LOG_DOMAIN, "ApngDecoder", "IsApng: Error during PNG reading (setjmp)");
        png_destroy_read_struct(&pngPtr, &infoPtr, nullptr);
        return false;
    }

    // Read header
    source->Init(pngPtr);
    // Set offset to skip PNG signature (8 bytes) since we already checked it
    // After png_set_sig_bytes, libpng expects data to start from offset 8
    static_cast<StreamSourceHarmony *>(source.get())->SetOffset(PNG_SIG_SIZE);
    png_set_sig_bytes(pngPtr, PNG_SIG_SIZE);

#ifndef USE_LIBPNG_APNG
    // 预编译 libpng 无 APNG 支持，需保留 acTL/fcTL/fdAT 为 unknown 供 stub 解析
    png_byte apngChunks[] = {'a', 'c', 'T', 'L', 0, 'f', 'c', 'T', 'L', 0, 'f', 'd', 'A', 'T', 0};
    png_set_keep_unknown_chunks(pngPtr, PNG_HANDLE_CHUNK_ALWAYS, apngChunks, APNG_CHUNK_COUNT);
#endif

    png_read_info(pngPtr, infoPtr);
    OH_LOG_Print(LOG_APP, LOG_DEBUG, APNG_LOG_DOMAIN, "ApngDecoder", "IsApng: png_read_info completed");

    // Check unknown chunks count
    png_unknown_chunkp unknowns;
    int numUnknowns = png_get_unknown_chunks(pngPtr, infoPtr, &unknowns);
    OH_LOG_Print(LOG_APP, LOG_DEBUG, APNG_LOG_DOMAIN, "ApngDecoder", "IsApng: Found %{public}d unknown chunks",
        numUnknowns);
    for (int i = 0; i < numUnknowns && i < MAX_LOG_UNKNOWN_CHUNKS; i++) { // Limit to 10 chunks for logging
        char chunkName[5] = {0};
        std::copy_n(unknowns[i].name, CHUNK_NAME_SIZE, chunkName);
        OH_LOG_Print(LOG_APP, LOG_DEBUG, APNG_LOG_DOMAIN, "ApngDecoder",
            "IsApng: Unknown chunk[%{public}d]: %{public}s, size=%{public}u", i, chunkName,
            static_cast<unsigned int>(unknowns[i].size));
    }

    // Read acTL
    png_uint_32 frames = 0;
    png_uint_32 plays = 0;
    bool hasAcTL = png_get_acTL(pngPtr, infoPtr, &frames, &plays) != 0;
    OH_LOG_Print(LOG_APP, LOG_INFO, APNG_LOG_DOMAIN, "ApngDecoder",
        "IsApng: png_get_acTL result=%{public}d, frames=%{public}u, plays=%{public}u", hasAcTL ? 1 : 0, frames, plays);
    png_destroy_read_struct(&pngPtr, &infoPtr, nullptr);
    return hasAcTL;
}
}