/*
 * Copyright (C) 2022-2026 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
 */

#include "context.hpp"

#include "device.hpp"
#include "driver.hpp"
#include "driver_handle.hpp"
#include "event.hpp"
#include "level_zero_driver/include/l0_exception.hpp"
#include "level_zero_driver/source/ext/disk_cache.hpp"
#include "metric.hpp"
#include "metric_query.hpp"
#include "metric_streamer.hpp"
#include "vpu_driver/source/device/hw_info.hpp"
#include "vpu_driver/source/device/vpu_device.hpp"
#include "vpu_driver/source/utilities/log.hpp"

#include <errno.h>
#include <limits>
#include <linux/sysinfo.h>
#include <string.h>
#include <sys/sysinfo.h>
#include <ze_api.h>
#include <zet_api.h>

namespace L0 {

Context::Context(DriverHandle *driverHandle, std::unique_ptr<VPU::VPUDeviceContext> ctx)
    : driverHandle(driverHandle)
    , ctx(std::move(ctx)) {
    if (this->ctx->getDeviceCapabilities().npuArch >= VPU::NPU40XX) {
        resourceCleaner = std::make_unique<ResourceCleaner>(this, idleTimeout);
    }
}

ze_result_t Context::destroy() {
    delete this;

    return ZE_RESULT_SUCCESS;
}

DriverHandle *Context::getDriverHandle() {
    return this->driverHandle;
}

ze_result_t Context::getStatus() {
    auto device = driverHandle->getPrimaryDevice();
    if (device == nullptr) {
        LOG_E("Driver handle failed to retrieve primary device");
        return ZE_RESULT_ERROR_DEVICE_LOST;
    }

    auto vpuDevice = device->getVPUDevice();
    if (vpuDevice == nullptr) {
        LOG_E("VPU device failed to be retrieved");
        return ZE_RESULT_ERROR_DEVICE_LOST;
    }

    return vpuDevice->isConnected() ? ZE_RESULT_SUCCESS : ZE_RESULT_ERROR_DEVICE_LOST;
}

ze_result_t Context::activateMetricGroups(zet_device_handle_t hDevice,
                                          uint32_t count,
                                          zet_metric_group_handle_t *phMetricGroups) {
    if (hDevice == nullptr) {
        LOG_E("Device handle is NULL");
        return ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
    }

    return L0::Device::fromHandle(hDevice)->activateMetricGroups(ctx->getFd(),
                                                                 count,
                                                                 phMetricGroups);
}

ze_result_t Context::createMetricQueryPool(zet_device_handle_t hDevice,
                                           zet_metric_group_handle_t hMetricGroup,
                                           const zet_metric_query_pool_desc_t *desc,
                                           zet_metric_query_pool_handle_t *phMetricQueryPool) {
    if (hDevice == nullptr || hMetricGroup == nullptr) {
        LOG_E("Device(%p) / MetricGroup(%p) handle is NULL.", hDevice, hMetricGroup);
        return ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
    }

    if (desc == nullptr || phMetricQueryPool == nullptr) {
        LOG_E("Desc(%p) / MetricQueryPool(%p) handle is NULL.", desc, phMetricQueryPool);
        return ZE_RESULT_ERROR_INVALID_NULL_POINTER;
    }

    if (desc->type > ZET_METRIC_QUERY_POOL_TYPE_EXECUTION) {
        LOG_E("Desc type value invalid: %x", desc->type);
        return ZE_RESULT_ERROR_INVALID_ENUMERATION;
    }

    if (desc->count == 0) {
        LOG_E("Desc count value invalid: %u", desc->count);
        return ZE_RESULT_ERROR_INVALID_SIZE;
    }

    auto device = Device::fromHandle(hDevice);
    if (!device->isMetricsLoaded()) {
        LOG_E("Device metrics is not initialized");
        return ZE_RESULT_ERROR_UNINITIALIZED;
    }

    auto metricGroup = MetricGroup::fromHandle(hMetricGroup);
    if (!metricGroup->isActivated()) {
        LOG_E("MetricGroup (%p) is not activated! Please activate metric group before QueryPool "
              "creation.",
              metricGroup);
        return ZE_RESULT_NOT_READY;
    }

    try {
        auto metricQueryPool = std::make_unique<MetricQueryPool>(this, *metricGroup, desc->count);
        *phMetricQueryPool = metricQueryPool.get();

        this->appendObject(std::move(metricQueryPool));

        LOG(CONTEXT, "MetricQueryPool created - %p", *phMetricQueryPool);
    } catch (const DriverError &err) {
        return err.result();
    }

    return ZE_RESULT_SUCCESS;
}

ze_result_t Context::metricStreamerOpen(zet_device_handle_t hDevice,
                                        zet_metric_group_handle_t hMetricGroup,
                                        zet_metric_streamer_desc_t *desc,
                                        ze_event_handle_t hNotificationEvent,
                                        zet_metric_streamer_handle_t *phMetricStreamer) {
    if (hDevice == nullptr) {
        LOG_E("Invalid hDevice pointer");
        return ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
    }

    if (hMetricGroup == nullptr) {
        LOG_E("Invalid hMetricGroup pointer");
        return ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
    }

    if (desc == nullptr) {
        LOG_E("Invalid desc pointer");
        return ZE_RESULT_ERROR_INVALID_NULL_POINTER;
    }

    if (phMetricStreamer == nullptr) {
        LOG_E("Invalid phMetricStreamer pointer");
        return ZE_RESULT_ERROR_INVALID_NULL_POINTER;
    }

    auto device = Device::fromHandle(hDevice);
    if (!device->isMetricsLoaded()) {
        LOG_E("Device metrics is not initialized");
        return ZE_RESULT_ERROR_UNINITIALIZED;
    }

    auto metricGroup = MetricGroup::fromHandle(hMetricGroup);
    if (!metricGroup->isActivated()) {
        LOG_E("MetricGroup (%p) is not activated.", metricGroup);
        return ZE_RESULT_NOT_READY;
    }

    if (desc->samplingPeriod < MetricContext::MIN_SAMPLING_RATE_NS) {
        LOG_E("Sampling rate is too low");
        return ZE_RESULT_ERROR_INVALID_ARGUMENT;
    }

    try {
        auto metricStreamer =
            std::make_unique<MetricStreamer>(this, metricGroup, desc, hNotificationEvent);
        *phMetricStreamer = metricStreamer.get();
        this->appendObject(std::move(metricStreamer));

        LOG(CONTEXT, "MetricStreamer created - %p", *phMetricStreamer);
    } catch (const DriverError &err) {
        return err.result();
    }

    return ZE_RESULT_SUCCESS;
}

ze_result_t Context::queryContextMemory(ze_graph_memory_query_type_t type,
                                        ze_graph_memory_query_t *query) {
    switch (type) {
    case ZE_GRAPH_QUERY_MEMORY_DDR: {
        struct sysinfo info = {};
        if (sysinfo(&info) < 0) {
            LOG_E("Failed to get total ram using sysinfo, errno: %i, str: %s",
                  errno,
                  strerror(errno));
        } else {
            query->total = info.totalram * info.mem_unit;
        }
        query->allocated = ctx->getAllocatedSize();
    } break;
    case ZE_GRAPH_QUERY_MEMORY_DRIVER_CACHE:
    case ZE_GRAPH_QUERY_MEMORY_PROGRAM_CACHE: {
        auto &diskCache = Driver::getInstance()->getDiskCache();
        query->allocated = diskCache.getCacheSize();
        query->total = diskCache.getMaxSize();
    } break;
    default:
        LOG_E("Unsupported type");
        return ZE_RESULT_ERROR_INVALID_ARGUMENT;
    }

    return ZE_RESULT_SUCCESS;
}

ze_result_t Context::setProperties(const ze_context_properties_npu_ext_t *pContextProperties) {
    bool enableIdleOptimizations =
        (pContextProperties->options & ZE_NPU_CONTEXT_OPTION_IDLE_OPTIMIZATIONS) != 0;
    if (enableIdleOptimizations) {
        if (ctx->getDeviceCapabilities().npuArch < VPU::NPU40XX) {
            return ZE_RESULT_SUCCESS;
        }
        std::unique_lock<std::mutex> lock(cleanerMutex);
        if (!resourceCleaner) {
            resourceCleaner = std::make_unique<ResourceCleaner>(this, idleTimeout);
        }
    } else {
        std::unique_lock<std::mutex> lock(cleanerMutex);
        resourceCleaner = nullptr;
    }

    return ZE_RESULT_SUCCESS;
}

ze_result_t Context::releaseMemory() {
    ctx->preemptionCachePrune();
    ctx->scratchCachePrune(std::numeric_limits<size_t>::max());
    return ZE_RESULT_SUCCESS;
}

void Context::setIdle() {
    std::unique_lock<std::mutex> lock(cleanerMutex);
    if (resourceCleaner) {
        resourceCleaner->setIdle();
    }
}

void Context::setIdlePruningTimeout(uint64_t timeout) {
    idleTimeout = std::chrono::milliseconds(timeout);
    std::unique_lock<std::mutex> lock(cleanerMutex);
    if (resourceCleaner) {
        resourceCleaner->setIdleTimeout(idleTimeout);
    }
}

ResourceCleaner::ResourceCleaner(Context *ctx, std::chrono::milliseconds timeout)
    : idleTimeout(timeout)
    , thread(
          [this](Context *ctx) {
              std::unique_lock<std::mutex> lock(mutex);
              auto timeout = std::chrono::time_point<std::chrono::steady_clock>::max();
              while (action != Action::BREAK) {
                  if (cv.wait_until(lock, timeout) == std::cv_status::timeout) {
                      ctx->releaseMemory();
                      timeout = std::chrono::time_point<std::chrono::steady_clock>::max();
                  } else if (action == Action::PRUNE_AFTER_TIMEOUT) {
                      timeout = std::chrono::steady_clock::now() + idleTimeout;
                  }
              }
          },
          ctx) {}

ResourceCleaner::~ResourceCleaner() {
    {
        std::unique_lock<std::mutex> lock(mutex);
        action = Action::BREAK;
    }
    cv.notify_one();
    thread.join();
}

void ResourceCleaner::setIdle() {
    {
        std::unique_lock<std::mutex> lock(mutex);
        action = Action::PRUNE_AFTER_TIMEOUT;
    }
    cv.notify_one();
}

void ResourceCleaner::setIdleTimeout(std::chrono::milliseconds timeout) {
    idleTimeout = timeout;
    cv.notify_one();
}

} // namespace L0
