#include "boottime_api/boottime_api.h"

#include "dbus_handler.hpp"
#include "gen.hpp"
#include "log.hpp"
#include "psm_handler.hpp"
#include "systemd_handler.hpp"

#include <fmt/printf.h>
#include <getopt.h> // For getopt_long

#include <boost/asio/io_service.hpp>
#include <sdbusplus/asio/connection.hpp>
#include <sdbusplus/asio/object_server.hpp>
#include <sdbusplus/bus.hpp>
#include <sdbusplus/server.hpp>

namespace btm = boot_time_monitor;

/*
 *  +-----------------+
 *  |      main       |
 *  +-----------------+
 *        |
 *        |
 *        v
 *  +-----------------+
 *  |   data_sources  |
 *  +-----------------+
 *        |   |   |
 *        |   |   psm_handler
 *        |   dbus_handler
 *        systemd_handler
 *        |
 *        | interface between data source and actual checkpoint/duration records
 *        v
 *  +-----------------+
 *  |      api        |
 *  +-----------------+
 *        |
 *        | store checkpoints/durations
 *        v
 *  +-----------------+
 *  |     files       |
 *  +-----------------+
 *
 */

constexpr int64_t kDefaultHostCount = 3; // Align with grpc-blobs default value
constexpr int64_t kDefaultBMCCount = 1;  // Default BMC count

void print_usage(const char* prog_name)
{
    fmt::print(stderr, "Usage: {} [options]\n", prog_name);
    fmt::print(stderr, "Options:\n");
    fmt::print(stderr,
               "  -h, --help             Show this help message and exit\n");
    fmt::print(
        stderr,
        "  --host_count <num>    Set number of hosts (optional, queries DBus if not set)\n");
    fmt::print(stderr,
               "  --bmc_count <num>     Set number of BMCs (default: 1)\n");
}

int main(int argc, char* argv[])
{
    int64_t hostCount = kDefaultHostCount;
    int64_t bmcCount = kDefaultBMCCount;

    const struct option long_options[] = {
        {"help", no_argument, nullptr, 'h'},
        {"host_count", required_argument, nullptr, 'o'}, // 'o' for hOst
        {"bmc_count", required_argument, nullptr, 'b'},  // 'b' for Bmc
        {nullptr, 0, nullptr, 0}};

    int opt;
    int option_index = 0;
    while ((opt = getopt_long(argc, argv, "ho:b:", long_options,
                              &option_index)) != -1)
    {
        switch (opt)
        {
            case 'h':
                print_usage(argv[0]);
                return 0;
            case 'o':
                try
                {
                    hostCount = std::stoi(optarg);
                    if (hostCount < 0)
                    {
                        fmt::print(
                            stderr,
                            "Warning: host_count should be non-negative. Using default {}.\n",
                            kDefaultHostCount);
                        hostCount = kDefaultHostCount;
                    }
                }
                catch (const std::exception& e)
                {
                    fmt::print(
                        stderr,
                        "Error parsing host_count: {}. Using default {}.\n",
                        e.what(), kDefaultHostCount);
                    hostCount = kDefaultHostCount;
                }
                break;
            case 'b':
                try
                {
                    bmcCount = std::stoi(optarg);
                    if (bmcCount <= 0)
                    {
                        fmt::print(
                            stderr,
                            "Warning: bmc_count must be positive. Using default {}.\n",
                            kDefaultBMCCount);
                        bmcCount = kDefaultBMCCount;
                    }
                    fmt::print(stdout, "Using provided bmc count: {}\n",
                               bmcCount);
                }
                catch (const std::exception& e)
                {
                    fmt::print(
                        stderr,
                        "Error parsing bmc_count: {}. Using default {}.\n",
                        e.what(), kDefaultBMCCount);
                    bmcCount = kDefaultBMCCount;
                }
                break;
            default: /* '?' */
                print_usage(argv[0]);
                return 1;
        }
    }

    fmt::print(stdout, "Host count: {}\n", hostCount);
    fmt::print(stdout, "BMC count: {}\n", bmcCount);

    boost::asio::io_service io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    conn->request_name("com.google.gbmc.boot_time_monitor");
    auto server = std::make_shared<sdbusplus::asio::object_server>(conn);
    auto& bus = static_cast<sdbusplus::bus::bus&>(*conn);

    std::vector<btm::NodeConfig> hostNodeConfigs;
    std::vector<btm::NodeConfig> bmcNodeConfigs;

    std::vector<std::unique_ptr<btm::dbus::Handler>> hostDbusHandlers;
    std::vector<std::unique_ptr<btm::dbus::Handler>> bmcDbusHandlers;

    std::vector<std::unique_ptr<btm::psm::Handler>> psmHandlers;

    std::shared_ptr<btm::api::BoottimeApi> api =
        std::make_shared<btm::api::BoottimeApi>();

    hostNodeConfigs.reserve(hostCount);
    for (int i = 0; i < hostCount; i++)
    {
        hostNodeConfigs.emplace_back(btm::gen::GenHostNodeName(i),
                                     std::string(btm::kBootTimeTagHost));
    }

    bmcNodeConfigs.reserve(bmcCount);
    for (int i = 0; i < bmcCount; i++)
    {
        bmcNodeConfigs.emplace_back(btm::gen::GenBmcNodeName(bmcCount, i),
                                    std::string(btm::kBootTimeTagBMC));
    }

    hostDbusHandlers.reserve(hostCount);
    psmHandlers.reserve(hostCount + 1); // extra psm handler for single host
    for (int i = 0; i < hostCount; i++)
    {
        auto& nodeConfig = hostNodeConfigs[i];
        absl::Status status = api->RegisterNode(nodeConfig);
        btm::log::LogIfError(status);

        hostDbusHandlers.emplace_back(std::make_unique<btm::dbus::Handler>(
            bus, nodeConfig, btm::gen::GenDbusConfig(nodeConfig), api));
        if (i == 0)
        {
            // Register an extra PSM handler for host 0, which is a common
            // convention for a single-host system. In a multi-host
            // configuration where this D-Bus service might not exist.
            psmHandlers.emplace_back(std::make_unique<btm::psm::Handler>(
                bus, hostNodeConfigs[0], btm::gen::GenPSMConfig(0), api));
        }
        psmHandlers.emplace_back(std::make_unique<btm::psm::Handler>(
            bus, nodeConfig, btm::gen::GenPSMConfig(i + 1),
            api)); // PSMConfig using 1-based index.
    }

    bmcDbusHandlers.reserve(bmcCount);
    for (auto& nodeConfig : bmcNodeConfigs)
    {
        absl::Status status = api->RegisterNode(nodeConfig);
        btm::log::LogIfError(status);

        bmcDbusHandlers.emplace_back(std::make_unique<btm::dbus::Handler>(
            bus, nodeConfig, btm::gen::GenDbusConfig(nodeConfig), api));
    }

    // Set up async systemd handler to update checkpoints/durations on this bmc.
    btm::systemd::Handler systemdHandler(conn, api, bmcNodeConfigs[0]);

    io.run();

    return 0;
}