// Copyright © Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier:  MIT

#pragma once

#include <hipdnn_data_sdk/logging/CallbackTypes.h>
#include <hipdnn_data_sdk/logging/LogLevel.hpp>
#include <hipdnn_data_sdk/logging/Logger.hpp>

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <initializer_list>
#include <iostream>
#include <mutex>
#include <sstream>
#include <string>
#include <vector>

namespace hipdnn_test_sdk::utilities
{

/**
 * @brief Recorded log entry with severity and message
 */
struct RecordedLog
{
    hipdnnSeverity_t severity;
    std::string message;
};

/**
 * @brief Thread-safe log storage with named instances
 *
 * This class provides centralized log storage for test verification.
 * Supports separate instances for test infrastructure (SHARED)
 * and user test code (ISOLATED) to prevent interference.
 */
class LogRecording
{
public:
    /**
     * @brief Named instance identifiers
     *
     * SHARED: Test harness infrastructure (initialized in main.cpp, uses
     *         logRecordingCallback() or logChainedRecordingCallback(), used
     *         by SharedLogRecorder)
     *         This instance is typically used to record logs intercepted from
     *         the UUT-specific logger and is shared across all tests run by
     *         the test harness.
     * ISOLATED: Individual test code (used by IsolatedLogRecorder)
     *           Does not capture logs generated by the UUT; each test must
     *           register the IsolatedLogRecorder callback function with
     *           the UUT in order to record logs. Logs are therefore isolated
     *           to one test (or test fixture) that registers its callback.
     */
    enum class Id
    {
        SHARED,
        ISOLATED
    };

    static LogRecording& instance(Id id = Id::ISOLATED)
    {
        static LogRecording s_shared;
        static LogRecording s_isolated;
        return (id == Id::SHARED) ? s_shared : s_isolated;
    }

    void startRecording()
    {
        _isRecording.store(true, std::memory_order_release);
    }

    void stopRecording()
    {
        _isRecording.store(false, std::memory_order_release);
    }

    bool isRecording() const
    {
        return _isRecording.load(std::memory_order_acquire);
    }

    void recordLog(hipdnnSeverity_t severity, const std::string& message)
    {
        if(!_isRecording.load(std::memory_order_acquire))
        {
            return;
        }

        {
            const std::lock_guard<std::mutex> lock(_logsMutex);
            _recordedLogs.push_back({severity, message});
        }
        _cvLogRecorded.notify_all();
    }

    void clearLogs()
    {
        const std::lock_guard<std::mutex> lock(_logsMutex);
        _recordedLogs.clear();
    }

    std::vector<RecordedLog> getRecordedLogs() const
    {
        const std::lock_guard<std::mutex> lock(_logsMutex);
        return _recordedLogs;
    }

    size_t getRecordedLogCount() const
    {
        const std::lock_guard<std::mutex> lock(_logsMutex);
        return _recordedLogs.size();
    }

    /**
     * @brief Wait until log count reaches target, with timeout
     * @param targetCount The target log count to wait for
     * @param timeout Maximum time to wait
     * @return true if target reached, false if timed out
     */
    [[nodiscard]] bool waitForLogCount(size_t targetCount, std::chrono::milliseconds timeout)
    {
        std::unique_lock<std::mutex> lock(_logsMutex);
        return _cvLogRecorded.wait_for(
            lock, timeout, [&] { return _recordedLogs.size() >= targetCount; });
    }

    /**
     * @brief Wait until logs containing all given texts are recorded, with a single shared timeout
     * @param texts The substrings to search for in log messages
     * @param timeout Maximum total time to wait for all texts to appear
     * @return true if all matching logs found, false if timed out
     */
    [[nodiscard]] bool waitForLogsContaining(std::initializer_list<std::string> texts,
                                             std::chrono::milliseconds timeout)
    {
        std::unique_lock<std::mutex> lock(_logsMutex);
        return _cvLogRecorded.wait_for(lock, timeout, [&] {
            for(const auto& text : texts)
            {
                bool found = false;
                for(const auto& log : _recordedLogs)
                {
                    if(log.message.find(text) != std::string::npos)
                    {
                        found = true;
                        break;
                    }
                }
                if(!found)
                {
                    return false;
                }
            }
            return true;
        });
    }

