/*

 * This file is part of PowerDNS or dnsdist.

 * Copyright -- PowerDNS.COM B.V. and its contributors

 *

 * This program is free software; you can redistribute it and/or modify

 * it under the terms of version 2 of the GNU General Public License as

 * published by the Free Software Foundation.

 *

 * In addition, for the avoidance of any doubt, permission is granted to

 * link this program with OpenSSL and to (re)distribute the binaries

 * produced as the result of such linking.

 *

 * This program is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 * GNU General Public License for more details.

 *

 * You should have received a copy of the GNU General Public License

 * along with this program; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

 */

#include "protozero-trace.hh"

#include "base64.hh"

#include "ednsoptions.hh"

#include "misc.hh"

#include <string>

#include <variant>

namespace pdns::trace

{

void AnyValue::encode(protozero::pbf_writer& writer) const

{

  if (std::holds_alternative<std::string>(*this)) {

    pdns::trace::encode(writer, 1, std::get<std::string>(*this), true);

  }

  else if (std::holds_alternative<bool>(*this)) {

    pdns::trace::encode(writer, 2, std::get<bool>(*this), true);

  }

  else if (std::holds_alternative<int64_t>(*this)) {

    pdns::trace::encode(writer, 3, std::get<int64_t>(*this), true);

  }

  else if (std::holds_alternative<double>(*this)) {

    pdns::trace::encode(writer, 4, std::get<double>(*this), true);

  }

  else if (std::holds_alternative<ArrayValue>(*this)) {

    protozero::pbf_writer sub{writer, 5};

    std::get<ArrayValue>(*this).encode(sub);

  }

  else if (std::holds_alternative<KeyValueList>(*this)) {

    protozero::pbf_writer sub{writer, 6};

    std::get<KeyValueList>(*this).encode(sub);

  }

  else if (std::holds_alternative<std::vector<uint8_t>>(*this)) {

    pdns::trace::encode(writer, 7, std::get<std::vector<uint8_t>>(*this), true);

  }

}

AnyValue AnyValue::decode(protozero::pbf_reader& reader)

{

  while (reader.next()) {

    switch (reader.tag()) {

    case 1:

      return AnyValue{reader.get_string()};

      break;

    case 2:

      return AnyValue{reader.get_bool()};

      break;

    case 3:

      return AnyValue{reader.get_int64()};

      break;

    case 4:

      return AnyValue{reader.get_double()};

      break;

    case 5: {

      protozero::pbf_reader arrayvalue = reader.get_message();

      return AnyValue{ArrayValue::decode(arrayvalue)};

      break;

    }

    case 6: {

      protozero::pbf_reader kvlist = reader.get_message();

      return AnyValue{KeyValueList::decode(kvlist)};

      break;

    }

    case 7: {

      auto value = reader.get_view();

      std::vector<uint8_t> data{};

      data.reserve(value.size());

      for (size_t i = 0; i < value.size(); ++i) {

        data.push_back(static_cast<uint8_t>(value.data()[i]));

      }

      return AnyValue{std::move(data)};

      break;

    }

    default:

      break;

    }

  }

  return {};

}

std::string AnyValue::toLogString() const

{

  if (std::holds_alternative<std::string>(*this)) {

    return "\"" + std::get<std::string>(*this) + "\"";

  }

  if (std::holds_alternative<bool>(*this)) {

    if (std::get<bool>(*this)) {

      return "true";

    }

    return "false";

  }

  if (std::holds_alternative<int64_t>(*this)) {

    return std::to_string(std::get<int64_t>(*this));

  }

  if (std::holds_alternative<double>(*this)) {

    return std::to_string(std::get<double>(*this));

  }

  if (std::holds_alternative<ArrayValue>(*this)) {

    std::string tmp = "[";

    for (auto const& val : std::get<ArrayValue>(*this).values) {

      tmp += val.toLogString() + ", ";

    }

    tmp.resize(tmp.size() - 2); // Strip ", "

    return tmp += "]";

  }

  if (std::holds_alternative<KeyValueList>(*this)) {

    std::string tmp = "{";

    for (auto const& val : std::get<KeyValueList>(*this).values) {

      tmp += val.key + ": " + val.value.toLogString() + ", ";

    }

    if (tmp.size() > 2) {

      tmp.resize(tmp.size() - 2); // Strip ", "

    }

    return tmp += "}";

  }

  if (std::holds_alternative<std::vector<uint8_t>>(*this)) {

    auto val = std::get<std::vector<uint8_t>>(*this);

    return Base64Encode(std::string(val.begin(), val.end()));

  }

  return "UNSUPPORTED TYPE";

}

void EntityRef::encode(protozero::pbf_writer& writer) const

{

  pdns::trace::encode(writer, 1, schema_url);

  pdns::trace::encode(writer, 2, type);

  for (auto const& element : id_keys) {

    pdns::trace::encode(writer, 3, element);

  }

  for (auto const& element : description_keys) {

    pdns::trace::encode(writer, 4, element);

  }

}

EntityRef EntityRef::decode(protozero::pbf_reader& reader)

{

  EntityRef ret;

  while (reader.next()) {

    switch (reader.tag()) {

    case 1:

      ret.schema_url = reader.get_string();

      break;

    case 2:

      ret.type = reader.get_string();

      break;

    case 3:

      ret.id_keys.emplace_back(reader.get_string());

      break;

    case 4:

      ret.description_keys.emplace_back(reader.get_string());

      break;

    default:

      break;

    }

  }

  return ret;

}

bool EntityRef::operator==(const EntityRef& rhs) const

{

  return (schema_url == rhs.schema_url && type == rhs.type && id_keys == rhs.id_keys && description_keys == rhs.description_keys);

}

void KeyValue::encode(protozero::pbf_writer& writer) const

{

  pdns::trace::encode(writer, 1, key);

  {

    protozero::pbf_writer val_sub{writer, 2};

    value.encode(val_sub);

  }

}

void Resource::encode(protozero::pbf_writer& writer) const

{

  pdns::trace::encode(writer, 1, attributes);

  pdns::trace::encode(writer, 2, dropped_attributes_count);

  pdns::trace::encode(writer, 3, entity_refs);

}

Resource Resource::decode(protozero::pbf_reader& reader)

{

  Resource ret;

  while (reader.next()) {

    switch (reader.tag()) {

    case 1: {

      auto sub = reader.get_message();

      ret.attributes.emplace_back(KeyValue::decode(sub));

      break;

    }

    case 2:

      ret.dropped_attributes_count = reader.get_uint32();

      break;

    case 3: {

      auto sub = reader.get_message();

      ret.entity_refs.emplace_back(EntityRef::decode(sub));

      break;

    }

    default:

      break;

    }

  }

  return ret;

}

bool Resource::operator==(const Resource& rhs) const

{

  return (attributes == rhs.attributes && dropped_attributes_count == rhs.dropped_attributes_count && entity_refs == rhs.entity_refs);

}

void InstrumentationScope::encode(protozero::pbf_writer& writer) const

{

  pdns::trace::encode(writer, 1, name);

  pdns::trace::encode(writer, 2, version);

  pdns::trace::encode(writer, 3, attributes);

  pdns::trace::encode(writer, 4, dropped_attributes_count);

}

InstrumentationScope InstrumentationScope::decode(protozero::pbf_reader& reader)

{

  InstrumentationScope ret;

  while (reader.next()) {

    switch (reader.tag()) {

    case 1:

      ret.name = reader.get_string();

      break;

    case 2:

      ret.version = reader.get_string();

      break;

    case 3: {

      auto sub = reader.get_message();

      ret.attributes.emplace_back(KeyValue::decode(sub));

      break;

    }

    case 4:

      ret.dropped_attributes_count = reader.get_uint32();

      break;

    default:

      break;

    }

  }

  return ret;

}

bool InstrumentationScope::operator==(const InstrumentationScope& rhs) const

{

  return (name == rhs.name && version == rhs.version && attributes == rhs.attributes && dropped_attributes_count == rhs.dropped_attributes_count);

}

void Status::encode(protozero::pbf_writer& writer) const

{

  pdns::trace::encode(writer, 2, message);

  pdns::trace::encode(writer, 3, uint32_t(code));

}

Status Status::decode(protozero::pbf_reader& reader)

{

  Status ret;

  while (reader.next()) {

    switch (reader.tag()) {

    case 2:

      ret.message = reader.get_string();

      break;

    case 3:

      ret.code = static_cast<StatusCode>(reader.get_uint32());

      break;

    default:

      break;

    }

  }

  return ret;

}

bool Status::operator==(const Status& rhs) const

{

  return (message == rhs.message && code == rhs.code);

}

void Span::Event::encode(protozero::pbf_writer& writer) const

{

  pdns::trace::encodeFixed(writer, 1, time_unix_nano);

  pdns::trace::encode(writer, 2, name);

  pdns::trace::encode(writer, 3, attributes);

  pdns::trace::encode(writer, 4, dropped_attributes_count);

}

Span::Event Span::Event::decode(protozero::pbf_reader& reader)

{

  Span::Event ret;

  while (reader.next()) {

    switch (reader.tag()) {

    case 1:

      ret.time_unix_nano = reader.get_fixed64();

      break;

    case 2:

      ret.name = reader.get_string();

      break;

    case 3: {

      auto sub = reader.get_message();

      ret.attributes.emplace_back(KeyValue::decode(sub));

      break;

    }

    case 4:

      ret.dropped_attributes_count = reader.get_uint32();

    default:

      break;

    }

  }

  return ret;

}

bool Span::Event::operator==(const Span::Event& rhs) const

{

  return (time_unix_nano == rhs.time_unix_nano && name == rhs.name && attributes == rhs.attributes && dropped_attributes_count == rhs.dropped_attributes_count);

}

void Span::Link::encode(protozero::pbf_writer& writer) const

{

  pdns::trace::encode(writer, 1, trace_id);

  pdns::trace::encode(writer, 2, span_id);

  pdns::trace::encode(writer, 3, trace_state);

  pdns::trace::encode(writer, 4, attributes);

  pdns::trace::encode(writer, 5, dropped_attributes_count);

  pdns::trace::encodeFixed(writer, 6, flags);

}

Span::Link Span::Link::decode(protozero::pbf_reader& reader)

{

  Link ret;

  while (reader.next()) {

    switch (reader.tag()) {

    case 1:

      ret.trace_id = decodeTraceID(reader);

      break;

    case 2:

      ret.span_id = decodeSpanID(reader);

      break;

    case 3:

      ret.trace_state = reader.get_string();

      break;

    case 4: {

      auto sub = reader.get_message();

      ret.attributes.emplace_back(KeyValue::decode(sub));

      break;

    }

    case 5:

      ret.dropped_attributes_count = reader.get_uint32();

      break;

    case 6:

      ret.flags = reader.get_fixed32();

    default:

      break;

    }

  }

  return ret;

}

bool Span::Link::operator==(const Span::Link& rhs) const

{

  return (trace_id == rhs.trace_id && span_id == rhs.span_id && trace_state == rhs.trace_state && attributes == rhs.attributes && dropped_attributes_count == rhs.dropped_attributes_count && flags == rhs.flags);

}

void Span::encode(protozero::pbf_writer& writer) const

{

  pdns::trace::encode(writer, 1, trace_id);

  pdns::trace::encode(writer, 2, span_id);

  pdns::trace::encode(writer, 3, trace_state);

  pdns::trace::encode(writer, 4, parent_span_id);

  pdns::trace::encode(writer, 5, name);

  pdns::trace::encode(writer, 6, uint32_t(kind));

  pdns::trace::encodeFixed(writer, 7, start_time_unix_nano);

  pdns::trace::encodeFixed(writer, 8, end_time_unix_nano);

  pdns::trace::encode(writer, 9, attributes);

  pdns::trace::encode(writer, 10, dropped_attributes_count);

  pdns::trace::encode(writer, 11, events);

  pdns::trace::encode(writer, 12, dropped_events_count);

  pdns::trace::encode(writer, 13, links);

  pdns::trace::encode(writer, 14, dropped_links_count);

  if (status.code != Status::StatusCode::STATUS_CODE_UNSET || !status.message.empty()) {

    protozero::pbf_writer sub{writer, 15};

    status.encode(sub);

  }

  pdns::trace::encodeFixed(writer, 16, flags);

}

Span Span::decode(protozero::pbf_reader& reader)

{

  Span ret;

  while (reader.next()) {

    switch (reader.tag()) {

    case 1:

      ret.trace_id = decodeTraceID(reader);

      break;

    case 2:

      ret.span_id = decodeSpanID(reader);

      break;

    case 3:

      ret.trace_state = reader.get_string();

      break;

    case 4:

      ret.parent_span_id = decodeSpanID(reader);

      break;

    case 5:

      ret.name = reader.get_string();

      break;

    case 6:

      ret.kind = static_cast<Span::SpanKind>(reader.get_uint32());

      break;

    case 7:

      ret.start_time_unix_nano = reader.get_fixed64();

      break;

    case 8:

      ret.end_time_unix_nano = reader.get_fixed64();

      break;

    case 9: {

      auto sub = reader.get_message();

      ret.attributes.emplace_back(KeyValue::decode(sub));

      break;

    }

    case 10:

      ret.dropped_attributes_count = reader.get_uint32();

      break;

    case 11: {

      auto sub = reader.get_message();

      ret.events.emplace_back(Span::Event::decode(sub));

      break;

    }

    case 12:

      ret.dropped_events_count = reader.get_uint32();

      break;

    case 13: {

      auto sub = reader.get_message();

      ret.links.emplace_back(Span::Link::decode(sub));

      break;

    }

    case 14:

      ret.dropped_links_count = reader.get_uint32();

      break;

    case 15: {

      auto sub = reader.get_message();

      ret.status = Status::decode(sub);

      break;

    }

    case 16:

      ret.flags = reader.get_fixed32();

      break;

    default:

      break;

    }

  }

  return ret;

}

bool Span::operator==(const Span& rhs) const

{

  return (trace_id == rhs.trace_id && span_id == rhs.span_id && trace_state == rhs.trace_state && parent_span_id == rhs.parent_span_id && name == rhs.name && kind == rhs.kind && start_time_unix_nano == rhs.start_time_unix_nano && end_time_unix_nano == rhs.end_time_unix_nano && attributes == rhs.attributes && dropped_attributes_count == rhs.dropped_attributes_count && events == rhs.events && dropped_events_count == rhs.dropped_events_count && links == rhs.links && dropped_links_count == rhs.dropped_links_count && status == rhs.status && flags == rhs.flags);

}

void ScopeSpans::encode(protozero::pbf_writer& writer) const

{

  {

    protozero::pbf_writer sub{writer, 1};

    scope.encode(sub);

  }

  pdns::trace::encode(writer, 2, spans);

  pdns::trace::encode(writer, 3, schema_url);

}

ScopeSpans ScopeSpans::decode(protozero::pbf_reader& reader)

{

  ScopeSpans ret;

  while (reader.next()) {

    switch (reader.tag()) {

    case 1: {

      auto sub = reader.get_message();

      ret.scope = InstrumentationScope::decode(sub);

      break;

    }

    case 2: {

      auto sub = reader.get_message();

      ret.spans.emplace_back(Span::decode(sub));

      break;

    }

    case 3:

      ret.schema_url = reader.get_string();

    default:

      break;

    }

  }

  return ret;

}

bool ScopeSpans::operator==(const ScopeSpans& rhs) const

{

  return (scope == rhs.scope && spans == rhs.spans && schema_url == rhs.schema_url);

}

void ResourceSpans::encode(protozero::pbf_writer& writer) const

{

  {

    protozero::pbf_writer sub{writer, 1};

    resource.encode(sub);

  }

  pdns::trace::encode(writer, 2, scope_spans);

  pdns::trace::encode(writer, 3, schema_url);

}

ResourceSpans ResourceSpans::decode(protozero::pbf_reader& reader)

{

  ResourceSpans ret;

  while (reader.next()) {

    switch (reader.tag()) {

    case 1: {

      protozero::pbf_reader sub = reader.get_message();

      ret.resource = Resource::decode(sub);

      break;

    }

    case 2: {

      protozero::pbf_reader sub = reader.get_message();

      ret.scope_spans.emplace_back(ScopeSpans::decode(sub));

      break;

    }

    case 3:

      ret.schema_url = reader.get_string();

    default:

      break;

    }

  }

  return ret;

}

bool ResourceSpans::operator==(const ResourceSpans& rhs) const

{

  return (resource == rhs.resource && scope_spans == rhs.scope_spans && schema_url == rhs.schema_url);

}

void TracesData::encode(protozero::pbf_writer& writer) const

{

  pdns::trace::encode(writer, 1, resource_spans);

}

TracesData TracesData::decode(protozero::pbf_reader& reader)

{

  TracesData ret;

  while (reader.next()) {

    switch (reader.tag()) {

    case 1: {

      auto sub = reader.get_message();

      ret.resource_spans.emplace_back(ResourceSpans::decode(sub));

      break;

    }

    default:

      break;

    }

  }

  return ret;

}

KeyValue KeyValue::decode(protozero::pbf_reader& reader)

{

  KeyValue value;

  while (reader.next()) {

    switch (reader.tag()) {

    case 1:

      value.key = reader.get_string();

      break;

    case 2: {

      protozero::pbf_reader sub = reader.get_message();

      value.value = AnyValue::decode(sub);

      break;

    }

    default:

      break;

    }

  }

  return value;

}

bool extractOTraceIDs(const EDNSOptionViewMap& map, EDNSOptionCode::EDNSOptionCodeEnum eoc, pdns::trace::InitialSpanInfo& span)

{

  // traceid gets set from edns options (if available and well-formed)

  // parent_span_id gets set from edns options (if available and well-formed)

  auto traceidset = extractOTraceIDs(map, eoc, span.trace_id, span.parent_span_id);

  return traceidset;

}

/**

 * @brief Extract the Trace and Span IDs from

 *

 * @param map an EDNSOptionViewMap where a TraceID option might be

 * @param eoc the EDNS option code where the TraceID might be, OTTRACEIDS by default

 * @param traceID will be set to the TraceID in the EDNS option, untouched otherwise

 * @param spanID will be set to the SpanID in the EDNS options, untouched otherwise

 * @return true if a traceid was found in the EDNS options

 */

bool extractOTraceIDs(const EDNSOptionViewMap& map, EDNSOptionCode::EDNSOptionCodeEnum eoc, pdns::trace::TraceID& traceID, pdns::trace::SpanID& spanID)

{

  EDNSOptionCode::EDNSOptionCodeEnum realEOC = eoc;

  if (realEOC == 0) {

    realEOC = EDNSOptionCode::TRACEPARENT;

  }

  bool traceidset = false;

  if (const auto& option = map.find(realEOC); option != map.end()) {

    const auto& value = option->second.values.at(0);

    const EDNSOTTraceRecordView data{reinterpret_cast<const uint8_t*>(value.content), value.size}; // NOLINT(cppcoreguidelines-pro-type-reinterpret-cast)

    // 1 byte version, then tracid then optional spanid

    uint8_t version{};

    if (data.getVersion(version) && version == 0) {

      if (data.getTraceID(traceID)) {

        traceidset = true;

      }

      (void)data.getSpanID(spanID);

    }

  }

  return traceidset;

}

std::string SpanID::toLogString() const

{

  return makeHexDump(std::string(this->begin(), this->end()), "");

}

std::string TraceID::toLogString() const

{

  return makeHexDump(std::string(this->begin(), this->end()), "");

}

} // namespace pdns::trace