#include "message.hpp"
#include <oxenc/endian.h>

#include "dns.hpp"
#include "srv_data.hpp"
#include <llarp/util/buffer.hpp>
#include <llarp/util/logging.hpp>
#include <llarp/net/ip.hpp>

#include <array>

namespace llarp
{
  namespace dns
  {
    static auto logcat = log::Cat("dns");

    bool
    MessageHeader::Encode(llarp_buffer_t* buf) const
    {
      if (!buf->put_uint16(id))
        return false;
      if (!buf->put_uint16(fields))
        return false;
      if (!buf->put_uint16(qd_count))
        return false;
      if (!buf->put_uint16(an_count))
        return false;
      if (!buf->put_uint16(ns_count))
        return false;
      return buf->put_uint16(ar_count);
    }

    bool
    MessageHeader::Decode(llarp_buffer_t* buf)
    {
      if (!buf->read_uint16(id))
        return false;
      if (!buf->read_uint16(fields))
        return false;
      if (!buf->read_uint16(qd_count))
        return false;
      if (!buf->read_uint16(an_count))
        return false;
      if (!buf->read_uint16(ns_count))
        return false;
      if (!buf->read_uint16(ar_count))
        return false;
      return true;
    }

    util::StatusObject
    MessageHeader::ToJSON() const
    {
      return util::StatusObject{};
    }

    Message::Message(Message&& other)
        : hdr_id(std::move(other.hdr_id))
        , hdr_fields(std::move(other.hdr_fields))
        , questions(std::move(other.questions))
        , answers(std::move(other.answers))
        , authorities(std::move(other.authorities))
        , additional(std::move(other.additional))
    {}

    Message::Message(const Message& other)
        : hdr_id(other.hdr_id)
        , hdr_fields(other.hdr_fields)
        , questions(other.questions)
        , answers(other.answers)
        , authorities(other.authorities)
        , additional(other.additional)
    {}

    Message::Message(const MessageHeader& hdr) : hdr_id(hdr.id), hdr_fields(hdr.fields)
    {
      questions.resize(size_t(hdr.qd_count));
      answers.resize(size_t(hdr.an_count));
      authorities.resize(size_t(hdr.ns_count));
      additional.resize(size_t(hdr.ar_count));
    }

    Message::Message(const Question& question) : hdr_id{0}, hdr_fields{}
    {
      questions.emplace_back(question);
    }

    bool
    Message::Encode(llarp_buffer_t* buf) const
    {
      MessageHeader hdr;
      hdr.id = hdr_id;
      hdr.fields = hdr_fields;
      hdr.qd_count = questions.size();
      hdr.an_count = answers.size();
      hdr.ns_count = 0;
      hdr.ar_count = 0;

      if (!hdr.Encode(buf))
        return false;

      for (const auto& question : questions)
        if (!question.Encode(buf))
          return false;

      for (const auto& answer : answers)
        if (!answer.Encode(buf))
          return false;

      return true;
    }

    bool
    Message::Decode(llarp_buffer_t* buf)
    {
      for (auto& qd : questions)
      {
        if (!qd.Decode(buf))
        {
          log::error(logcat, "failed to decode question");
          return false;
        }
        log::debug(logcat, "question: {}", qd);
      }
      for (auto& an : answers)
      {
        if (not an.Decode(buf))
        {
          log::debug(logcat, "failed to decode answer");
          return false;
        }
      }
      return true;
    }

    util::StatusObject
    Message::ToJSON() const
    {
      std::vector<util::StatusObject> ques;
      std::vector<util::StatusObject> ans;
      for (const auto& q : questions)
      {
        ques.push_back(q.ToJSON());
      }
      for (const auto& a : answers)
      {
        ans.push_back(a.ToJSON());
      }
      return util::StatusObject{{"questions", ques}, {"answers", ans}};
    }

    OwnedBuffer
    Message::ToBuffer() const
    {
      std::array<byte_t, 1500> tmp;
      llarp_buffer_t buf{tmp};
      if (not Encode(&buf))
        throw std::runtime_error("cannot encode dns message");
      return OwnedBuffer::copy_used(buf);
    }