    LogRecording(const LogRecording&) = delete;
    LogRecording& operator=(const LogRecording&) = delete;

private:
    LogRecording() = default;
    ~LogRecording() = default;

    mutable std::mutex _logsMutex;
    std::condition_variable _cvLogRecorded;
    std::atomic<bool> _isRecording{false};
    std::vector<RecordedLog> _recordedLogs;
};

/**
 * @brief Manages simple output of recorded logs for developer visibility
 *
 * This singleton handles filtering and outputting logs from the test harness
 * based on the log level configured during initialize() (via HIPDNN_LOG_LEVEL).
 *
 * Responsibilities:
 * - Store log level (fixed at startup)
 * - Filter test output based on configured level
 * - Chain to log output callback
 */
class LogRecordingOutput
{
public:
    static LogRecordingOutput& instance()
    {
        static LogRecordingOutput s_instance;
        return s_instance;
    }

    void initialize(hipdnnSeverity_t hipdnnLogLevel, hipdnnCallback_t logOutputCallback)
    {
        _minLogLevel = hipdnnLogLevel;
        _logOutputCallback = logOutputCallback;
    }

    void outputToChainedCallback(hipdnnSeverity_t severity, const char* message, bool force) const
    {
        if(_logOutputCallback != nullptr)
        {
            if(force || (_minLogLevel != HIPDNN_SEV_OFF && severity >= _minLogLevel))
            {
                _logOutputCallback(severity, message);
            }
        }
    }

    LogRecordingOutput(const LogRecordingOutput&) = delete;
    LogRecordingOutput& operator=(const LogRecordingOutput&) = delete;

private:
    LogRecordingOutput() = default;
    ~LogRecordingOutput() = default;

    hipdnnSeverity_t _minLogLevel{HIPDNN_SEV_OFF};
    hipdnnCallback_t _logOutputCallback{nullptr};
};

inline const char* severityToString(hipdnnSeverity_t severity)
{
    switch(severity)
    {
    case HIPDNN_SEV_INFO:
        return "info";
    case HIPDNN_SEV_WARN:
        return "warn";
    case HIPDNN_SEV_ERROR:
        return "error";
    case HIPDNN_SEV_FATAL:
        return "fatal";
    case HIPDNN_SEV_OFF:
        return "off";
    default:
        return "UNKNOWN";
    }
}

/**
 * @brief Pre-made logging callback for use by SharedLogRecorder that records all
 * logs and forwards them to LogRecordingOutput with force = false so that the
 * chained callback is called only if the HIPDNN_LOG_LEVEL allows.
 */
inline void logRecordingCallback(hipdnnSeverity_t severity, const char* message)
{
    if(message == nullptr)
    {
        return;
    }
    // Always record for test harness infrastructure (ignores level)
    LogRecording::instance(LogRecording::Id::SHARED).recordLog(severity, std::string(message));

    // Forward to chained callback if test harness level allows
    LogRecordingOutput::instance().outputToChainedCallback(severity, message, false);
}

/**
 * @brief Pre-made logging callback for use by SharedLogRecorder that records all
 * logs and forwards them to LogRecordingOutput with force = true so that the
 * chained callback is always called.
 */
inline void logChainedRecordingCallback(hipdnnSeverity_t severity, const char* message)
{
    if(message == nullptr)
    {
        return;
    }
    // Always record for test harness infrastructure (ignores level)
    LogRecording::instance(LogRecording::Id::SHARED).recordLog(severity, std::string(message));

    // Forward to chained callback
    LogRecordingOutput::instance().outputToChainedCallback(severity, message, true);
}

/**
 * @brief Simple std::cerr console output callback for use by LogRecordingOutput
 */
inline void simpleStderrOutputCallback([[maybe_unused]] hipdnnSeverity_t severity,
                                       const char* message)
{
#ifndef DISABLE_TEST_LOGGING
    if(message != nullptr)
    {
        std::cerr << "[" << severityToString(severity) << "] " << message << '\n';
    }
#else
    (void)message;
#endif
}

/**
 * @brief Abstract base class for RAII log recording in tests
 *
 * Provides shared verification methods for log queries.
 * Use SharedLogRecorder or IsolatedLogRecorder depending on your use case.
 */
class LogRecorderBase
{
public:
    virtual ~LogRecorderBase()
    {
        LogRecording::instance(_recordingId).stopRecording();
        LogRecording::instance(_recordingId).clearLogs();

        // Always restore UUT log level
        hipdnn_data_sdk::logging::setLogLevel(_savedLogLevel);
    }

