/*
 * Copyright (c) 2019-2022 Arm Limited.
 *
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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
 *
 * 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.
 */

/****************************************************************************
 * Includes
 ****************************************************************************/

#include "FreeRTOS.h"
#include "queue.h"
#include "semphr.h"
#include "task.h"

#include <inttypes.h>
#include <stdio.h>

#include "ethosu_core_interface.h"
#include "message_handler.hpp"
#include "message_queue.hpp"
#include <mailbox.hpp>

#if defined(MHU_V2)
#include <mhu_v2.hpp>
#elif defined(MHU_JUNO)
#include <mhu_juno.hpp>
#else
#include <mhu_dummy.hpp>
#endif

/* Disable semihosting */
__asm(".global __use_no_semihosting\n\t");

using namespace EthosU;
using namespace MessageHandler;

/****************************************************************************
 * Defines
 ****************************************************************************/

// Nr. of tasks to process inferences with, reserves driver & runs inference (Normally 1 per NPU, but not a must)
#if defined(ETHOSU_NPU_COUNT)
constexpr size_t NUM_PARALLEL_TASKS = ETHOSU_NPU_COUNT;
#else
constexpr size_t NUM_PARALLEL_TASKS = 1;
#endif

// TensorArena static initialisation
constexpr size_t arenaSize = TENSOR_ARENA_SIZE / NUM_PARALLEL_TASKS;

__attribute__((section(".bss.tensor_arena"), aligned(16))) uint8_t tensorArena[NUM_PARALLEL_TASKS][arenaSize];

// Message queue from remote host
__attribute__((section("ethosu_core_in_queue"))) MessageQueue::Queue<1000> inputMessageQueue;

// Message queue to remote host
__attribute__((section("ethosu_core_out_queue"))) MessageQueue::Queue<1000> outputMessageQueue;

namespace {

SemaphoreHandle_t messageNotify;
QueueHandle_t inferenceInputQueue;
QueueHandle_t inferenceOutputQueue;

// Mailbox driver
#ifdef MHU_V2
Mailbox::MHUv2 mailbox(MHU_TX_BASE_ADDRESS, MHU_RX_BASE_ADDRESS); // txBase, rxBase
#elif defined(MHU_JUNO)
Mailbox::MHUJuno mailbox(MHU_BASE_ADDRESS);
#else
Mailbox::MHUDummy mailbox;
#endif

} // namespace

/****************************************************************************
 * Mutex & Semaphore
 ****************************************************************************/

extern "C" {

void *ethosu_mutex_create(void) {
    return xSemaphoreCreateMutex();
}

void ethosu_mutex_lock(void *mutex) {
    SemaphoreHandle_t handle = reinterpret_cast<SemaphoreHandle_t>(mutex);
    xSemaphoreTake(handle, portMAX_DELAY);
}

void ethosu_mutex_unlock(void *mutex) {
    SemaphoreHandle_t handle = reinterpret_cast<SemaphoreHandle_t>(mutex);
    xSemaphoreGive(handle);
}

void *ethosu_semaphore_create(void) {
    return xSemaphoreCreateBinary();
}

void ethosu_semaphore_take(void *sem) {
    SemaphoreHandle_t handle = reinterpret_cast<SemaphoreHandle_t>(sem);
    xSemaphoreTake(handle, portMAX_DELAY);
}

void ethosu_semaphore_give(void *sem) {
    SemaphoreHandle_t handle = reinterpret_cast<SemaphoreHandle_t>(sem);
    xSemaphoreGiveFromISR(handle, NULL);
}
}

/****************************************************************************
 * Application
 ****************************************************************************/

namespace {

#ifdef MHU_IRQ
void mailboxIrqHandler() {
    mailbox.handleMessage();
}
#endif

void inferenceTask(void *pvParameters) {
    printf("Starting inference task\n");

    uint8_t *arena = reinterpret_cast<uint8_t *>(pvParameters);
    InferenceHandler process(arena, arenaSize, inferenceInputQueue, inferenceOutputQueue, messageNotify);
    process.run();
}

void messageTask(void *) {
    printf("Starting input message task\n");

    IncomingMessageHandler process(*inputMessageQueue.toQueue(),
                                   *outputMessageQueue.toQueue(),
                                   mailbox,
                                   inferenceInputQueue,
                                   inferenceOutputQueue,
                                   messageNotify);

#ifdef MHU_IRQ
    // Register mailbox interrupt handler
    NVIC_SetVector((IRQn_Type)MHU_IRQ, (uint32_t)&mailboxIrqHandler);
    NVIC_EnableIRQ((IRQn_Type)MHU_IRQ);
#endif

    process.run();
}

} // namespace

// FreeRTOS application. NOTE: Additional tasks may require increased heap size.
int main() {
    BaseType_t ret;

    if (!mailbox.verifyHardware()) {
        printf("Failed to verify mailbox hardware\n");
        return 1;
    }

    // Create message queues for inter process communication
    messageNotify        = xSemaphoreCreateBinary();
    inferenceInputQueue  = xQueueCreate(10, sizeof(ethosu_core_inference_req));
    inferenceOutputQueue = xQueueCreate(10, sizeof(ethosu_core_inference_rsp));

    // Task for handling incoming messages from the remote host
    ret = xTaskCreate(messageTask, "messageTask", 1024, nullptr, 2, nullptr);
    if (ret != pdPASS) {
        printf("Failed to create 'messageTask'\n");
        return ret;
    }

    // One inference task for each NPU
    for (size_t n = 0; n < NUM_PARALLEL_TASKS; n++) {
        ret = xTaskCreate(inferenceTask, "inferenceTask", 8 * 1024, &tensorArena[n], 3, nullptr);
        if (ret != pdPASS) {
            printf("Failed to create 'inferenceTask%d'\n", n);
            return ret;
        }
    }

    // Start Scheduler
    vTaskStartScheduler();

    return 1;
}