// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "hoth_updater_cli.hpp"

#include "google3/host_commands.h"

#include "message_util.hpp"

#include <boost/endian/arithmetic.hpp>
#include <sdbusplus/bus.hpp>
#include <sdbusplus/exception.hpp>
#include <sdbusplus/message.hpp>
#include <stdplus/print.hpp>
#include <stdplus/raw.hpp>

#include <chrono>
#include <cstdint>
#include <cstdio>
#include <exception>
#include <format>
#include <fstream>
#include <functional>
#include <iostream>
#include <iterator>
#include <span>
#include <stdexcept>
#include <string_view>
#include <thread>
#include <variant>
#include <vector>

using namespace std::chrono_literals;

namespace google::hoth::tools
{

namespace
{

// Normal update time:
// root@mvbbj12-nfd01:/tmp# time ./hoth_updater  --hoth_id=<HOTH_ID> update
// --image=<IMAGE_FILE_NAME>
// --address=0x01ef0000 real    0m8.952s user    0m0.290s sys     0m0.030s
constexpr auto kCallTimeout = 15s;
constexpr auto kRetryDelay = 30s;
constexpr auto kAttemptLimit = 3;

std::string getHothService(std::string_view hoth_id)
{
    std::string service = "xyz.openbmc_project.Control.Hoth";
    if (!hoth_id.empty() && hoth_id != "bmc")
    {
        service += ".";
        service += hoth_id;
    }
    return service;
}

std::string getHothObject(std::string_view hoth_id)
{
    std::string object = "/xyz/openbmc_project/Control/Hoth";
    if (!hoth_id.empty() && hoth_id != "bmc")
    {
        object += "/";
        object += hoth_id;
    }
    return object;
}

sdbusplus::message::message hothMessage(
    sdbusplus::bus::bus& bus, std::string_view hoth_id, const char* method)
{
    std::string service = getHothService(hoth_id);
    std::string object = getHothObject(hoth_id);
    return bus.new_method_call(service.c_str(), object.c_str(),
                               "xyz.openbmc_project.Control.Hoth", method);
}

sdbusplus::message::message hothPropertyMessage(sdbusplus::bus::bus& bus,
                                                std::string_view hoth_id)
{
    std::string service = getHothService(hoth_id);
    std::string object = getHothObject(hoth_id);
    return bus.new_method_call(service.c_str(), object.c_str(),
                               "org.freedesktop.DBus.Properties", "Get");
}

std::vector<std::string> getHothIds(const Args& args)
{
    std::vector<std::string> hoth_ids = splitString(args.hothId, ',');
    if (hoth_ids.empty())
    {
        hoth_ids.emplace_back("");
    }

    return hoth_ids;
}

template <typename T>
std::optional<T> getHothStateProperty(
    sdbusplus::bus::bus& bus, std::string_view hoth_id,
    std::string_view property,
    std::optional<sdbusplus::SdBusDuration> timeout = std::nullopt)
{
    try
    {
        sdbusplus::message::message msg = hothPropertyMessage(bus, hoth_id);
        msg.append("xyz.openbmc_project.Control.Hoth.State", property);
        std::variant<T> value{};
        sdbusplus::message::message resp =
            bus.call(msg, timeout.value_or(kCallTimeout));
        resp.read(value);
        return std::get<T>(value);
    }
    catch (const std::exception& ex)
    {
        return std::nullopt;
    }
}

template <typename T>
inline std::string optionalToString(const std::optional<T>& value)
{
    return value.has_value() ? std::format("{}", *value) : "n/a";
}

// This function will not print anything here. The logs will be printed
// to stderr of hothd.
void printHothLogsInHothd(
    sdbusplus::bus::bus& bus, std::string_view hoth_id,
    std::optional<sdbusplus::SdBusDuration> timeout = std::nullopt)
{
    sdbusplus::message::message msg =
        hothMessage(bus, hoth_id, "CollectHothLogs");
    bus.call(msg, timeout.value_or(kCallTimeout));
}

void emitPackageInstallationDetails(
    const std::vector<HothInstallStatus>& install_status,
    std::string_view notes)
{
    stdplus::print(stdout, "firmware_status {{\n");
    for (const HothInstallStatus& dev_install : install_status)
    {
        stdplus::print(stdout, "  device_status {{\n");
        stdplus::print(stdout, "    devpath: \"{}\"\n", dev_install.devpath);
        stdplus::print(stdout, "    failed: {}\n",
                       dev_install.failed ? "true" : "false");
        stdplus::print(stdout,
                       "    machine_readable_failure_symptom_key: \"{}\"\n",
                       dev_install.failure_symptom_key);
        stdplus::print(stdout, "    failure: \"{}\"\n", dev_install.failure);
        stdplus::print(stdout, "  }}\n");
    }

    stdplus::print(stdout, "}}\n");
    stdplus::print(stdout, "notes: \"{}\"\n", notes);
}

} // namespace

std::string HothUpdaterCLI::getFailureSymptomKey(sdbusplus::bus::bus& bus,
                                                 std::string_view hoth_id)
{
    std::string failure_symptom;

    try
    {
        // Test the transport by issuing a simple host command
        (void)getHothVersion(bus, hoth_id);

        // If we reach here, then we were able to query the
        // hoth version successfully and the transport then
        // must be working fine.  This must be some other errors.
        // TODO(sportley): We should have a better solution for detecting
        // these failures.
        failure_symptom = "rot-install-generic-failure";
        printHothLogsInHothd(bus, hoth_id);
    }
    catch (const sdbusplus::exception::SdBusError& ex)
    {
        if (std::string_view(ex.name()) ==
            "xyz.openbmc_project.Control.Hoth.Error.InterfaceError")
        {
            auto transport = getTransportType(bus, hoth_id);
            if (transport.has_value())
            {
                failure_symptom =
                    std::format("rot-install-{}-failure", transport.value());
            }
            else
            {
                failure_symptom = std::format("rot-install-generic-failure");
            }
        }
        else
        {
            failure_symptom = std::format("rot-install-generic-failure");
        }
    }
    catch (...)
    {
        failure_symptom = std::format("rot-install-generic-failure");
    }

    return failure_symptom;
}

std::vector<uint8_t> HothUpdaterCLI::sendHostCommand(
    sdbusplus::bus::bus& bus, std::string_view hoth_id,
    const std::span<const uint8_t> command)
{
    sdbusplus::message::message msg =
        hothMessage(bus, hoth_id, "SendHostCommand");
    msg.append(command);
    sdbusplus::message::message resp = bus.call(msg, kCallTimeout);
    std::vector<uint8_t> result;
    resp.read(result);
    return result;
}

void HothUpdaterCLI::updateFirmware(sdbusplus::bus::bus& bus,
                                    std::string_view hoth_id,
                                    const std::span<const uint8_t> image)
{
    sdbusplus::message::message msg =
        hothMessage(bus, hoth_id, "UpdateFirmware");
    msg.append(image);
    stdplus::print(stderr, "Sending update firmare request\n");
    auto start_time = std::chrono::steady_clock::now();
    try
    {
        bus.call(msg, kCallTimeout);
    }
    catch (const std::exception& ex)
    {
        stdplus::print(stderr, "Exception caught during updateFirmware: {}\n",
                       ex.what());
        auto end_time = std::chrono::steady_clock::now();
        auto duration = end_time - start_time;
        auto milliseconds =
            std::chrono::duration_cast<std::chrono::milliseconds>(duration)
                .count();
        stdplus::print(stderr, "Time spent in bus call: {}ms\n", milliseconds);
        throw;
    }
    auto end_time = std::chrono::steady_clock::now();
    auto duration = end_time - start_time;
    auto milliseconds =
        std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();
    stdplus::print(stderr, "Time spent in bus call: {}ms\n", milliseconds);
}

void HothUpdaterCLI::spiWrite(
    sdbusplus::bus::bus& bus, std::string_view hoth_id,
    const std::span<const uint8_t> image, std::optional<uint32_t> address)
{
    for (int attempt = 0; attempt < kAttemptLimit; attempt++)
    {
        try
        {
            sdbusplus::message::message msg =
                hothMessage(bus, hoth_id, "SpiWrite");
            if (address)
            {
                msg.append(*address); // u
                msg.append(image);    // ay
            }
            bus.call(msg, kCallTimeout);
            return;
        }
        catch (const std::exception& ex)
        {
            stdplus::print(stderr, "Exeption caught: {}\n", ex.what());
            stdplus::print(stderr, "Will rety in {} seconds\n",
                           kRetryDelay.count());
            std::this_thread::sleep_for(kRetryDelay);
        }
    }
    throw std::runtime_error("Retry attempt limit exhausted.");
}

FirmwareUpdateStatus HothUpdaterCLI::getFirmwareUpdateStatus(
    sdbusplus::bus::bus& bus, std::string_view hoth_id)
{
    sdbusplus::message::message msg =
        hothMessage(bus, hoth_id, "GetFirmwareUpdateStatus");
    auto start_time = std::chrono::steady_clock::now();
    try
    {
        sdbusplus::message::message reply = bus.call(msg, kCallTimeout);
        std::string rsp;
        reply.read(rsp);
        return sdbusplus::xyz::openbmc_project::Control::server::Hoth::
            convertFirmwareUpdateStatusFromString(rsp);
    }
    catch (const std::exception& ex)
    {
        stdplus::print(stderr,
                       "Exception caught during getFirmwareUpdateStatus: {}\n",
                       ex.what());
        auto end_time = std::chrono::steady_clock::now();
        auto duration = end_time - start_time;
        auto milliseconds =
            std::chrono::duration_cast<std::chrono::milliseconds>(duration)
                .count();
        stdplus::print(stderr, "Time spent in bus call: {}ms\n", milliseconds);
        throw;
    }
}

std::vector<uint8_t> HothUpdaterCLI::readFileIntoByteArray(
    std::string_view filename)
{
    std::ifstream image_file;
    image_file.exceptions(std::ios::failbit);
    image_file.open(std::string(filename), std::ios::binary);
    std::vector<uint8_t> image(std::istreambuf_iterator<char>(image_file), {});
    return image;
}

void HothUpdaterCLI::doUpdate(const Args& args)
{
    sdbusplus::bus::bus bus = sdbusplus::bus::new_default();
    doUpdateLogic(bus, args);
}

// TODO(sportley): Refactor this to take something besides Args.
//  It's convenient for plumbing through to spiWrite + updateFirmware
//  but runs the risk of using args.hothId instead of hoth_id.
bool HothUpdaterCLI::startUpdate(sdbusplus::bus::bus& bus,
                                 std::string_view hoth_id, const Args& args,
                                 const std::vector<uint8_t>& image)
{
    bool ret = true;
    try
    {
        if (args.updateMethod == "spi")
        {
            std::string spi_msg;
            if (args.address)
            {
                spi_msg += std::format("at 0x{:x}", *args.address);
            }
            stdplus::print(stderr, "Writing to SPI {}, hoth_id: {}\n", spi_msg,
                           hoth_id);
            spiWrite(bus, hoth_id, image, args.address);
        }
        else if (args.updateMethod == "update_firmware")
        {
            stdplus::print(stderr, "Sending 'update_firmware to hoth_id: {}\n",
                           hoth_id);
            updateFirmware(bus, hoth_id, image);
        }
        else
        {
            ret = false;
            stdplus::print(
                stderr,
                "Invalid update method, only 'spi' and 'update_firmware' are "
                "supported.");
        }
    }
    catch (const std::exception& ex)
    {
        stdplus::print(
            stderr, "Caught dBus level exception when performing update: {}\n",
            ex.what());
        ret = false;
    }

    return ret;
}

void HothUpdaterCLI::doUpdateLogic(sdbusplus::bus::bus& bus, const Args& args)
{
    std::vector<uint8_t> image = readFileIntoByteArray(args.imageFilename);

    std::vector<HothInstallStatus> install_status;
    auto end_time = std::chrono::steady_clock::now() + 5min;

    std::string notes = std::format("updateMethod={}", args.updateMethod);
    notes += std::format(" | hothId={}", args.hothId);

    const std::vector<std::string> hoth_ids = getHothIds(args);

    // Start the update for every hothd
    for (const auto& hoth_id : hoth_ids)
    {
        auto status = HothInstallStatus{
            .failed = false,
            .devpath = "",
            .failure_symptom_key = "",
            .failure = "",
        };

        if (!startUpdate(bus, hoth_id, args, image))
        {
            status.failed = true;
            status.failure_symptom_key = "rot-hothd-dbus-failure";
            status.failure =
                "dbus failure when starting update.  Is hothd running?";
        }

        install_status.push_back(status);
    }

    // Start the update for every hothd instance, we run
    // these all in parallel for performance.
    while (std::chrono::steady_clock::now() < end_time)
    {
        bool done = true;
        for (std::size_t i = 0; i < hoth_ids.size(); i++)
        {
            if (install_status[i].failed)
            {
                // We already know that this install failed,
                // no sense in checking the status again...
                continue;
            }

            auto status = FirmwareUpdateStatus::Error;
            try
            {
                status = getFirmwareUpdateStatus(bus, hoth_ids[i]);
            }
            catch (const std::exception& ex)
            {
                install_status[i].failed = true;
                install_status[i].failure_symptom_key =
                    "rot-hothd-dbus-failure";
                install_status[i].failure = std::format(
                    "dBus exception when fetching update status: {}",
                    ex.what());
            }

            if (status == FirmwareUpdateStatus::InProgress)
            {
                // Still waiting
                done = false;
                continue;
            }

            if (status != FirmwareUpdateStatus::Done)
            {
                install_status[i].failed = true;
                install_status[i].failure_symptom_key =
                    getFailureSymptomKey(bus, hoth_ids[i]);
                install_status[i].failure =
                    "hothd is responding, but signaled a firmware update failure.";
            }
        }

        if (done)
        {
            break;
        }

        std::this_thread::sleep_for(100ms);
    }

    emitPackageInstallationDetails(install_status, notes);

    for (const auto& status : install_status)
    {
        if (status.failed)
        {
            throw std::runtime_error("Update failed");
        }
    }
}

std::optional<std::string> HothUpdaterCLI::getTransportType(
    sdbusplus::bus::bus& bus, std::string_view hoth_id)
{
    return getHothStateProperty<std::string>(bus, hoth_id, "TransportType");
}

HothVersionStringsRsp HothUpdaterCLI::getHothVersion(sdbusplus::bus::bus& bus,
                                                     std::string_view hoth_id)
{
    google::hoth::internal::ReqHeader header;
    google::hoth::internal::populateReqHeader(EC_CMD_GET_VERSION, 0, nullptr, 0,
                                              &header);

    auto header_span = stdplus::raw::asSpan<uint8_t>(header);
    std::vector<uint8_t> command(header_span.begin(), header_span.end());
    std::vector<uint8_t> resp_bytes = sendHostCommand(bus, hoth_id, command);

    auto response = stdplus::raw::copyFrom<HothVersionStringsRsp>(resp_bytes);

    if (response.header.result != 0)
    {
        throw std::runtime_error("Failed to get versions");
    }

    return response;
}

void HothUpdaterCLI::doFirmwareVersion(const Args& args)
{
    sdbusplus::bus::bus bus = sdbusplus::bus::new_default();
    doFirmwareVersionLogic(bus, args);
}

void HothUpdaterCLI::doFirmwareVersionLogic(sdbusplus::bus::bus& bus,
                                            const Args& args)
{
    auto response = getHothVersion(bus, args.hothId);

    if (args.ro)
    {
        std::cout << response.version.version_string_ro;
    }
    else
    {
        std::cout << response.version.version_string_rw;
    }
}

std::vector<std::string> splitString(const std::string& s, const char delim)
{
    size_t idx = 0;
    std::vector<std::string> ret;
    while (idx < s.size())
    {
        std::size_t delim_pos = s.find_first_of(delim, idx);
        if (delim_pos == std::string::npos)
        {
            break;
        }
        ret.emplace_back(s.substr(idx, delim_pos - idx));
        idx = delim_pos + 1;
    }
    if (idx < s.size())
    {
        ret.emplace_back(s.substr(idx));
    }
    return ret;
}

HothActivationStatistics HothUpdaterCLI::getHothActivationStatistics(
    sdbusplus::bus::bus& bus, std::string_view hoth_id)
{
    HothActivationStatistics status;

    // Try query the statistics properties on dBus.
    // Each individual property could fail. Ignore errors and try populate as
    // much as possible.
    status.rw_failure_code = getHothStateProperty<uint32_t>(
        bus, hoth_id, "FirmwareUpdateFailureCode");
    status.rw_failed_minor = getHothStateProperty<uint32_t>(
        bus, hoth_id, "FirmwareUpdateFailedMinor");
    status.ro_failure_code = getHothStateProperty<uint32_t>(
        bus, hoth_id, "BootloaderUpdateFailureCode");
    status.reset_flags =
        getHothStateProperty<uint32_t>(bus, hoth_id, "ResetFlags");
    status.uptime_us = getHothStateProperty<uint64_t>(bus, hoth_id, "UpTime");

    return status;
}

void HothUpdaterCLI::doActivationCheck(const Args& args)
{
    sdbusplus::bus::bus bus = sdbusplus::bus::new_default();
    doActivationCheckLogic(bus, args);
}

void HothUpdaterCLI::doActivationCheckLogic(sdbusplus::bus::bus& bus,
                                            const Args& args)
{
    stdplus::print(stdout, "installed_version: \"{}\"\n",
                   args.expectedRwVersion);

    std::vector<std::string> hoth_ids = getHothIds(args);

    for (const std::string& hoth_id : hoth_ids)
    {
        auto response = getHothVersion(bus, hoth_id);

        // Strip out the hoth family in the version string, because it's not
        // guaranteed to be returned completely.
        std::string_view version_stripped(response.version.version_string_rw);
        size_t split_pos = version_stripped.find_first_of('/');
        if (split_pos != std::string::npos)
        {
            version_stripped = version_stripped.substr(0, split_pos);
        }

        stdplus::print(
            stdout,
            "activated_versions {{\n"
            "    key: \"{}\"\n"
            "    value: \"{}\"\n"
            "}}\n",
            (hoth_id.empty() ? "active" : hoth_id), version_stripped);

        if (args.expectedRwVersion != version_stripped)
        {
            HothActivationStatistics actv_status =
                getHothActivationStatistics(bus, hoth_id);
            stdplus::print(stdout,
                           "notes: \"Running RW version {}"
                           " does not match expected version {}, Status: rw={} "
                           "(ver={}) ro={} rst_flags={}, uptime={}us\"\n",
                           version_stripped, args.expectedRwVersion,
                           optionalToString(actv_status.rw_failure_code),
                           optionalToString(actv_status.rw_failed_minor),
                           optionalToString(actv_status.ro_failure_code),
                           optionalToString(actv_status.reset_flags),
                           optionalToString(actv_status.uptime_us));
            throw std::runtime_error("Activation check failed");
        }
    }
}

void HothUpdaterCLI::doEnableSecurebootEnforcement(const Args& args)
{
    sdbusplus::bus::bus bus = sdbusplus::bus::new_default();
    doEnableSecurebootEnforcementLogic(bus, args);
}

void HothUpdaterCLI::doEnableSecurebootEnforcementLogic(
    sdbusplus::bus::bus& bus, const Args& args)
{
    google::hoth::internal::ReqHeader header;
    secure_boot_enforcement_state payload = {.enabled = 1, .reserved0 = {0}};
    google::hoth::internal::populateReqHeader(
        EC_CMD_BOARD_SPECIFIC_BASE +
            EC_PRV_CMD_HOTH_SET_SECURE_BOOT_ENFORCEMENT,
        0, &payload, sizeof(payload), &header);

    auto header_span = stdplus::raw::asSpan<uint8_t>(header);
    auto payload_span = stdplus::raw::asSpan<uint8_t>(payload);
    std::vector<uint8_t> command;
    command.insert(command.end(), header_span.begin(), header_span.end());
    command.insert(command.end(), payload_span.begin(), payload_span.end());

    const std::vector<std::string> hoth_ids = getHothIds(args);

    for (const std::string& hoth_id : hoth_ids)
    {
        std::vector<uint8_t> resp_bytes =
            sendHostCommand(bus, hoth_id, command);

        auto response =
            stdplus::raw::copyFrom<google::hoth::internal::RspHeader>(
                resp_bytes);

        if (response.result != 0)
        {
            throw std::runtime_error(
                "Failed to enable secure boot enforcement");
        }
    }
}

void setupCLIApp(CLI::App& app, HothUpdaterCLI& cli, Args& args)
{
    app.require_subcommand(1);
    app.add_option("--hoth_id", args.hothId,
                   "Hoth IDs, comma-delimited for activation-check");
    auto* update = app.add_subcommand("update", "Update Hoth image");
    update->add_option("--image", args.imageFilename, "Firmware image path")
        ->required()
        ->check(CLI::ExistingFile);
    update
        ->add_option("--method", args.updateMethod,
                     "Update method, can be spi|update_firmware")
        ->required();
    update->add_option("--address", args.address, "SPI address for data")
        ->check(CLI::NonNegativeNumber);
    update->callback([&args, &cli] { cli.doUpdate(args); });

    auto* version =
        app.add_subcommand("firmware-version", "Get firmware version");

    version->add_flag("--ro,!--rw", args.ro, "Select ro/rw partition")
        ->required();
    version->callback([&args, &cli] { cli.doFirmwareVersion(args); });

    auto* activation_check = app.add_subcommand(
        "activation-check", "Generate package activation check textproto");
    activation_check->add_flag("--rw", args.expectedRwVersion,
                               "Expected RW version");
    // RO version check not implemented yet.
    activation_check->callback([&args, &cli] { cli.doActivationCheck(args); });

    auto* enable_secureboot = app.add_subcommand(
        "enable-secureboot-enforcement", "Enable secure boot enforcement");
    enable_secureboot->callback([&args, &cli] {
        cli.doEnableSecurebootEnforcement(args);
    });
}

} // namespace google::hoth::tools