    // Non-copyable, non-movable
    LogRecorderBase(const LogRecorderBase&) = delete;
    LogRecorderBase& operator=(const LogRecorderBase&) = delete;
    LogRecorderBase(LogRecorderBase&&) = delete;
    LogRecorderBase& operator=(LogRecorderBase&&) = delete;

    /**
     * @brief Get the log level that was active before this recorder was created
     */
    hipdnnSeverity_t getSavedLogLevel() const
    {
        return _savedLogLevel;
    }

    // === Instance Verification Methods ===

    bool hasLogContaining(const std::string& text) const
    {
        auto logs = LogRecording::instance(_recordingId).getRecordedLogs();
        for(const auto& log : logs)
        {
            if(log.message.find(text) != std::string::npos)
            {
                return true;
            }
        }
        return false;
    }

    bool hasLogContaining(hipdnnSeverity_t severity, const std::string& text) const
    {
        auto logs = LogRecording::instance(_recordingId).getRecordedLogs();
        for(const auto& log : logs)
        {
            if(log.severity == severity && log.message.find(text) != std::string::npos)
            {
                return true;
            }
        }
        return false;
    }

    size_t getRecordedLogCount() const
    {
        return LogRecording::instance(_recordingId).getRecordedLogCount();
    }

    size_t countLogsAtLevel(hipdnnSeverity_t severity) const
    {
        auto logs = LogRecording::instance(_recordingId).getRecordedLogs();
        size_t count = 0;
        for(const auto& log : logs)
        {
            if(log.severity == severity)
            {
                ++count;
            }
        }
        return count;
    }

    std::string getRecordedLogsAsString(size_t maxLogs = 0) const
    {
        auto logs = LogRecording::instance(_recordingId).getRecordedLogs();

        if(logs.empty())
        {
            return "(No logs captured.)\n";
        }

        std::ostringstream oss;
        const size_t logsToOutput = (maxLogs == 0 || maxLogs > logs.size()) ? logs.size() : maxLogs;

        for(size_t i = 0; i < logsToOutput; ++i)
        {
            oss << "[" << severityToString(logs[i].severity) << "] " << logs[i].message << "\n";
        }

        if(maxLogs > 0 && logs.size() > maxLogs)
        {
            const size_t skipped = logs.size() - maxLogs;
            oss << "(Skipped " << skipped << " additional log" << (skipped > 1 ? "s" : "")
                << ".)\n";
        }

        return oss.str();
    }

    std::vector<RecordedLog> getRecordedLogs() const
    {
        return LogRecording::instance(_recordingId).getRecordedLogs();
    }

    /**
     * @brief Clear all recorded logs for the current log recorder.
     */
    void clearLogs()
    {
        LogRecording::instance(_recordingId).clearLogs();
    }

    /**
     * @brief Wait until log count reaches target, with timeout
     * @param targetCount The target log count to wait for
     * @param timeout Maximum time to wait
     * @return true if target reached, false if timed out
     */
    [[nodiscard]] bool waitForLogCount(size_t targetCount, std::chrono::milliseconds timeout)
    {
        return LogRecording::instance(_recordingId).waitForLogCount(targetCount, timeout);
    }

