/*
 * Copyright 2017, 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.
 */

#ifndef CODEC2_BUFFER_H_

#define CODEC2_BUFFER_H_

#include <C2Buffer.h>
#include <C2Config.h>

#include <binder/IMemory.h>
#include <media/hardware/VideoAPI.h>
#include <media/stagefright/foundation/ABuffer.h>
#include <media/MediaCodecBuffer.h>

namespace android {

namespace hardware {
class HidlMemory;
namespace cas {
namespace native {
namespace V1_0 {
struct SharedBuffer;
}  // namespace V1_0
}  // namespace native
}  // namespace cas
namespace drm {
namespace V1_0 {
struct SharedBuffer;
}  // namespace V1_0
}  // namespace drm
}  // namespace hardware

class Codec2Buffer : public MediaCodecBuffer {
public:
    using MediaCodecBuffer::MediaCodecBuffer;
    ~Codec2Buffer() override = default;

    sp<ABuffer> getImageData() const { return mImageData; }

    virtual void clearC2BufferRefs() {}

protected:
    /**
     * canCopy() implementation for linear buffers.
     */
    bool canCopyLinear(const std::shared_ptr<C2Buffer> &buffer) const;

    /**
     * copy() implementation for linear buffers.
     */
    bool copyLinear(const std::shared_ptr<C2Buffer> &buffer);

    /**
     * sets MediaImage data for flexible graphic buffers
     */
    void setImageData(const sp<ABuffer> &imageData);

    sp<ABuffer> mImageData;
};

/**
 * MediaCodecBuffer implementation on top of local linear buffer. This cannot
 * cross process boundary so asC2Buffer() returns only nullptr.
 */
class LocalLinearBuffer : public Codec2Buffer {
public:
    using Codec2Buffer::Codec2Buffer;

    std::shared_ptr<C2Buffer> asC2Buffer() override { return nullptr; }
    bool canCopy(const std::shared_ptr<C2Buffer> &buffer) const override;
    bool copy(const std::shared_ptr<C2Buffer> &buffer) override;
};

/**
 * MediaCodecBuffer implementation to be used only as a dummy wrapper around a
 * C2Buffer object.
 */
class DummyContainerBuffer : public Codec2Buffer {
public:
    DummyContainerBuffer(
            const sp<AMessage> &format,
            const std::shared_ptr<C2Buffer> &buffer = nullptr);

    std::shared_ptr<C2Buffer> asC2Buffer() override;
    void clearC2BufferRefs() override;
    bool canCopy(const std::shared_ptr<C2Buffer> &buffer) const override;
    bool copy(const std::shared_ptr<C2Buffer> &buffer) override;

private:
    std::shared_ptr<C2Buffer> mBufferRef;
};

/**
 * MediaCodecBuffer implementation wraps around C2LinearBlock.
 */
class LinearBlockBuffer : public Codec2Buffer {
public:
    /**
     * Allocate a new LinearBufferBlock wrapping around C2LinearBlock object.
     *
     * \param   format  mandatory buffer format for MediaCodecBuffer
     * \param   block   C2LinearBlock object to wrap around.
     * \return          LinearBlockBuffer object with writable mapping.
     *                  nullptr if unsuccessful.
     */
    static sp<LinearBlockBuffer> Allocate(
            const sp<AMessage> &format, const std::shared_ptr<C2LinearBlock> &block);

    virtual ~LinearBlockBuffer() = default;

    std::shared_ptr<C2Buffer> asC2Buffer() override;
    bool canCopy(const std::shared_ptr<C2Buffer> &buffer) const override;
    bool copy(const std::shared_ptr<C2Buffer> &buffer) override;

private:
    LinearBlockBuffer(
            const sp<AMessage> &format,
            C2WriteView &&writeView,
            const std::shared_ptr<C2LinearBlock> &block);
    LinearBlockBuffer() = delete;

    C2WriteView mWriteView;
    std::shared_ptr<C2LinearBlock> mBlock;
};

/**
 * MediaCodecBuffer implementation wraps around C2ConstLinearBlock.
 */
class ConstLinearBlockBuffer : public Codec2Buffer {
public:
    /**
     * Allocate a new ConstLinearBlockBuffer wrapping around C2Buffer object.
     *
     * \param   format  mandatory buffer format for MediaCodecBuffer
     * \param   buffer  linear C2Buffer object to wrap around.
     * \return          ConstLinearBlockBuffer object with readable mapping.
     *                  nullptr if unsuccessful.
     */
    static sp<ConstLinearBlockBuffer> Allocate(
            const sp<AMessage> &format, const std::shared_ptr<C2Buffer> &buffer);

    virtual ~ConstLinearBlockBuffer() = default;