    void Message::AddServFail(RR_TTL_t)
    {
      if (questions.size())
      {
        hdr_fields |= flags_RCODEServFail;
        // authorative response with recursion available
        hdr_fields |= flags_QR | flags_AA | flags_RA;
        // don't allow recursion on this request
        hdr_fields &= ~flags_RD;
      }
    }

    static constexpr uint16_t
    reply_flags(uint16_t setbits)
    {
      return setbits | flags_QR | flags_AA | flags_RA;
    }

    void
    Message::AddINReply(llarp::huint128_t ip, bool isV6, RR_TTL_t ttl)
    {
      if (questions.size())
      {
        hdr_fields = reply_flags(hdr_fields);
        ResourceRecord rec;
        rec.rr_name = questions[0].qname;
        rec.rr_class = qClassIN;
        rec.ttl = ttl;
        if (isV6)
        {
          rec.rr_type = qTypeAAAA;
          ip.ToV6(rec.rData);
        }
        else
        {
          const auto addr = net::TruncateV6(ip);
          rec.rr_type = qTypeA;
          rec.rData.resize(4);
          oxenc::write_host_as_big(addr.h, rec.rData.data());
        }
        answers.emplace_back(std::move(rec));
      }
    }

    void
    Message::AddAReply(std::string name, RR_TTL_t ttl)
    {
      if (questions.size())
      {
        hdr_fields = reply_flags(hdr_fields);

        const auto& question = questions[0];
        answers.emplace_back();
        auto& rec = answers.back();
        rec.rr_name = question.qname;
        rec.rr_type = question.qtype;
        rec.rr_class = qClassIN;
        rec.ttl = ttl;
        std::array<byte_t, 512> tmp = {{0}};
        llarp_buffer_t buf(tmp);
        if (EncodeNameTo(&buf, name))
        {
          buf.sz = buf.cur - buf.base;
          rec.rData.resize(buf.sz);
          memcpy(rec.rData.data(), buf.base, buf.sz);
        }
      }
    }

    void
    Message::AddNSReply(std::string name, RR_TTL_t ttl)
    {
      if (not questions.empty())
      {
        hdr_fields = reply_flags(hdr_fields);

        const auto& question = questions[0];
        answers.emplace_back();
        auto& rec = answers.back();
        rec.rr_name = question.qname;
        rec.rr_type = qTypeNS;
        rec.rr_class = qClassIN;
        rec.ttl = ttl;
        std::array<byte_t, 512> tmp = {{0}};
        llarp_buffer_t buf(tmp);
        if (EncodeNameTo(&buf, name))
        {
          buf.sz = buf.cur - buf.base;
          rec.rData.resize(buf.sz);
          memcpy(rec.rData.data(), buf.base, buf.sz);
        }
      }
    }

    void
    Message::AddCNAMEReply(std::string name, RR_TTL_t ttl)
    {
      if (questions.size())
      {
        hdr_fields = reply_flags(hdr_fields);

        const auto& question = questions[0];
        answers.emplace_back();
        auto& rec = answers.back();
        rec.rr_name = question.qname;
        rec.rr_type = qTypeCNAME;
        rec.rr_class = qClassIN;
        rec.ttl = ttl;
        std::array<byte_t, 512> tmp = {{0}};
        llarp_buffer_t buf(tmp);
        if (EncodeNameTo(&buf, name))
        {
          buf.sz = buf.cur - buf.base;
          rec.rData.resize(buf.sz);
          memcpy(rec.rData.data(), buf.base, buf.sz);
        }
      }
    }

