/*
 * 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 "dnsdist-opentelemetry.hh"
#include "dnsdist-ecs.hh"

#include <memory>
#include <vector>

#ifndef DISABLE_PROTOBUF
#include "protozero-trace.hh"
#endif

namespace pdns::trace::dnsdist
{

TracesData Tracer::getTracesData()
{
#ifdef DISABLE_PROTOBUF
  return 0;
#else
  auto traceid = getTraceID();
  {
    auto data = d_data.read_only_lock();
    auto otTrace = pdns::trace::TracesData{
      .resource_spans = {
        {.resource = {
           .attributes = {
             {"service.name", {"dnsdist"}},
           }},
         .scope_spans = {{.scope = {
                            .name = "dnsdist/queryFromFrontend",
                            .version = PACKAGE_VERSION,
                            .attributes = {data->d_attributes.cbegin(), data->d_attributes.cend()},
                          },
                          .spans = {}}}}}};

    for (auto const& span : data->d_spans) {
      otTrace.resource_spans.at(0).scope_spans.at(0).spans.push_back(
        {
          .trace_id = traceid,
          .span_id = span.span_id == data->d_oldAndNewRootSpanID.oldID ? data->d_oldAndNewRootSpanID.newID : span.span_id,
          .parent_span_id = span.parent_span_id == data->d_oldAndNewRootSpanID.oldID ? data->d_oldAndNewRootSpanID.newID : span.parent_span_id,
          .name = span.name,
          .kind = pdns::trace::Span::SpanKind::SPAN_KIND_SERVER,
          .start_time_unix_nano = span.start_time_unix_nano,
          .end_time_unix_nano = span.end_time_unix_nano,
          .attributes = span.attributes,
        });
    }
    return otTrace;
  }
#endif
}

std::string Tracer::getOTProtobuf()
{
#ifdef DISABLE_PROTOBUF
  return 0;
#else
  // TODO: Should we close all spans?
  return getTracesData().encode();
#endif
}

SpanID Tracer::addSpan([[maybe_unused]] const std::string& name)
{
#ifdef DISABLE_PROTOBUF
  return 0;
#else
  return addSpan(name, getLastSpanID());
#endif
}

SpanID Tracer::addSpan([[maybe_unused]] const std::string& name, [[maybe_unused]] const SpanID& parentSpanID)
{
#ifdef DISABLE_PROTOBUF
  return 0;
#else
  auto spanID = pdns::trace::SpanID::getRandomSpanID();
  {
    auto data = d_data.lock();
    data->d_spans.push_back({
      .name = name,
      .span_id = spanID,
      .parent_span_id = parentSpanID,
      .start_time_unix_nano = pdns::trace::timestamp(),
      .end_time_unix_nano = 0,
      .attributes = {},
    });

    data->d_spanIDStack.emplace_back(spanID);
  }
  return spanID;
#endif
}

void Tracer::setTraceID([[maybe_unused]] const TraceID& traceID)
{
#ifndef DISABLE_PROTOBUF
  d_data.lock()->d_traceid = traceID;
#endif
}

void Tracer::setRootSpanID([[maybe_unused]] const SpanID& spanID)
{
#ifndef DISABLE_PROTOBUF
  SpanID oldRootSpanID{pdns::trace::s_emptySpanID};
  {
    auto data = d_data.read_only_lock();
    if (const auto& spans = data->d_spans; !spans.empty()) {
      auto iter = std::find_if(spans.cbegin(), spans.cend(), [](const auto& span) { return span.parent_span_id == pdns::trace::s_emptySpanID; });
      if (iter != spans.cend()) {
        oldRootSpanID = iter->span_id;
      }
    }
  }

  if (oldRootSpanID == pdns::trace::s_emptySpanID) {
    return;
  }

  {
    auto data = d_data.lock();
    data->d_oldAndNewRootSpanID.oldID = oldRootSpanID;
    data->d_oldAndNewRootSpanID.newID = spanID;
  }
#endif
}

// TODO: Figure out what to do with duplicate keys
bool Tracer::setTraceAttribute([[maybe_unused]] const std::string& key, [[maybe_unused]] const AnyValue& value)
{
#ifdef DISABLE_PROTOBUF
  // always successful
  return true;
#else
  d_data.lock()->d_attributes.push_back({key, value});
  return true;
#endif
}

void Tracer::closeSpan([[maybe_unused]] const SpanID& spanID)
{
#ifndef DISABLE_PROTOBUF
  auto data = d_data.lock();
  auto& spans = data->d_spans;
  auto spanIt = std::find_if(
    spans.rbegin(),
    spans.rend(),
    [&spanID](const miniSpan& span) { return span.span_id == spanID; });
  if (spanIt != spans.rend() && spanIt->end_time_unix_nano == 0) {
    spanIt->end_time_unix_nano = pdns::trace::timestamp();

    // Only closers are allowed, so this can never happen
    assert(!data->d_spanIDStack.empty());
    assert(data->d_spanIDStack.back() == spanID);
    data->d_spanIDStack.pop_back();
  }
#endif
}

void Tracer::setRootSpanAttribute([[maybe_unused]] const std::string& key, [[maybe_unused]] const AnyValue& value)
{
#ifndef DISABLE_PROTOBUF
  setSpanAttribute(getRootSpanID(), key, value);
#endif
}

// TODO: Figure out what to do with duplicate keys
void Tracer::setSpanAttribute([[maybe_unused]] const SpanID& spanid, [[maybe_unused]] const std::string& key, [[maybe_unused]] const AnyValue& value)
{
#ifndef DISABLE_PROTOBUF
  auto data = d_data.lock();
  auto& spans = data->d_spans;
  if (auto iter = std::find_if(spans.rbegin(),
                               spans.rend(),
                               [&spanid](const auto& span) { return span.span_id == spanid; });
      iter != spans.rend()) {
    iter->attributes.push_back({key, value});
  }
#endif
}

SpanID Tracer::getRootSpanID()
{
#ifdef DISABLE_PROTOBUF
  return 0;
#else
  auto data = d_data.read_only_lock();
  if (data->d_oldAndNewRootSpanID.newID != pdns::trace::s_emptySpanID) {
    return data->d_oldAndNewRootSpanID.newID;
  }

  if (auto& spans = data->d_spans; !spans.empty()) {
    auto iter = std::find_if(spans.cbegin(), spans.cend(), [](const auto& span) { return span.parent_span_id == pdns::trace::s_emptySpanID; });
    if (iter != spans.cend()) {
      return iter->span_id;
    }
  }
  return pdns::trace::s_emptySpanID;
#endif
}

SpanID Tracer::getLastSpanID()
{
#ifdef DISABLE_PROTOBUF
  return 0;
#else
  auto data = d_data.read_only_lock();
  if (data->d_spanIDStack.empty()) {
    return pdns::trace::s_emptySpanID;
  }
  if (data->d_spanIDStack.size() == 1 && data->d_spanIDStack.front() == data->d_oldAndNewRootSpanID.oldID) {
    return data->d_oldAndNewRootSpanID.newID;
  }
  return data->d_spanIDStack.back();
#endif
}

SpanID Tracer::getLastSpanIDForName([[maybe_unused]] const std::string& name)
{
#ifdef DISABLE_PROTOBUF
  return 0;
#else
  auto data = d_data.read_only_lock();
  if (auto& spans = data->d_spans; !spans.empty()) {
    if (auto iter = std::find_if(spans.rbegin(),
                                 spans.rend(),
                                 [name](const miniSpan& span) { return span.name == name; });
        iter != spans.rend()) {
      return iter->span_id;
    }
  }
  return pdns::trace::s_emptySpanID;
#endif
}

TraceID Tracer::getTraceID()
{
#ifdef DISABLE_PROTOBUF
  return 0;
#else
  auto data = d_data.lock();
  auto& traceID = data->d_traceid;
  if (traceID == pdns::trace::s_emptyTraceID) {
    traceID.makeRandom();
  }
  return traceID;
#endif
}

Tracer::Closer Tracer::getCloser([[maybe_unused]] const SpanID& spanid)
{
#ifdef DISABLE_PROTOBUF
  return Tracer::Closer();
#else
  return {shared_from_this(), spanid};
#endif
}

Tracer::Closer Tracer::openSpan([[maybe_unused]] const std::string& name)
{
#ifdef DISABLE_PROTOBUF
  return Tracer::Closer();
#else
  auto spanid = addSpan(name);
  return getCloser(spanid);
#endif
}

Tracer::Closer Tracer::openSpan([[maybe_unused]] const std::string& name, [[maybe_unused]] const SpanID& parentSpanID)
{
#ifdef DISABLE_PROTOBUF
  return Tracer::Closer();
#else
  auto spanid = addSpan(name, parentSpanID);
  return getCloser(spanid);
#endif
}

SpanID Tracer::Closer::getSpanID() const
{
#ifdef DISABLE_PROTOBUF
  return 0;
#else
  return d_spanID;
#endif
}

void Tracer::Closer::setAttribute([[maybe_unused]] const std::string& key, [[maybe_unused]] const AnyValue& value)
{
#ifdef DISABLE_PROTOBUF
  return;
#else
  return d_tracer->setSpanAttribute(d_spanID, key, value);
#endif
}

std::vector<uint8_t> makeEDNSTraceParentOption(std::shared_ptr<Tracer> tracer)
{
  std::vector<uint8_t> ret;
#ifndef DISABLE_PROTOBUF
  if (tracer == nullptr) {
    return ret;
  }
  ret.reserve(27);
  ret.push_back(0); // Version
  ret.push_back(0); // Reserved
  auto traceId = tracer->getTraceID();
  ret.insert(ret.end(), traceId.begin(), traceId.end());
  auto spanId = tracer->getLastSpanID();
  ret.insert(ret.end(), spanId.begin(), spanId.end());
  ret.push_back(0); // Flags
#endif
  return ret;
}

bool addTraceparentEdnsOptionToPacketBuffer(PacketBuffer& origBuf, const std::shared_ptr<Tracer>& tracer, const size_t qnameWireLength, const size_t proxyProtocolPayloadSize, const uint16_t traceparentOptionCode, const bool isTCP)
{
#ifndef DISABLE_PROTOBUF
  if (tracer == nullptr) {
    return false;
  }
  // buf contains the whole DNS query without PROXY protocol and TCP length header
  PacketBuffer buf{origBuf.begin() + proxyProtocolPayloadSize + (isTCP ? 2 : 0), origBuf.end()};

  uint16_t optRDPosition;
  size_t remaining;
  bool queryHadEdns = ::dnsdist::getEDNSOptionsStart(buf, qnameWireLength, &optRDPosition, &remaining) == 0;
  if (queryHadEdns) {
    size_t optLen = buf.size() - optRDPosition - remaining;
    removeEDNSOptionFromOPT(reinterpret_cast<char*>(buf.data() + optRDPosition), &optLen, traceparentOptionCode);
  }

  auto opt = pdns::trace::dnsdist::makeEDNSTraceParentOption(tracer);
  bool ednsAdded{false};
  bool optionAdded{false};
  uint16_t maxEdnsSize = queryHadEdns ? (uint16_t)(buf.at(optRDPosition - 6) << 8) + buf.at(optRDPosition - 5) : 512;
  setEDNSOption(buf, traceparentOptionCode, std::string(opt.begin(), opt.end()), isTCP ? std::numeric_limits<uint16_t>().max() : maxEdnsSize, ednsAdded, optionAdded);

  if (isTCP) {
    const std::array<uint8_t, 2> sizeBytes{static_cast<uint8_t>(buf.size() / 256), static_cast<uint8_t>(buf.size() % 256)};
    buf.insert(buf.begin(), sizeBytes.begin(), sizeBytes.end());
  }

  // Resize the buffer to remove the existing packet, but keep any PROXYv2 data
  origBuf.resize(proxyProtocolPayloadSize);
  // Insert the new query into the buffer
  origBuf.insert(origBuf.end(), buf.begin(), buf.end());

  return ednsAdded;
#else
  return false;
#endif
}

} // namespace pdns::trace::dnsdist