    std::shared_ptr<C2Buffer> asC2Buffer() override;
    void clearC2BufferRefs() override;

private:
    ConstLinearBlockBuffer(
            const sp<AMessage> &format,
            C2ReadView &&readView,
            const std::shared_ptr<C2Buffer> &buffer);
    ConstLinearBlockBuffer() = delete;

    C2ReadView mReadView;
    std::shared_ptr<C2Buffer> mBufferRef;
};

/**
 * MediaCodecBuffer implementation wraps around C2GraphicBlock.
 *
 * This object exposes the underlying bits via accessor APIs and "image-data"
 * metadata, created automatically at allocation time.
 */
class GraphicBlockBuffer : public Codec2Buffer {
public:
    /**
     * Allocate a new GraphicBlockBuffer wrapping around C2GraphicBlock object.
     * If |block| is not in good color formats, it allocates YV12 local buffer
     * and copies the content over at asC2Buffer().
     *
     * \param   format  mandatory buffer format for MediaCodecBuffer
     * \param   block   C2GraphicBlock object to wrap around.
     * \param   alloc   a function to allocate backing ABuffer if needed.
     * \return          GraphicBlockBuffer object with writable mapping.
     *                  nullptr if unsuccessful.
     */
    static sp<GraphicBlockBuffer> Allocate(
            const sp<AMessage> &format,
            const std::shared_ptr<C2GraphicBlock> &block,
            std::function<sp<ABuffer>(size_t)> alloc);

    virtual ~GraphicBlockBuffer() = default;

    std::shared_ptr<C2Buffer> asC2Buffer() override;

private:
    GraphicBlockBuffer(
            const sp<AMessage> &format,
            const sp<ABuffer> &buffer,
            C2GraphicView &&view,
            const std::shared_ptr<C2GraphicBlock> &block,
            const sp<ABuffer> &imageData,
            bool wrapped);
    GraphicBlockBuffer() = delete;

    inline MediaImage2 *imageData() { return (MediaImage2 *)mImageData->data(); }

    C2GraphicView mView;
    std::shared_ptr<C2GraphicBlock> mBlock;
    const bool mWrapped;
};

/**
 * MediaCodecBuffer implementation wraps around VideoNativeMetadata.
 */
class GraphicMetadataBuffer : public Codec2Buffer {
public:
    /**
     * Construct a new GraphicMetadataBuffer with local linear buffer for
     * VideoNativeMetadata.
     *
     * \param   format      mandatory buffer format for MediaCodecBuffer
     */
    GraphicMetadataBuffer(
            const sp<AMessage> &format, const std::shared_ptr<C2Allocator> &alloc);
    virtual ~GraphicMetadataBuffer() = default;

    std::shared_ptr<C2Buffer> asC2Buffer() override;

private:
    GraphicMetadataBuffer() = delete;

    std::shared_ptr<C2Allocator> mAlloc;
};

/**
 * MediaCodecBuffer implementation wraps around graphic C2Buffer object.
 *
 * This object exposes the underlying bits via accessor APIs and "image-data"
 * metadata, created automatically at allocation time.
 */
class ConstGraphicBlockBuffer : public Codec2Buffer {
public:
    /**
     * Allocate a new ConstGraphicBlockBuffer wrapping around C2Buffer object.
     * If |buffer| is not in good color formats, it allocates YV12 local buffer
     * and copies the content of |buffer| over to expose.
     *
     * \param   format  mandatory buffer format for MediaCodecBuffer
     * \param   buffer  graphic C2Buffer object to wrap around.
     * \param   alloc   a function to allocate backing ABuffer if needed.
     * \return          ConstGraphicBlockBuffer object with readable mapping.
     *                  nullptr if unsuccessful.
     */
    static sp<ConstGraphicBlockBuffer> Allocate(
            const sp<AMessage> &format,
            const std::shared_ptr<C2Buffer> &buffer,
            std::function<sp<ABuffer>(size_t)> alloc);

    /**
     * Allocate a new ConstGraphicBlockBuffer which allocates YV12 local buffer
     * and copies the content of |buffer| over to expose.
     *
     * \param   format  mandatory buffer format for MediaCodecBuffer
     * \param   alloc   a function to allocate backing ABuffer if needed.
     * \return          ConstGraphicBlockBuffer object with no wrapping buffer.
     */
    static sp<ConstGraphicBlockBuffer> AllocateEmpty(
            const sp<AMessage> &format,
            std::function<sp<ABuffer>(size_t)> alloc);

    virtual ~ConstGraphicBlockBuffer() = default;

    std::shared_ptr<C2Buffer> asC2Buffer() override;
    void clearC2BufferRefs() override;
    bool canCopy(const std::shared_ptr<C2Buffer> &buffer) const override;
    bool copy(const std::shared_ptr<C2Buffer> &buffer) override;

private:
    ConstGraphicBlockBuffer(
            const sp<AMessage> &format,
            const sp<ABuffer> &aBuffer,
            std::unique_ptr<const C2GraphicView> &&view,
            const std::shared_ptr<C2Buffer> &buffer,
            const sp<ABuffer> &imageData,
            bool wrapped);
    ConstGraphicBlockBuffer() = delete;

