/* Copyright 2017 - 2026 R. Thomas
 * Copyright 2017 - 2026 Quarkslab
 *
 * 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 <sstream>
#include "spdlog/fmt/fmt.h"

#include "LIEF/BinaryStream/SpanStream.hpp"

#include "LIEF/MachO/FunctionVariants.hpp"

#include "MachO/Structures.hpp"

#include "logging.hpp"
#include "internal_utils.hpp"


namespace LIEF::MachO {

static_assert(sizeof(details::runtime_table_entry_t) == 8);

using FLAGS = FunctionVariants::RuntimeTableEntry::FLAGS;
using KIND = FunctionVariants::RuntimeTable::KIND;

static const char* pretty_name(FLAGS f) {
  switch (f) {

// NOLINTNEXTLINE(cppcoreguidelines-macro-usage)
#define FUNCTION_VARIANT_FLAG(name, _, pretty)                                    \
  case FLAGS::name: return pretty;
#include "LIEF/MachO/FunctionVariants/Arm64.def"
#include "LIEF/MachO/FunctionVariants/PerProcess.def"
#include "LIEF/MachO/FunctionVariants/SystemWide.def"
#include "LIEF/MachO/FunctionVariants/X86_64.def"
#undef FUNCTION_VARIANT_FLAG
    case FLAGS::UNKNOWN: return "UNKNOWN";
  }
  return "UNKNOWN";
}

FunctionVariants::FunctionVariants(const details::linkedit_data_command& cmd) :
  LoadCommand::LoadCommand{LoadCommand::TYPE(cmd.cmd), cmd.cmdsize},
  data_offset_{cmd.dataoff},
  data_size_{cmd.datasize} {}

std::ostream& FunctionVariants::print(std::ostream& os) const {
  LoadCommand::print(os) << '\n';
  auto runtime_tbl = runtime_table();

  for (size_t i = 0; i < runtime_tbl.size(); ++i) {
    const RuntimeTable& entry = runtime_tbl[i];
    os << fmt::format("Table #{}\n", i) << indent(entry.to_string(), 2) << '\n';
  }

  return os;
}

FunctionVariants::RuntimeTableEntry::RuntimeTableEntry(
    const details::runtime_table_entry_t& entry
) :
  another_table_(entry.another_table),
  impl_(entry.impl) {
  flag_bit_nums_.fill(0);
  std::copy(std::begin(entry.flag_bit_nums), std::end(entry.flag_bit_nums),
            std::begin(flag_bit_nums_));
}

span<const uint8_t> get_relevant_flags(span<const uint8_t> raw) {
  auto it = std::find_if(raw.begin(), raw.end(), [](uint8_t v) { return v == 0; });
  if (it == raw.end()) {
    return raw;
  }
  size_t distance = std::distance(raw.begin(), it);

  return raw.subspan(0, distance);
}

std::vector<FLAGS> get_high_level_flags(span<const uint8_t> flags, KIND kind) {
  using RuntimeTableEntry = FunctionVariants::RuntimeTableEntry;

  std::vector<FLAGS> out;
  out.reserve(flags.size());

  for (uint8_t flag : flags) {
    switch (kind) {
      case KIND::UNKNOWN: out.push_back((FLAGS)flag); break;

      case KIND::PER_PROCESS:
        out.push_back((FLAGS)(flag | RuntimeTableEntry::F_PER_PROCESS));
        break;

      case KIND::SYSTEM_WIDE:
        out.push_back((FLAGS)(flag | RuntimeTableEntry::F_SYSTEM_WIDE));
        break;

      case KIND::ARM64:
        out.push_back((FLAGS)(flag | RuntimeTableEntry::F_ARM64));
        break;

      case KIND::X86_64:
        out.push_back((FLAGS)(flag | RuntimeTableEntry::F_X86_64));
        break;
    }
  }
  return out;
}

std::string FunctionVariants::RuntimeTableEntry::to_string() const {
  std::ostringstream oss;
  if (another_table()) {
    oss << fmt::format("Table: #{}", impl());
  } else {
    oss << fmt::format("Function: {:#010x}", impl());
  }
  std::vector<FLAGS> flags = this->flags();
  if (flags.empty()) {
    oss << " 00: default";
  } else {
    std::vector<uint8_t> raw;
    std::vector<const char*> formated;
    for (FLAGS f : flags) {
      raw.push_back(get_raw(f));
      formated.push_back(pretty_name(f));
    }
    oss << fmt::format(" {:02x}: {}", fmt::join(raw, " "),
                       fmt::join(formated, ", "));
  }

  return oss.str();
}

std::string FunctionVariants::RuntimeTable::to_string() const {
  std::ostringstream oss;
  oss << fmt::format("Namespace: {} ({})", MachO::to_string(kind()),
                     (uint32_t)kind())
      << '\n';
  for (const RuntimeTableEntry& entry : entries()) {
    oss << "  " << entry << '\n';
  }
  return oss.str();
}

std::vector<FunctionVariants::RuntimeTable>
    FunctionVariants::parse_payload(SpanStream& stream) {
  std::vector<FunctionVariants::RuntimeTable> result;

  auto tableCount = stream.read<uint32_t>();
  if (!tableCount) {
    LIEF_DEBUG("Failed to read FunctionVariants.OnDiskFormat.tableCount");
    return result;
  }

  int64_t max_entries = stream.size() / sizeof(details::runtime_table_entry_t);

  for (size_t i = 0; i < *tableCount; ++i) {
    if (max_entries <= 0) {
      break;
    }

    auto tableOffset = stream.read<uint32_t>();
    if (!tableOffset) {
      LIEF_DEBUG("Failed to read FunctionVariants.OnDiskFormat.tableOffsets[{}]",
                 i);
      return result;
    }

    {
      ScopedStream runtime_table_strm(stream, *tableOffset);
      if (auto entry = parse_entry(*runtime_table_strm, max_entries)) {
        max_entries -= entry->entries().size();
        result.push_back(std::move(*entry));
      } else {
        LIEF_DEBUG("Failed to parse entry #{} at offset: {:#08x}", i,
                   *tableOffset);
      }
    }
  }
  return result;
}

result<FunctionVariants::RuntimeTable>
    FunctionVariants::parse_entry(BinaryStream& stream, uint64_t max_entries) {
  auto kind = stream.read<uint32_t>();

  if (!kind) {
    LIEF_DEBUG("Failed to read FunctionVariantsRuntimeTable.kind");
    return make_error_code(kind.error());
  }

  RuntimeTable runtime_table((RuntimeTable::KIND)*kind, stream.pos());

  auto count = stream.read<uint32_t>();

  if (!count) {
    LIEF_DEBUG("Failed to read FunctionVariantsRuntimeTable.count");
    return make_error_code(count.error());
  }

  for (size_t i = 0; i < *count; ++i) {
    if (i > max_entries) {
      LIEF_WARN("Too many FunctionVariants. Stopping");
      return runtime_table;
    }
    auto raw_entry = stream.read<details::runtime_table_entry_t>();
    if (!raw_entry) {
      LIEF_DEBUG("Failed to read FunctionVariantsRuntimeTable.entries[{}]", i);
      return runtime_table;
    }

    RuntimeTableEntry entry(*raw_entry);

    span<const uint8_t> flags = get_relevant_flags(entry.flag_bit_nums());
    std::vector<FLAGS> hl_flags =
        get_high_level_flags(flags, runtime_table.kind());
    entry.set_flags(std::move(hl_flags));

    runtime_table.add(std::move(entry));
  }
  return runtime_table;
}

const char* to_string(FunctionVariants::RuntimeTable::KIND kind) {
  using KIND = FunctionVariants::RuntimeTable::KIND;
  switch (kind) {
    default:
    case KIND::UNKNOWN: return "UNKNOWN";
    case KIND::PER_PROCESS: return "PER_PROCESS";
    case KIND::SYSTEM_WIDE: return "SYSTEM_WIDE";
    case KIND::ARM64: return "ARM64";
    case KIND::X86_64: return "X86_64";
  }
}

const char* to_string(FLAGS f) {
  switch (f) {
#define FUNCTION_VARIANT_FLAG(X, __, _)                                           \
  case FLAGS::X: return #X;
#include "LIEF/MachO/FunctionVariants/Arm64.def"
#include "LIEF/MachO/FunctionVariants/PerProcess.def"
#include "LIEF/MachO/FunctionVariants/SystemWide.def"
#include "LIEF/MachO/FunctionVariants/X86_64.def"
#undef FUNCTION_VARIANT_FLAG
    case FLAGS::UNKNOWN: return "UNKNOWN";
  }
  return "UNKNOWN";
}

}