#include "systemd_handler.hpp"

#include "log.hpp" // provide btm::log::LogIfError

#include <fmt/printf.h>

#include <boost/interprocess/file_mapping.hpp>
#include <boost/interprocess/mapped_region.hpp>

#include <fstream>
#include <regex>

namespace boot_time_monitor
{
namespace systemd
{

namespace btm = boot_time_monitor;

constexpr std::string_view kSystemdFirmwareTime = "FirmwareTimestampMonotonic";
constexpr std::string_view kSystemdLoaderTime = "LoaderTimestampMonotonic";
constexpr std::string_view kSystemdKernelTime = "KernelTimestampMonotonic";
constexpr std::string_view kSystemdInitRDTime = "InitRDTimestampMonotonic";
constexpr std::string_view kSystemdUserspaceTime =
    "UserspaceTimestampMonotonic";
constexpr std::string_view kSystemdFinishTime = "FinishTimestampMonotonic";

struct SystemdTimestamps
{
    uint64_t firmware = 0;
    uint64_t loader = 0;
    uint64_t kernel = 0;
    uint64_t initRD = 0;
    uint64_t userspace = 0;
    uint64_t finish = 0;
};

// Support negative duration if needed.
struct SystemdDurations
{
    int64_t firmware = 0;
    int64_t loader = 0;
    int64_t kernel = 0;
    int64_t initRD = 0;
    int64_t userspace = 0;
};

namespace
{
// --- Helper Functions ---

/**
 * @brief Calculates boot stage durations from systemd timestamps.
 * @param times A struct containing all collected monotonic timestamps.
 * @return A struct containing the calculated duration for each stage in ms.
 *
 * The calculation logic is based on systemd-analyze:
 * https://github.com/systemd/systemd/blob/82b7bf8c1c8c6ded6f56b43998c803843a3b944b/src/analyze/analyze-time-data.c#L176-L186
 * Special calculation for kernel duration:
 * https://github.com/systemd/systemd/blob/82b7bf8c1c8c6ded6f56b43998c803843a3b944b/src/analyze/analyze-time-data.c#L84-L87
 */
SystemdDurations CalculateSystemdDuration(const SystemdTimestamps& times)
{
    // Durations are calculated by subtracting the start times of subsequent
    // stages. Cast to signed to prevent underflow if timestamps are out of
    // order.
    auto to_ms = [](int64_t us) { return us > 0 ? us / 1000 : 0; };

    int64_t firmware_us = static_cast<int64_t>(times.firmware);
    int64_t loader_us = static_cast<int64_t>(times.loader);
    int64_t initRD_us = static_cast<int64_t>(times.initRD);
    int64_t userspace_us = static_cast<int64_t>(times.userspace);
    int64_t finish_us = static_cast<int64_t>(times.finish);

    SystemdDurations durations;
    durations.firmware = to_ms(firmware_us - loader_us);
    durations.loader = to_ms(loader_us);
    // Kernel time ends when initrd is done, or when userspace starts if no
    // initrd.
    durations.kernel = (initRD_us > 0) ? to_ms(initRD_us) : to_ms(userspace_us);
    durations.initRD = (initRD_us > 0) ? to_ms(userspace_us - initRD_us) : 0;
    durations.userspace = to_ms(finish_us - userspace_us);

    return durations;
}

/**
 * @brief Removes the "TimestampMonotonic" suffix from a property name.
 * @param name The full D-Bus property name.
 * @return The base name for the boot stage.
 */
std::string GetStageName(std::string_view name)
{
    constexpr std::string_view suffix = "TimestampMonotonic";
    if (name.ends_with(suffix))
    {
        name.remove_suffix(suffix.length());
    }
    return std::string(name);
}

} // namespace

// --- Class Implementation ---

Handler::Handler(std::shared_ptr<sdbusplus::asio::connection> conn,
                 std::shared_ptr<btm::api::IBoottimeApi> api,
                 btm::NodeConfig bmcNode) :
    mApi(std::move(api)), mBMCFinishTimer(conn->get_io_context()),
    mDbusConn(std::move(conn)), mNode(std::move(bmcNode))
{
    // On startup, check if the system has already finished booting.
    if (auto finishTime = QuerySystemdTimestamp(kSystemdFinishTime);
        finishTime && *finishTime > 0)
    {
        // If the system is booted, ensure any pending reboot state is cleared.
        if (mApi->IsNodeRebooting(mNode))
        {
            fmt::print(
                stderr,
                "System already booted; marking reboot as complete for '{}'\n",
                mNode.node_name);
            btm::log::LogIfError(mApi->NotifyNodeComplete(mNode));
        }
        else
        {
            fmt::print(stderr, "System already booted. No action needed.\n");
        }
        return;
    }

    // If not booted, start polling for boot completion.
    fmt::print(stderr,
               "System not yet booted. Waiting for systemd to finish...\n");
    mBMCFinishTimer.expires_at(std::chrono::steady_clock::now());
    mBMCFinishTimer.async_wait(
        [this](const auto& ec) { UpdateBmcCheckpoints(ec); });
}

std::optional<uint64_t>
    Handler::QuerySystemdTimestamp(std::string_view propertyName)
{
    auto method = mDbusConn->new_method_call(
        "org.freedesktop.systemd1",        // Systemd Service
        "/org/freedesktop/systemd1",       // Systemd Path
        "org.freedesktop.DBus.Properties", // Systemd Iface
        "Get");                            // Systemd Func
    method.append("org.freedesktop.systemd1.Manager", propertyName);

    try
    {
        auto reply = mDbusConn->call(method);
        std::variant<uint64_t> result;
        reply.read(result);
        return std::get<uint64_t>(result);
    }
    catch (const sdbusplus::exception::SdBusError& e)
    {
        fmt::print(stderr, "Failed to get systemd property '{}': {}\n",
                   propertyName, e.what());
        return std::nullopt;
    }
}

void Handler::RestartPolling()
{
    mBMCFinishTimer.expires_after(kCheckBmcFinishInterval);
    mBMCFinishTimer.async_wait(
        [this](const auto& ec) { UpdateBmcCheckpoints(ec); });
}

void Handler::UpdateBmcCheckpoints(const boost::system::error_code& ec)
{
    if (ec)
    {
        if (ec != boost::asio::error::operation_aborted)
        {
            fmt::print(stderr, "BMC checkpoint timer error: {}\n",
                       ec.message());
        }
        return;
    }

    // 1. Check if systemd has finished booting. If not, poll again.
    std::optional<uint64_t> finishTimestamp =
        QuerySystemdTimestamp(kSystemdFinishTime);
    if (!finishTimestamp || *finishTimestamp == 0)
    {
        RestartPolling();
        return;
    }

    // 2. Systemd is ready. Query all other timestamps.
    SystemdTimestamps times;
    times.finish = *finishTimestamp;

    auto queryAndAssign = [&](const std::string_view name,
                              uint64_t& destination) -> bool {
        if (auto val = QuerySystemdTimestamp(name))
        {
            destination = *val;
            return true;
        }
        return false;
    };
    if (!queryAndAssign(kSystemdFirmwareTime, times.firmware))
    {
        RestartPolling();
        return;
    }
    if (!queryAndAssign(kSystemdLoaderTime, times.loader))
    {
        RestartPolling();
        return;
    }
    if (!queryAndAssign(kSystemdKernelTime, times.kernel))
    {
        RestartPolling();
        return;
    }
    if (!queryAndAssign(kSystemdInitRDTime, times.initRD))
    {
        RestartPolling();
        return;
    }
    if (!queryAndAssign(kSystemdUserspaceTime, times.userspace))
    {
        RestartPolling();
        return;
    }

    fmt::print(stderr, "Successfully collected all systemd boot timestamps.\n");

    // 3. All data is collected. Calculate and report the durations.
    SystemdDurations durationMap = CalculateSystemdDuration(times);

#ifdef NPCM7XX_OR_NEWER
    // NPCM7XX or newer Nuvoton BMC has a register that starts counting from
    // SoC power on. Also uptime starts counting when kernel is up. Thus we
    // can get (Firmware + Loader) time by `value[SEC_CNT_ADDR] - uptime`.
    constexpr uint32_t SEC_CNT_ADDR = 0xF0801068;
    if (auto powerOnSec = readMem4Bytes(SEC_CNT_ADDR); powerOnSec)
    {
        if (auto upTimeMS = getUpTimeInMs(); upTimeMS)
        {
            // This register provides a more accurate power-on to kernel-up
            // time.
            durationMap.firmware = (powerOnSec.value() * 1000) -
                                   upTimeMS.value();
            durationMap.loader = 0; // The combined time is stored in firmware.
        }
    }
#endif

    // 4. Report the calculated durations.
    btm::log::LogIfError(mApi->SetNodeDuration(
        mNode, GetStageName(kSystemdFirmwareTime), durationMap.firmware));
    btm::log::LogIfError(mApi->SetNodeDuration(
        mNode, GetStageName(kSystemdLoaderTime), durationMap.loader));
    btm::log::LogIfError(mApi->SetNodeDuration(
        mNode, GetStageName(kSystemdKernelTime), durationMap.kernel));
    btm::log::LogIfError(mApi->SetNodeDuration(
        mNode, GetStageName(kSystemdInitRDTime), durationMap.initRD));
    btm::log::LogIfError(mApi->SetNodeDuration(
        mNode, GetStageName(kSystemdUserspaceTime), durationMap.userspace));

    // 5. Finalize the boot process.
    btm::log::LogIfError(mApi->SetNodeCheckpoint(mNode, "UserspaceEnd", 0, 0));
    btm::log::LogIfError(mApi->NotifyNodeComplete(mNode));
}

#ifdef NPCM7XX_OR_NEWER
std::optional<uint32_t> Handler::readMem4Bytes(uint32_t target)
{
    // Typically the pageSize will be 4K/8K for 32 bit operating systems
    uint32_t pageSize = boost::interprocess::mapped_region::get_page_size();
    const uint32_t kPageMask = pageSize - 1;

    uint32_t pageOffset = target & (~kPageMask);
    uint32_t offsetInPage = target & kPageMask;

    // Map `/dev/mem` to a region.
    std::unique_ptr<boost::interprocess::file_mapping> fileMapping;
    std::unique_ptr<boost::interprocess::mapped_region> MappedRegion;
    try
    {
        fileMapping = std::make_unique<boost::interprocess::file_mapping>(
            "/dev/mem", boost::interprocess::read_only);
        // No need to unmap in the end.
        MappedRegion = std::make_unique<boost::interprocess::mapped_region>(
            *fileMapping, boost::interprocess::read_only, pageOffset,
            pageSize * 2);

        // MappedRegion->get_address() returns (void*) which needs to covert
        // into (char*) to make `+ offsetInPage` work.
        // Then converts it again into (uint32_t*) since we read 4 bytes.
        return *(reinterpret_cast<uint32_t*>(
            static_cast<char*>(MappedRegion->get_address()) + offsetInPage));
    }
    catch (const std::exception& e)
    {
        fmt::print(stderr, "Error reading /dev/mem at address {:#x}: {}\n",
                   target, e.what());
        return std::nullopt;
    }

    fmt::print(stderr, "[{}]: Shouldn't go to this line.\n", __FUNCTION__);
    return std::nullopt;
}

std::optional<int64_t> Handler::getUpTimeInMs()
{
    std::ifstream fin("/proc/uptime");
    if (!fin.is_open())
    {
        fmt::print(stderr,
                   "Failed to open /proc/uptime to read system uptime.\n");
        return std::nullopt;
    }
    double uptimeSec = 0.0;
    fin >> uptimeSec;
    if (fin.fail())
    {
        fmt::print(stderr,
                   "Failed to parse system uptime from /proc/uptime.\n");
        return std::nullopt;
    }

    return static_cast<int64_t>(uptimeSec * 1000);
}
#endif // NPCM7XX_OR_NEWER

} // namespace systemd
} // namespace boot_time_monitor