    sp<ABuffer> mImageData;
    std::unique_ptr<const C2GraphicView> mView;
    std::shared_ptr<C2Buffer> mBufferRef;
    const bool mWrapped;
};

/**
 * MediaCodecBuffer implementation wraps around C2LinearBlock for component
 * and IMemory for client. Underlying C2LinearBlock won't be mapped for secure
 * usecases..
 */
class EncryptedLinearBlockBuffer : public Codec2Buffer {
public:
    /**
     * Construct a new EncryptedLinearBufferBlock wrapping around C2LinearBlock
     * object and writable IMemory region.
     *
     * \param   format      mandatory buffer format for MediaCodecBuffer
     * \param   block       C2LinearBlock object to wrap around.
     * \param   memory      IMemory object to store encrypted content.
     * \param   heapSeqNum  Heap sequence number from ICrypto; -1 if N/A
     */
    EncryptedLinearBlockBuffer(
            const sp<AMessage> &format,
            const std::shared_ptr<C2LinearBlock> &block,
            const sp<IMemory> &memory,
            int32_t heapSeqNum = -1);
    EncryptedLinearBlockBuffer() = delete;

    virtual ~EncryptedLinearBlockBuffer() = default;

    std::shared_ptr<C2Buffer> asC2Buffer() override;

    /**
     * Fill the source buffer structure with appropriate value based on
     * internal IMemory object.
     *
     * \param source  source buffer structure to fill.
     */
    void fillSourceBuffer(
            hardware::drm::V1_0::SharedBuffer *source);
    void fillSourceBuffer(
            hardware::cas::native::V1_0::SharedBuffer *source);

    /**
     * Copy the content of |decrypted| into C2LinearBlock inside. This shall
     * only be called in non-secure usecases.
     *
     * \param   decrypted   decrypted content to copy from.
     * \param   length      length of the content
     * \return  true        if successful
     *          false       otherwise.
     */
    bool copyDecryptedContent(const sp<IMemory> &decrypted, size_t length);

    /**
     * Copy the content of internal IMemory object into C2LinearBlock inside.
     * This shall only be called in non-secure usecases.
     *
     * \param   length      length of the content
     * \return  true        if successful
     *          false       otherwise.
     */
    bool copyDecryptedContentFromMemory(size_t length);

    /**
     * Return native handle of secure buffer understood by ICrypto.
     *
     * \return secure buffer handle
     */
    native_handle_t *handle() const;

    class MappedBlock {
    public:
        explicit MappedBlock(const std::shared_ptr<C2LinearBlock> &block);
        virtual ~MappedBlock();
        bool copyDecryptedContent(const sp<IMemory> &decrypted, size_t length);
    private:
        C2WriteView mView;
    };

    void getMappedBlock(std::unique_ptr<MappedBlock> * const mappedBlock) const;

private:

    std::shared_ptr<C2LinearBlock> mBlock;
    sp<IMemory> mMemory;
    sp<hardware::HidlMemory> mHidlMemory;
    int32_t mHeapSeqNum;
};

/**
 * Get HDR metadata from Gralloc4 handle.
 *
 * \param[in]   handle      handle of the allocation
 * \param[out]  staticInfo  HDR static info to be filled. Ignored if null;
 *                          if |handle| is invalid or does not contain the metadata,
 *                          the shared_ptr is reset.
 * \param[out]  dynamicInfo HDR dynamic info to be filled. Ignored if null;
 *                          if |handle| is invalid or does not contain the metadata,
 *                          the shared_ptr is reset.
 * \return C2_OK if successful
 */
c2_status_t GetHdrMetadataFromGralloc4Handle(
        const C2Handle *const handle,
        std::shared_ptr<C2StreamHdrStaticMetadataInfo::input> *staticInfo,
        std::shared_ptr<C2StreamHdrDynamicMetadataInfo::input> *dynamicInfo);

/**
 * Set metadata to Gralloc4 handle.
 *
 * \param[in]   dataSpace   Dataspace to set.
 * \param[in]   staticInfo  HDR static info to set. Ignored if null or invalid.
 * \param[in]   dynamicInfo HDR dynamic info to set. Ignored if null or invalid.
 * \param[out]  handle      handle of the allocation.
 * \return C2_OK if successful
 */
c2_status_t SetMetadataToGralloc4Handle(
        const android_dataspace_t dataSpace,
        const std::shared_ptr<const C2StreamHdrStaticMetadataInfo::output> &staticInfo,
        const std::shared_ptr<const C2StreamHdrDynamicMetadataInfo::output> &dynamicInfo,
        const C2Handle *const handle);

}  // namespace android

#endif  // CODEC2_BUFFER_H_