    void
    Message::AddMXReply(std::string name, uint16_t priority, RR_TTL_t ttl)
    {
      if (questions.size())
      {
        hdr_fields = reply_flags(hdr_fields);

        const auto& question = questions[0];
        answers.emplace_back();
        auto& rec = answers.back();
        rec.rr_name = question.qname;
        rec.rr_type = qTypeMX;
        rec.rr_class = qClassIN;
        rec.ttl = ttl;
        std::array<byte_t, 512> tmp = {{0}};
        llarp_buffer_t buf(tmp);
        buf.put_uint16(priority);
        if (EncodeNameTo(&buf, name))
        {
          buf.sz = buf.cur - buf.base;
          rec.rData.resize(buf.sz);
          memcpy(rec.rData.data(), buf.base, buf.sz);
        }
      }
    }

    void
    Message::AddSRVReply(std::vector<SRVData> records, RR_TTL_t ttl)
    {
      hdr_fields = reply_flags(hdr_fields);

      const auto& question = questions[0];

      for (const auto& srv : records)
      {
        if (not srv.IsValid())
        {
          AddNXReply();
          return;
        }

        answers.emplace_back();
        auto& rec = answers.back();
        rec.rr_name = question.qname;
        rec.rr_type = qTypeSRV;
        rec.rr_class = qClassIN;
        rec.ttl = ttl;

        std::array<byte_t, 512> tmp = {{0}};
        llarp_buffer_t buf(tmp);

        buf.put_uint16(srv.priority);
        buf.put_uint16(srv.weight);
        buf.put_uint16(srv.port);

        std::string target;
        if (srv.target == "")
        {
          // get location of second dot (after service.proto) in qname
          size_t pos = question.qname.find(".");
          pos = question.qname.find(".", pos + 1);

          target = question.qname.substr(pos + 1);
        }
        else
        {
          target = srv.target;
        }

        if (not EncodeNameTo(&buf, target))
        {
          AddNXReply();
          return;
        }

        buf.sz = buf.cur - buf.base;
        rec.rData.resize(buf.sz);
        memcpy(rec.rData.data(), buf.base, buf.sz);
      }
    }

    void
    Message::AddTXTReply(std::string str, RR_TTL_t ttl)
    {
      auto& rec = answers.emplace_back();
      rec.rr_name = questions[0].qname;
      rec.rr_class = qClassIN;
      rec.rr_type = qTypeTXT;
      rec.ttl = ttl;
      std::array<byte_t, 1024> tmp{};
      llarp_buffer_t buf(tmp);
      while (not str.empty())
      {
        const auto left = std::min(str.size(), size_t{256});
        const auto sub = str.substr(0, left);
        uint8_t byte = left;
        *buf.cur = byte;
        buf.cur++;
        if (not buf.write(sub.begin(), sub.end()))
          throw std::length_error("text record too big");
        str = str.substr(left);
      }
      buf.sz = buf.cur - buf.base;
      rec.rData.resize(buf.sz);
      std::copy_n(buf.base, buf.sz, rec.rData.data());
    }

    void Message::AddNXReply(RR_TTL_t)
    {
      if (questions.size())
      {
        answers.clear();
        authorities.clear();
        additional.clear();

        // authorative response with recursion available
        hdr_fields = reply_flags(hdr_fields);
        // don't allow recursion on this request
        hdr_fields &= ~flags_RD;
        hdr_fields |= flags_RCODENameError;
      }
    }

    std::string
    Message::ToString() const
    {
      return fmt::format(
          "[DNSMessage id={:x} fields={:x} questions={{{}}} answers={{{}}} authorities={{{}}} "
          "additional={{{}}}]",
          hdr_id,
          hdr_fields,
          fmt::format("{}", fmt::join(questions, ",")),
          fmt::format("{}", fmt::join(answers, ",")),
          fmt::format("{}", fmt::join(authorities, ",")),
          fmt::format("{}", fmt::join(additional, ",")));
    }

    std::optional<Message>
    MaybeParseDNSMessage(llarp_buffer_t buf)
    {
      MessageHeader hdr{};
      if (not hdr.Decode(&buf))
        return std::nullopt;

      Message msg{hdr};
      if (not msg.Decode(&buf))
        return std::nullopt;
      return msg;
    }
  }  // namespace dns
}  // namespace llarp