#include "ssh_actions_plugin.h"

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "absl/strings/substitute.h"
#include "absl/synchronization/notification.h"
#include "absl/time/time.h"
#include <source_location>
#include "redfish_query_engine/http/codes.h"
#include "nlohmann/json.hpp"
#include "proxy.h"
#include "proxy_builder.h"
#include "proxy_config.pb.h"
#include "redfish_plugin.h"
#include "request_response.h"
#include "ssh_client.h"

namespace milotic {

using SshActionsConfig = ::milotic_grpc_proxy::Plugin::SshActions;

SshActionsPlugin::SshActionsPlugin(const SshActionsConfig& config)
    : common_command_options_({.timeout = absl::Seconds(10),
                               .retry_interval = absl::Milliseconds(100),
                               .use_pty = false}) {
  if (config.execute_timeout_sec() > 0) {
    common_command_options_.timeout =
        absl::Seconds(config.execute_timeout_sec());
  }
  if (config.execute_retry_interval_ms() > 0) {
    common_command_options_.retry_interval =
        absl::Milliseconds(config.execute_retry_interval_ms());
  }
  for (const SshActionsConfig::Action& action : config.actions()) {
    // sanity check
    if (action.args_count() >= 3 || action.args_count() < 0) {
      LOG(ERROR) << "Invalid 'args_count' configuration for command: '"
                 << action.ssh_command() << "', skipping.";
      continue;
    }

    actions_[action.action_path()] = {
        .command = action.ssh_command(),
        .timeout = absl::Seconds(action.timeout_sec()),
        .args_count = static_cast<int>(action.args_count())};
  }
}

absl::Status SshActionsPlugin::Initialize(Proxy* proxy) {
  ssh_client_ = proxy->GetResources().ssh_client.get();
  if (ssh_client_ == nullptr) {
    return absl::UnavailableError("SSH client not available");
  }
  proxy_ = proxy;
  return absl::OkStatus();
}

absl::string_view SshActionsPlugin::GetResourceId(absl::string_view url) {
  return url.substr(proxy_->GetNetworkEndpoint().length());
}

RedfishPlugin::RequestAction SshActionsPlugin::PreprocessRequest(
    RedfishPlugin::RequestVerb verb, ProxyRequest& request) {
  if (verb == RedfishPlugin::RequestVerb::kPost &&
      actions_.contains(GetResourceId(request.uri))) {
    VLOG(1) << "Handle " << request.uri;
    return RedfishPlugin::RequestAction::kHandle;
  }
  return RedfishPlugin::RequestAction::kNext;
}

namespace {
ProxyResponse ErrorResponse(
    const absl::Status& status,
    std::source_location location = std::source_location::current()) {
  LOG(ERROR).AtLocation(location.file_name(), static_cast<int>(location.line())) << status;
  return ProxyResponse{
      {.code = ecclesia::HttpResponseCode::HTTP_CODE_ERROR,
       .body = status.ToString(absl::StatusToStringMode::kWithNoExtraData),
       .headers = {{"Content-Type", "text/plain"}}}};
}
}  // namespace

absl::StatusOr<std::string> SshActionsPlugin::GetCommand(
    const SshActionsPlugin::Action& action,
    std::unique_ptr<ProxyRequest> request) {
  // Parse the optional command arguments from the request body.
  std::string command;
  auto obj = nlohmann::json::parse(request->body, nullptr,
                                   /* allow_exceptions */ false,
                                   /* ignore_comments */ false);

  auto args = obj.find("args");
  int args_count = 0;
  if (args != std::end(obj)) {
    if (args->is_array()) {
      args_count = static_cast<int>(args->size());
    } else {
      args_count = 1;
    }
  }

  if (args_count != action.args_count) {
    return absl::InvalidArgumentError(absl::StrFormat(
        "Got %d arguments, expected %d.", args_count, action.args_count));
  }

  switch (args_count) {
    case 0:
      command = action.command;
      break;

    case 1: {
      std::string* string_ptr = args->is_array()
                                    ? args->at(0).get_ptr<std::string*>()
                                    : args->get_ptr<std::string*>();
      if (string_ptr == nullptr) {
        return absl::InvalidArgumentError("Invalid 'args' argument type");
      }
      command = absl::Substitute(action.command, *string_ptr);
    } break;

    case 2: {
      std::string* string_ptr[2] = {args->at(0).get_ptr<std::string*>(),
                                    args->at(1).get_ptr<std::string*>()};
      if (string_ptr[0] == nullptr || string_ptr[1] == nullptr) {
        return absl::InvalidArgumentError("Invalid 'args' argument type");
      }
      command =
          absl::Substitute(action.command, *string_ptr[0], *string_ptr[1]);
    } break;

    default:
      return absl::InternalError(absl::StrFormat(
          "args_count value %d is not handled in code", action.args_count));
      break;
  }
  return command;
}

absl::StatusOr<ProxyResponse> SshActionsPlugin::HandleRequest(
    RedfishPlugin::RequestVerb verb, std::unique_ptr<ProxyRequest> request) {
  class Callbacks : public ExecutionCallbacks {
   public:
    virtual absl::Status OnStdout(ExecutionContext* context,
                                  absl::string_view data) {
      stdout_buffer_.append(data);
      return absl::OkStatus();
    }
    virtual absl::Status OnStderr(ExecutionContext* context,
                                  absl::string_view data) {
      stderr_buffer_.append(data);
      return absl::OkStatus();
    }
    virtual void OnDone(ExecutionContext* context, absl::StatusOr<int> result) {
      result_ = std::move(result);
      done_.Notify();
    }

    static absl::StatusOr<nlohmann::json> GetJsonOutput(
        Callbacks&& callbacks, absl::Duration timeout) {
      if (!callbacks.done_.WaitForNotificationWithTimeout(timeout)) {
        return absl::DeadlineExceededError(
            "Timeout waiting for command to complete");
      }
      if (!callbacks.result_.ok()) {
        LOG(ERROR) << "Failed to run command: " << callbacks.result_.status();
        return callbacks.result_.status();
      }
      return nlohmann::json{
          {"StdOut", std::move(callbacks.stdout_buffer_)},
          {"StdErr", std::move(callbacks.stderr_buffer_)},
          {"Result", *callbacks.result_},
      };
    }

   private:
    std::string stdout_buffer_;
    std::string stderr_buffer_;
    absl::StatusOr<int> result_;
    absl::Notification done_;
  } execution_callbacks;

  // `at` is safe to call. We have already validated that the action is in the
  // map in PreprocessRequest.
  const Action& action = actions_.at(GetResourceId(request->uri));

  auto command = GetCommand(action, std::move(request));
  if (!command.ok()) {
    return ErrorResponse(command.status());
  }

  absl::StatusOr<std::unique_ptr<ExecutionContext>> context =
      ssh_client_->ExecuteCommand(command.value(), &execution_callbacks,
                                  common_command_options_);
  if (!context.ok()) {
    return ErrorResponse(context.status());
  }
  absl::StatusOr<nlohmann::json> result =
      Callbacks::GetJsonOutput(std::move(execution_callbacks), action.timeout);
  if (!result.ok()) {
    return ErrorResponse(result.status());
  }
  return ProxyResponse(
      ecclesia::HttpResponseCode::HTTP_CODE_REQUEST_OK, result.value().dump(),
      {{"Content-Type", "application/json"}, {"OData-Version", "4.0"}});
}

REGISTER_REDFISH_PLUGIN(ssh_actions, SshActionsPlugin);

}  // namespace milotic