/**
 * Copyright (C) 2020 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.
 */
#include "../includes/omxUtils.h"
#define FRAME_WIDTH 2000
#define FRAME_HEIGHT 2000
#define FRAME_RATE (30 << 16)
#define BUFFER_SIZE 12
#define BUFFER_COUNT 2
#define EMPTY_BUFFER_DONE_CALLBACK_TIMEOUT_SEC 30
extern bool mUseTreble;
extern int numCallbackEmptyBufferDone;
sp<IAllocator> mAllocator = IAllocator::getService("ashmem");
int allocateHidlPortBuffers(OMX_U32 portIndex, Vector<Buffer> *buffers,
                            int BufferSize) {
    buffers->clear();
    OMX_PARAM_PORTDEFINITIONTYPE def;
    int err = omxUtilsGetParameter(portIndex, &def);
    omxExitOnError(err);
    for (OMX_U32 i = 0; i < def.nBufferCountActual; ++i) {
        Buffer buffer;
        buffer.mFlags = 0;
        bool success;
        auto transStatus = mAllocator->allocate(BufferSize, [&success, &buffer](
                bool s,
                hidl_memory const& m) {
            success = s;
            buffer.mHidlMemory = m;
        });
        omxExitOnError(!transStatus.isOk());
        omxExitOnError(!success);
        buffers->push(buffer);
    }
    return OK;
}
int main() {
    status_t err;
    /* Initialize OMX for the specified codec                                 */
    status_t ret = omxUtilsInit((char *) "OMX.google.h264.encoder");
    omxExitOnError(ret);
    int allCallbacksReceivedEmptyBufferDone = 0;
    /* Get OMX input port parameters                                          */
    OMX_PARAM_PORTDEFINITIONTYPE *params =
            (OMX_PARAM_PORTDEFINITIONTYPE *) malloc(
                    sizeof(OMX_PARAM_PORTDEFINITIONTYPE));
    params->nPortIndex = OMX_UTILS_IP_PORT;
    params->format.video.nFrameWidth = FRAME_WIDTH;
    params->format.video.nFrameHeight = FRAME_HEIGHT;
    params->format.video.xFramerate = FRAME_RATE;
    params->format.video.eCompressionFormat = OMX_VIDEO_CodingUnused;
    params->format.video.eColorFormat = OMX_COLOR_FormatAndroidOpaque;
    params->nBufferSize = BUFFER_SIZE;
    params->nBufferCountActual = params->nBufferCountMin = BUFFER_COUNT;
    err = omxUtilsSetParameter(OMX_UTILS_IP_PORT, params);
    memset(params, 0, sizeof(OMX_PARAM_PORTDEFINITIONTYPE));
    err = omxUtilsGetParameter(OMX_UTILS_IP_PORT, params);
    /* prepare input port buffers                                             */
    int inMemSize = params->nBufferCountActual * params->nBufferSize;
    int inBufferCnt = params->nBufferCountActual;
    int inBufferSize = inMemSize / inBufferCnt;
    IOMX::buffer_id *inBufferId = new IOMX::buffer_id[inBufferCnt];
    omxUtilsGetParameter(OMX_UTILS_OP_PORT, params);
    /* prepare output port buffers                                            */
    int outMemSize = params->nBufferCountActual * params->nBufferSize;
    int outBufferCnt = params->nBufferCountActual;
    int outBufferSize = outMemSize / outBufferCnt;
    Vector<Buffer> inputBuffers;
    Vector<Buffer> outputBuffers;
    /* Register input buffers with OMX component                              */
    if (mUseTreble) {
        allocateHidlPortBuffers(OMX_UTILS_IP_PORT, &inputBuffers, inBufferSize);
        for (int i = 0; i < inBufferCnt; i++) {
            inBufferId[i] = inputBuffers[i].mID;
            err = omxUtilsUseBuffer(OMX_UTILS_IP_PORT,
                                    inputBuffers[i].mHidlMemory,
                                    &inBufferId[i]);
        }
    } else {
        sp < MemoryDealer > dealerIn = new MemoryDealer(inMemSize);
        for (int i = 0; i < inBufferCnt; i++) {
            sp < IMemory > memory = dealerIn->allocate(inBufferSize);
            memset(memory->pointer(), 0xCF, inBufferSize);
            OMXBuffer omxBuf(memory);
            err = omxUtilsUseBuffer(OMX_UTILS_IP_PORT, omxBuf, &inBufferId[i]);
        }
    }
    IOMX::buffer_id *outBufferId = new IOMX::buffer_id[outBufferCnt];
    /* Register output buffers with OMX component                             */
    if (mUseTreble) {
        allocateHidlPortBuffers(OMX_UTILS_OP_PORT, &outputBuffers,
                                outBufferSize);
        for (int i = 0; i < outBufferCnt; i++) {
            outBufferId[i] = outputBuffers[i].mID;
            err = omxUtilsUseBuffer(OMX_UTILS_OP_PORT,
                                    outputBuffers[i].mHidlMemory,
                                    &outBufferId[i]);
        }
    } else {
        sp < MemoryDealer > dealerOut = new MemoryDealer(outMemSize);
        for (int i = 0; i < outBufferCnt; i++) {
            sp < IMemory > memoryout = dealerOut->allocate(outBufferSize);
            OMXBuffer omxBuf(memoryout);
            err = omxUtilsUseBuffer(OMX_UTILS_OP_PORT, omxBuf, &outBufferId[i]);
        }
    }
    /* Do OMX State change to Idle                                            */
    err = omxUtilsSendCommand(OMX_CommandStateSet, OMX_StateIdle);
    /* Do OMX State change to Executing                                       */
    err = omxUtilsSendCommand(OMX_CommandStateSet, OMX_StateExecuting);
    for (int i = 0; i < inBufferCnt; i++) {
        OMXBuffer omxBuf(0, inBufferSize);
        err = omxUtilsEmptyBuffer(inBufferId[i], omxBuf, 0, 0, -1);
    }
    for (int i = 0; i < outBufferCnt; i++) {
        OMXBuffer omxBuf(0, outBufferSize);
        err = omxUtilsFillBuffer(outBufferId[i], omxBuf, -1);
    }
    /* Do OMX State chage to Idle                                             */
    omxUtilsSendCommand(OMX_CommandStateSet, OMX_StateIdle);
    time_t currentTime = time(NULL);
    time_t waitTimeInSeconds = EMPTY_BUFFER_DONE_CALLBACK_TIMEOUT_SEC;
    time_t endTime = currentTime + waitTimeInSeconds;
    while (currentTime < endTime) {
        if (numCallbackEmptyBufferDone == inBufferCnt) {
            allCallbacksReceivedEmptyBufferDone = 1;
            break;
        }
        currentTime = time(NULL);
    }
    if (!allCallbacksReceivedEmptyBufferDone) {
        ALOGE("Exiting the app");
        exit (EXIT_FAILURE);
    }
    /* Free input and output buffers                                          */
    for (int i = 0; i < inBufferCnt; i++) {
        omxUtilsFreeBuffer(OMX_UTILS_IP_PORT, inBufferId[i]);
    }
    for (int i = 0; i < outBufferCnt; i++) {
        omxUtilsFreeBuffer(OMX_UTILS_OP_PORT, outBufferId[i]);
    }
    /* Free OMX resources                                                     */
    omxUtilsFreeNode();
    return EXIT_SUCCESS;
}