    /**
     * @brief Wait until logs containing all given texts are recorded, with a single shared timeout
     * @param texts The substrings to search for in log messages
     * @param timeout Maximum total time to wait for all texts to appear
     * @return true if all matching logs found, false if timed out
     */
    [[nodiscard]] bool waitForLogsContaining(std::initializer_list<std::string> texts,
                                             std::chrono::milliseconds timeout)
    {
        return LogRecording::instance(_recordingId).waitForLogsContaining(texts, timeout);
    }

protected:
    /**
     * @brief Protected constructor - use factory methods in derived classes
     *
     * @param captureLevel Level to set UUT to
     * @param shouldChangeLevel If true, changes UUT level; if false, preserves it
     * @param recordingId Which LogRecording instance to use
     */
    LogRecorderBase(hipdnnSeverity_t captureLevel,
                    bool shouldChangeLevel,
                    LogRecording::Id recordingId)
        : _savedLogLevel(hipdnn_data_sdk::logging::getLogLevel())
        , _recordingId(recordingId)
    {
        // Change UUT log level if requested
        if(shouldChangeLevel)
        {
            hipdnn_data_sdk::logging::setLogLevel(captureLevel);
        }

        // Clear and start recording
        LogRecording::instance(_recordingId).clearLogs();
        LogRecording::instance(_recordingId).startRecording();
    }

    hipdnnSeverity_t _savedLogLevel{HIPDNN_SEV_OFF};
    LogRecording::Id _recordingId;
};

/**
 * @brief RAII log recorder for tests using the globally-registered logging callback
 * and always restores the original UUT log level on destruction.
 *
 * Use this recorder in unit tests where the test harness (main.cpp) has already
 * registered logRecordingCallback() or logChainedRecordingCallback() with the UUT's
 * logger. This LogRecorder is shared across all tests using the harness.
 *
 * Two modes available:
 * 1. **withCurrentLevel()** - Preserves UUT's current log level
 * 2. **withOverrideLevel()** - Overrides UUT's log level
 */
class SharedLogRecorder : public LogRecorderBase
{
public:
    /**
     * @brief Create recorder that PRESERVES current UUT log level
     *
     * The UUT's log level is NOT changed. Only logs that pass the UUT's
     * macro-level filtering will be captured.
     */
    static SharedLogRecorder withCurrentLevel()
    {
        const hipdnnSeverity_t currentLevel = hipdnn_data_sdk::logging::getLogLevel();
        return {currentLevel, false}; // Don't change level
    }

    /**
     * @brief Create recorder that OVERRIDES UUT log level
     *
     * Use this when testing: "Does code contain this log statement?"
     *
     * @param overrideLevel Level to set UUT to (default: INFO = log everything)
     */
    static SharedLogRecorder withOverrideLevel(hipdnnSeverity_t overrideLevel = HIPDNN_SEV_INFO)
    {
        return {overrideLevel, true}; // Change level
    }

private:
    SharedLogRecorder(hipdnnSeverity_t level, bool shouldChange)
        : LogRecorderBase(level, shouldChange, LogRecording::Id::SHARED)
    {
    }
};

/**
 * @brief RAII log recorder for tests that register their own callback
 * and always restores the original UUT log level on destruction.
 *
 * Use this recorder in integration tests where you need to register a log recording
 * callback explicitly with the UUT. Logs are isolated to only this recorder.
 *
 * NOTE: There should be only one instance of the IsolatedLogRecorder per unit
 * test as all instances of the IsolatedLogRecorder will share the same log
 * recording buffer internally.
 *
 * This class provides a static IsolatedLogRecorder::getIsolatedRecordingCallback()
 * method to obtain the callback function for registration with the UUT's APIs.
 *
 * Two modes available:
 * 1. **withCurrentLevel()** - Preserves UUT's current log level
 * 2. **withOverrideLevel()** - Overrides UUT's log level
 */
class IsolatedLogRecorder : public LogRecorderBase
{
public:
    /**
     * @brief Create recorder that PRESERVES current UUT log level
     *
     * The UUT's log level is NOT changed. Only logs that pass the UUT's
     * macro-level filtering will be captured.
     */
    static IsolatedLogRecorder withCurrentLevel()
    {
        const hipdnnSeverity_t currentLevel = hipdnn_data_sdk::logging::getLogLevel();
        return {currentLevel, false}; // Don't change level
    }

    /**
     * @brief Create recorder that OVERRIDES UUT log level
     *
     * @param overrideLevel Level to set UUT to (default: INFO = log everything)
     */
    static IsolatedLogRecorder withOverrideLevel(hipdnnSeverity_t overrideLevel = HIPDNN_SEV_INFO)
    {
        return {overrideLevel, true}; // Change level
    }

