/**
 * Copyright (c) 2020-2023 Hailo Technologies Ltd. All rights reserved.
 * Distributed under the MIT license (https://opensource.org/licenses/MIT)
 **/
/**
 * @file async_infer_runner.cpp
 * @brief Implemention of the async HL infer
 **/

#include <iostream>

#include "common/utils.hpp"
#include "common/os_utils.hpp"
#include "hailo/event.hpp"
#include "hailo/hailort_defaults.hpp"
#include "hailo/hailort_common.hpp"
#include "net_flow/pipeline/async_infer_runner.hpp"
#include "net_flow/pipeline/pipeline_internal.hpp"
#include "net_flow/ops/op_metadata.hpp"

namespace hailort
{

Expected<std::shared_ptr<AsyncPipeline>> AsyncPipeline::create_shared()
{
    auto async_pipeline_ptr = make_shared_nothrow<AsyncPipeline>();
    CHECK_NOT_NULL_AS_EXPECTED(async_pipeline_ptr, HAILO_OUT_OF_HOST_MEMORY);
    return async_pipeline_ptr;
}

AsyncPipeline::AsyncPipeline() : m_is_multi_planar(false) {}

void AsyncPipeline::add_element_to_pipeline(std::shared_ptr<PipelineElement> pipeline_element)
{
    m_pipeline_elements.push_back(pipeline_element);
}

void AsyncPipeline::set_async_hw_element(std::shared_ptr<AsyncHwElement> async_hw_element)
{
    m_async_hw_element = async_hw_element;
}

void AsyncPipeline::add_entry_element(std::shared_ptr<PipelineElement> pipeline_element, const std::string &input_name)
{
    assert(!contains(m_entry_elements, input_name));
    m_entry_elements[input_name] = pipeline_element;
}

void AsyncPipeline::add_last_element(std::shared_ptr<PipelineElement> pipeline_element, const std::string &output_name)
{
    assert(!contains(m_last_elements, output_name));
    m_last_elements[output_name] = pipeline_element;
}

void AsyncPipeline::set_build_params(ElementBuildParams &build_params)
{
    m_build_params = build_params;
}

void AsyncPipeline::shutdown(hailo_status error_status)
/*
Async pipeline shutdown handling:
Shutdown the pipeline is considered unrecoverable, so the entire pipeline and core op (or net_group) are assumed to be unusable*.
Shutdown can originate from internal sources:
- Errors from the core.
- Errors from getting buffers from one of the buffer pools
- Error from enqueue/dequeue buffers in a AsyncPushQueue element

or external sources:
- User request to abort.

The flow in case of shutdown is:
1. Set the pipeline_status to error_status - will block new infer requests from coming
2. Shutdown threads in all elements
3. Dequeue all user buffers with the error_status**

*  - TODO: add support for resume flow
** - if there are user buffers in an AsyncPushQueueElement (inside the queue itself) - the element will clear them after stop_thread() is complete


*/
{
    if (HAILO_STREAM_ABORTED_BY_USER == error_status) {
        LOGGER__INFO("Pipeline was aborted by user. Shutting it down");
    } else {
        LOGGER__ERROR("Shutting down the pipeline with status {}", error_status);
    }
    m_build_params.pipeline_status->store(error_status);
    auto status = m_build_params.shutdown_event->signal();
    if (HAILO_SUCCESS != status) {
        LOGGER__CRITICAL("Executing pipeline shutdown failed with status {}", status);
    }

    for (auto &element : get_entry_elements()) {
        status = element.second->terminate(error_status);
        if (HAILO_SUCCESS != status) {
            LOGGER__CRITICAL("Executing pipeline terminate failed with status {}", status);
        }
    }

    for (auto &element : get_entry_elements()) {
        status = element.second->dequeue_user_buffers(error_status);
        if (HAILO_SUCCESS != status) {
            LOGGER__CRITICAL("Dequeueing external buffers failed with status {}", status);
        }
    }
}

const std::vector<std::shared_ptr<PipelineElement>>& AsyncPipeline::get_pipeline() const
{
    return m_pipeline_elements;
}

const std::unordered_map<std::string, std::shared_ptr<PipelineElement>>& AsyncPipeline::get_entry_elements() const
{
    return m_entry_elements;
}

const std::unordered_map<std::string, std::shared_ptr<PipelineElement>>& AsyncPipeline::get_last_elements() const
{
    return m_last_elements;
}

const std::shared_ptr<AsyncHwElement> AsyncPipeline::get_async_hw_element()
{
    return m_async_hw_element;
}

const ElementBuildParams AsyncPipeline::get_build_params()
{
    return m_build_params;
}

void AsyncPipeline::set_as_multi_planar()
{
    m_is_multi_planar = true;
}

bool AsyncPipeline::is_multi_planar()
{
    return m_is_multi_planar;
}

Expected<std::shared_ptr<AsyncInferRunnerImpl>> AsyncInferRunnerImpl::create(std::shared_ptr<ConfiguredNetworkGroup> net_group,
    const std::unordered_map<std::string, hailo_format_t> &inputs_formats, const std::unordered_map<std::string, hailo_format_t> &outputs_formats,
    const uint32_t timeout)
{
    auto pipeline_status = make_shared_nothrow<std::atomic<hailo_status>>(HAILO_SUCCESS);
    CHECK_AS_EXPECTED(nullptr != pipeline_status, HAILO_OUT_OF_HOST_MEMORY);

    auto async_pipeline_expected = PipelineBuilder::create_pipeline(net_group, inputs_formats, outputs_formats, timeout, pipeline_status);
    CHECK_EXPECTED(async_pipeline_expected);

    auto async_infer_runner_ptr = make_shared_nothrow<AsyncInferRunnerImpl>(async_pipeline_expected.release(), pipeline_status);
    CHECK_NOT_NULL_AS_EXPECTED(async_infer_runner_ptr, HAILO_OUT_OF_HOST_MEMORY);

    auto status = async_infer_runner_ptr->start_pipeline();
    CHECK_SUCCESS_AS_EXPECTED(status);

    return async_infer_runner_ptr;
}

AsyncInferRunnerImpl::AsyncInferRunnerImpl(std::shared_ptr<AsyncPipeline> async_pipeline, std::shared_ptr<std::atomic<hailo_status>> pipeline_status) :
    m_async_pipeline(async_pipeline),
    m_is_activated(false),
    m_is_aborted(false),
    m_pipeline_status(pipeline_status)
{}

AsyncInferRunnerImpl::~AsyncInferRunnerImpl()
{
    (void)stop_pipeline();
}

hailo_status AsyncInferRunnerImpl::stop_pipeline()
{
    hailo_status status = HAILO_SUCCESS;
    if (m_is_activated) {
        m_is_activated = false;
        for (auto &entry_element : m_async_pipeline->get_entry_elements()) {
            status = entry_element.second->deactivate();
            if (HAILO_SUCCESS != status) {
                LOGGER__WARNING("Failed deactivate of element {} status {}", entry_element.second->name(), status);
            }

            auto should_clear_abort = (!m_is_aborted);
            status = entry_element.second->post_deactivate(should_clear_abort);
            if (HAILO_SUCCESS != status) {
                LOGGER__WARNING("Failed post deactivate of element {} status {}", entry_element.second->name(), status);
            }
        }
    }
    return status;
}

hailo_status AsyncInferRunnerImpl::start_pipeline()
{
    hailo_status status = HAILO_SUCCESS;
    for (auto &entry_element : m_async_pipeline->get_entry_elements()) {
        status = entry_element.second->activate();
        CHECK_SUCCESS(status);
    }

    m_is_activated = true;

    return status;
}

void AsyncInferRunnerImpl::abort()
{
    m_is_aborted = true;
    m_async_pipeline->shutdown(HAILO_STREAM_ABORTED_BY_USER);
    return;
}

Expected<bool> AsyncInferRunnerImpl::can_push_buffers()
{
    for (auto &last_element : m_async_pipeline->get_last_elements()) {
        auto can_push_buffer = last_element.second->can_push_buffer_upstream(last_element.first);
        CHECK_EXPECTED(can_push_buffer);
        if (!can_push_buffer.release()) {
            return false;
        }
    }

    for (auto &entry_element : m_async_pipeline->get_entry_elements()) {
        auto can_push_buffer = entry_element.second->can_push_buffer_downstream(entry_element.first);
        CHECK_EXPECTED(can_push_buffer);
        if (!can_push_buffer.release()) {
            return false;
        }
    }

    return true;
}

hailo_status AsyncInferRunnerImpl::async_infer()
{
    hailo_status status = m_async_pipeline->get_build_params().pipeline_status->load();
    CHECK_SUCCESS(status, "Can't handle infer request since Pipeline status is {}.", status);

    auto pools_are_ready = can_push_buffers();
    CHECK_EXPECTED_AS_STATUS(pools_are_ready);
    CHECK(pools_are_ready.release(), HAILO_QUEUE_IS_FULL, "Can't handle infer request since a queue in the pipeline is full.");

    for (auto &last_element : m_async_pipeline->get_last_elements()) {
        assert(contains(m_output_buffers, last_element.first));
        auto output_buffer = m_output_buffers.at(last_element.first);
        auto read_done = m_read_dones.at(last_element.first);
        // TODO: handle the non-recoverable case where one buffer is enqueued successfully and the second isn't (HRT-11783)
        status = last_element.second->enqueue_execution_buffer(output_buffer, read_done);
        CHECK_SUCCESS(status);
    }

    for (auto &entry_element : m_async_pipeline->get_entry_elements()) {
        assert(contains(m_input_buffers, entry_element.first));
        entry_element.second->sinks()[0].run_push_async(std::move(m_input_buffers.at(entry_element.first)));
    }
    return HAILO_SUCCESS;
}

void AsyncInferRunnerImpl::add_element_to_pipeline(std::shared_ptr<PipelineElement> pipeline_element)
{
    m_async_pipeline->add_element_to_pipeline(pipeline_element);
}

void AsyncInferRunnerImpl::add_entry_element(std::shared_ptr<PipelineElement> pipeline_element, const std::string &input_name)
{
    m_async_pipeline->add_entry_element(pipeline_element, input_name);
}

void AsyncInferRunnerImpl::add_last_element(std::shared_ptr<PipelineElement> pipeline_element, const std::string &output_name)
{
    m_async_pipeline->add_last_element(pipeline_element, output_name);
}

std::unordered_map<std::string, std::shared_ptr<PipelineElement>> AsyncInferRunnerImpl::get_entry_elements()
{
    return m_async_pipeline->get_entry_elements();
}

std::unordered_map<std::string, std::shared_ptr<PipelineElement>> AsyncInferRunnerImpl::get_last_elements()
{
    return m_async_pipeline->get_last_elements();
}

void AsyncInferRunnerImpl::set_input(const std::string &input_name, MemoryView &&input_buffer, TransferDoneCallbackAsyncInfer &write_done)
{
    m_input_buffers[input_name] = PipelineBuffer(std::move(input_buffer), write_done);
}

void AsyncInferRunnerImpl::set_input(const std::string &input_name, hailo_pix_buffer_t input_buffer, TransferDoneCallbackAsyncInfer &write_done)
{
    // If only one plane is passed, address it as memview
    if (1 == input_buffer.number_of_planes) {
        m_input_buffers[input_name] = PipelineBuffer(MemoryView(input_buffer.planes[0].user_ptr, input_buffer.planes[0].bytes_used), write_done);
    } else if (m_async_pipeline->is_multi_planar()) {
        // If model is multi-planar
        m_input_buffers[input_name] = PipelineBuffer(std::move(input_buffer), write_done);
    } else {
        // Other cases - return error, as on async flow we do not support copy to new buffer
        LOGGER__ERROR("HEF was compiled for single input layer, while trying to pass non-contiguous planes buffers.");
        m_input_buffers[input_name] = PipelineBuffer(HAILO_INVALID_OPERATION, write_done);
    }
}

void AsyncInferRunnerImpl::set_output(const std::string &output_name, MemoryView &&output_buffer, TransferDoneCallbackAsyncInfer &read_done)
{
    m_output_buffers[output_name] = std::move(output_buffer);
    m_read_dones[output_name] = read_done;
}

std::vector<std::shared_ptr<PipelineElement>> AsyncInferRunnerImpl::get_pipeline() const
{
    return m_async_pipeline->get_pipeline();
}

std::string AsyncInferRunnerImpl::get_pipeline_description() const
{
    std::stringstream pipeline_str;
    pipeline_str << "Async infer pipeline description:\n";
    for (const auto &element : get_pipeline()) {
	    pipeline_str << " >> " << element->description();
    }
    return pipeline_str.str();
}

hailo_status AsyncInferRunnerImpl::get_pipeline_status() const
{
    return m_pipeline_status->load();
}

std::shared_ptr<AsyncPipeline> AsyncInferRunnerImpl::get_async_pipeline() const
{
    return m_async_pipeline;
}

} /* namespace hailort */