    /**
     * @brief Get the callback function for registration with logging APIs
     *
     * Register this callback with the UUT's APIs to capture logs in the
     * ISOLATED recording instance.
     *
     * @return Callback function pointer suitable for logging API registration
     */
    static hipdnnCallback_t getIsolatedRecordingCallback()
    {
        return &isolatedRecordingCallback;
    }

    /**
     * @brief Get the user callback function for registration with user callback APIs
     *
     * Register this callback with hipdnnSetUserLogCallback_ext() or the frontend
     * setUserLogCallback() to capture logs in the ISOLATED recording instance.
     * The userHandle parameter is ignored — any non-null handle can be used.
     *
     * @return User callback function pointer suitable for hipdnnSetUserLogCallback_ext()
     */
    static hipdnnUserLogCallback_t getIsolatedUserRecordingCallback()
    {
        return &isolatedUserRecordingCallback;
    }

private:
    IsolatedLogRecorder(hipdnnSeverity_t level, bool shouldChange)
        : LogRecorderBase(level, shouldChange, LogRecording::Id::ISOLATED)
    {
    }

    static void isolatedRecordingCallback(hipdnnSeverity_t severity, const char* message)
    {
        if(message != nullptr)
        {
            LogRecording::instance(LogRecording::Id::ISOLATED)
                .recordLog(severity, std::string(message));
        }
    }

    static void isolatedUserRecordingCallback(hipdnnUserLogCallbackHandle_t /*userHandle*/,
                                              hipdnnSeverity_t severity,
                                              const char* message)
    {
        if(message != nullptr)
        {
            LogRecording::instance(LogRecording::Id::ISOLATED)
                .recordLog(severity, std::string(message));
        }
    }
};

/**
 * @brief Initialize test logging infrastructure for frontend unit tests
 *
 * This should be called from main() before RUN_ALL_TESTS() in unit test suites
 * that use the data SDK logger.
 *
 * This function will capture the current log level from HIPDNN_LOG_LEVEL env var
 * and will register the test recording callback, logRecordingCallback(), with the data SDK
 * logger. If the HIPDNN_LOG_LEVEL env var is defined and not set to off, then the test harness
 * will forward logs to chainedCallback, which defaults to simpleStderrOutputCallback which
 * output logs of the configured level using std::cerr.
 *
 * Any subsequent calls to hipdnn_data_sdk::logging::registerLoggingCallback()
 * after calling this function will replace the log recording callback, likely rendering
 * the test harness unable to record logs until hipdnn_data_sdk::logging::registerLoggingCallback()
 * is called again to set the data SDK logger to use the logRecordingCallback() callback (or
 * hipdnn_data_sdk::logging::registerLoggingCallback(logRecordingCallback) is called directly).
 */
inline hipdnnCallback_t initializeTestLogRecordingShared()
{
    // Initialize log level from environment
    hipdnn_data_sdk::logging::initializeLogLevel();

    // Capture test harness log level and set up log output filtering
    LogRecordingOutput::instance().initialize(hipdnn_data_sdk::logging::getLogLevel(),
                                              simpleStderrOutputCallback);

    // Register test recording callback
    hipdnn_data_sdk::logging::registerLoggingCallback(logRecordingCallback);

    return logRecordingCallback;
}

/**
 * @brief Always forwards messages to chainedCallback regardless of logging level
 * captured when logging was initialized in the test main(). Otherwise identical
 * to initializeTestLogRecordingShared() and should be called from main() before
 * RUN_ALL_TESTS() in unit test suites that use the data SDK logger.
 */
inline hipdnnCallback_t initializeChainedTestLogRecordingShared(hipdnnCallback_t chainedCallback)
{
    // Initialize log level from environment
    hipdnn_data_sdk::logging::initializeLogLevel();

    // Capture test harness log level and set up log output filtering
    LogRecordingOutput::instance().initialize(hipdnn_data_sdk::logging::getLogLevel(),
                                              chainedCallback);

    // Register test recording callback
    hipdnn_data_sdk::logging::registerLoggingCallback(logChainedRecordingCallback);

    return logChainedRecordingCallback;
}

} // namespace hipdnn_test_sdk::utilities
