/* Goal: a fully standards compliant basic authoritative server. In <500 lines.
   Non-goals: notifications, slaving zones, name compression, edns,
              performance
*/
#include <cstdint>
#include <string>
#include <vector>
#include <deque>
#include <map>
#include <stdexcept>
#include "sclasses.hh"
#include "dns.hh"
#include "safearray.hh"
#include <thread>

using namespace std;

typedef std::string dnslabel;

enum class RCode 
{
  Noerror = 0, Servfail = 2, Nxdomain = 3, Notimp = 4, Refused = 5
};

enum class DNSType : uint16_t
{
  A = 1, NS = 2, CNAME = 5, SOA=6, AAAA = 28, IXFR = 251, AXFR = 252, ANY = 255
};

enum class DNSSection
{
  Question, Answer, Authority, Additional
};

typedef deque<dnslabel> dnsname;
// this should perform escaping rules!
static std::ostream & operator<<(std::ostream &os, const dnsname& d)
{
  for(const auto& l : d) {
    os<<l<<".";
  }
  return os;
}

struct RRSet
{
  vector<string> contents;
  uint32_t ttl{3600};
};

struct CIStringCompare: public std::binary_function<string, string, bool>
{
  bool operator()(const string& a, const string& b) const
  {
    return strcasecmp(a.c_str(), b.c_str()) < 0; // XXX locale pain, plus embedded zeros
  }
};


struct DNSNode
{
  const DNSNode* find(dnsname& name, dnsname& last, bool* passedZonecut=0) const;
  DNSNode* add(dnsname name);
  map<dnslabel, DNSNode, CIStringCompare> children;
  map<DNSType, RRSet > rrsets;
  
  DNSNode* zone{0}; // if this is set, this node is a zone
};

const DNSNode* DNSNode::find(dnsname& name, dnsname& last, bool* passedZonecut) const
{
  cout<<"find for '"<<name<<"', last is now '"<<last<<"'"<<endl;
  if(!last.empty() && passedZonecut && rrsets.count(DNSType::NS)) {
    *passedZonecut=true;
  }

  if(name.empty()) {
    cout<<"Empty lookup, returning this node or 0"<<endl;
    if(!zone && rrsets.empty()) // only root zone can have this
      return 0;
    else
      return this;
  }
  cout<<"Children at this node: ";
  for(const auto& c: children) cout <<"'"<<c.first<<"' ";
  cout<<endl;
  auto iter = children.find(name.back());
  cout<<"Looked for child called '"<<name.back()<<"'"<<endl;
  if(iter == children.end()) {
    cout<<"Found nothing, trying wildcard"<<endl;
    iter = children.find("*");
    if(iter == children.end()) {
      cout<<"Still nothing, returning leaf"<<endl;
      return this;
    }
    else {
      cout<<"Had wildcard match, following"<<endl;
    }
  }
  cout<<"Had match, continuing to child '"<<iter->first<<"'"<<endl;
  last.push_front(name.back());
  name.pop_back();
  return iter->second.find(name, last, passedZonecut);
}

DNSNode* DNSNode::add(dnsname name) 
{
  cout<<"Add for '"<<name<<"'"<<endl;
  if(name.size() == 1) {
    cout<<"Last label, adding "<<name.front()<<endl;
    return &children[name.front()];
  }

  auto back = name.back();
  name.pop_back();
  auto iter = children.find(back);

  if(iter == children.end()) {
    cout<<"Inserting new child for "<<back<<endl;
    return children[back].add(name);
  }
  return iter->second.add(name);
}

struct DNSMessage
{
  struct dnsheader dh=dnsheader{};
  SafeArray<500> payload;

  dnsname getName();
  void getQuestion(dnsname& name, DNSType& type);
  void setQuestion(const dnsname& name, DNSType type);
  void putRR(DNSSection section, const dnsname& name, DNSType type, uint32_t ttl, const std::string& rr);

  string serialize() const;
}; // __attribute__((packed));

dnsname DNSMessage::getName()
{
  dnsname name;
  for(;;) {
    uint8_t labellen=payload.getUInt8();
    if(labellen > 63)
      throw std::runtime_error("Got a compressed label");
    if(!labellen) // end of dnsname
      break;
    dnslabel label = payload.getBlob(labellen);
    name.push_back(label);
  }
  return name;
}

void DNSMessage::getQuestion(dnsname& name, DNSType& type)
{
  name=getName();
  type=(DNSType)payload.getUInt16();
}

void putName(auto& payload, const dnsname& name)
{
  for(const auto& l : name) {
    if(l.size() > 63)
      throw std::runtime_error("Can't emit a label larger than 63 characters");
    payload.putUInt8(l.size());
    payload.putBlob(l);
  }
  payload.putUInt8(0);
}

void DNSMessage::putRR(DNSSection section, const dnsname& name, DNSType type, uint32_t ttl, const std::string& content)
{
  putName(payload, name);
  payload.putUInt16((int)type); payload.putUInt16(1);
  payload.putUInt32(ttl);
  payload.putUInt16(content.size()); // check for overflow!
  payload.putBlob(content);

  switch(section) {
    case DNSSection::Question:
      throw runtime_error("Can't add questions to a DNS Message with putRR");
    case DNSSection::Answer:
      dh.ancount = htons(ntohs(dh.ancount) + 1);
      break;
    case DNSSection::Authority:
      dh.nscount = htons(ntohs(dh.nscount) + 1);
      break;
    case DNSSection::Additional:
      dh.arcount = htons(ntohs(dh.arcount) + 1);
      break;
  }
}

void DNSMessage::setQuestion(const dnsname& name, DNSType type)
{
  putName(payload, name);
  payload.putUInt16((uint16_t)type);
  payload.putUInt16(1); // class
}

string DNSMessage::serialize() const
{
  return string((const char*)this, (const char*)this + sizeof(dnsheader) + payload.payloadpos);
}


static_assert(sizeof(DNSMessage) == 516, "dnsmessage size must be 516");

std::string serializeDNSName(const dnsname& dn)
{
  std::string ret;
  for(const auto & l : dn) {
    ret.append(1, l.size());
    ret+=l;
  }
  ret.append(1, (char)0);
  return ret;
}

std::string serializeSOARecord(const dnsname& mname, const dnsname& rname, uint32_t serial, uint32_t minimum=3600, uint32_t refresh=10800, uint32_t retry=3600, uint32_t expire=604800)
{
  SafeArray<256> sa;
  putName(sa, mname);
  putName(sa, rname);
  sa.putUInt32(serial);
  sa.putUInt32(refresh);
  sa.putUInt32(retry);
  sa.putUInt32(expire);
  sa.putUInt32(minimum);
  
  return sa.serialize();
}

std::string serializeARecord(const std::string& src)
{
  ComboAddress ca(src);
  if(ca.sin4.sin_family != AF_INET)
    throw std::runtime_error("Could not convert '"+src+"' to an IPv4 address");
  auto p = (const char*)&ca.sin4.sin_addr.s_addr;
  return std::string(p, p+4);
}

std::string serializeAAAARecord(const std::string& src)
{
  ComboAddress ca(src);
  if(ca.sin4.sin_family != AF_INET6)
    throw std::runtime_error("Could not convert '"+src+"' to an IPv6 address");
  auto p = (const char*)ca.sin6.sin6_addr.s6_addr;
  return std::string(p, p+16);
}


dnsname operator+(const dnsname& a, const dnsname& b)
{
  dnsname ret=a;
  for(const auto& l : b)
    ret.push_back(l);
  return ret;
}

bool processQuestion(const DNSNode& zones, DNSMessage& dm, const ComboAddress& local, const ComboAddress& remote, DNSMessage& response)
try
{
  dnsname name;
  DNSType type;
  dm.getQuestion(name, type);
  cout<<"Received a query from "<<remote.toStringWithPort()<<" for "<<name<<" and type "<<(int)type<<endl;
  
  response.dh = dm.dh;
  response.dh.ad = 0;
  response.dh.ra = 0;
  response.dh.aa = 0;
  response.dh.qr = 1;
  response.dh.ancount = response.dh.arcount = response.dh.nscount = 0;
  response.setQuestion(name, type);
  
  if(type == DNSType::AXFR) {
      cout<<"Query was for AXFR or IXFR over UDP, can't do that"<<endl;
      response.dh.rcode = (int)RCode::Servfail;
      return true;
  }

  if(dm.dh.opcode != 0) {
    cout<<"Query had non-zero opcode "<<dm.dh.opcode<<", sending NOTIMP"<<endl;
    response.dh.rcode = (int)RCode::Notimp;
    return true;
  }
    
  dnsname zone;
  auto fnd = zones.find(name, zone);
  if(fnd && fnd->zone) {
    cout<<"---\nBest zone: "<<zone<<", name now "<<name<<", loaded: "<<(void*)fnd->zone<<endl;

    response.dh.aa = 1; 
    
    auto bestzone = fnd->zone;
    dnsname searchname(name), lastnode;
    bool passedZonecut=false;
    auto node = bestzone->find(searchname, lastnode, &passedZonecut);
    if(passedZonecut)
      response.dh.aa = false;
    
    if(!node) {
      cout<<"Found nothing in zone '"<<zone<<"' for lhs '"<<name<<"'"<<endl;
    }
    else if(!searchname.empty()) {
      cout<<"This was a partial match, searchname now "<<searchname<<endl;
      
      for(const auto& rr: node->rrsets) {
        cout<<"  Have type "<<(int)rr.first<<endl;
      }
      auto iter = node->rrsets.find(DNSType::NS);
      if(iter != node->rrsets.end() && passedZonecut) {
        cout<<"Have delegation"<<endl;
        const auto& rrset = iter->second;
        for(const auto& rr : rrset.contents) {
          response.putRR(DNSSection::Answer, lastnode+zone, DNSType::NS, rrset.ttl, rr);
        }
        // should do additional processing here
      }
      else {
        cout<<"This is an NXDOMAIN situation"<<endl;
        const auto& rrset = fnd->zone->rrsets[DNSType::SOA];
        response.dh.rcode = (int)RCode::Nxdomain;
        response.putRR(DNSSection::Authority, zone, DNSType::SOA, rrset.ttl, rrset.contents[0]);
      }
    }
    else {
      cout<<"Found something in zone '"<<zone<<"' for lhs '"<<name<<"', searchname now '"<<searchname<<"', lastnode '"<<lastnode<<"', passedZonecut="<<passedZonecut<<endl;
      
      auto iter = node->rrsets.cbegin();
      if(type == DNSType::ANY) {
        for(const auto& t : node->rrsets) {
          const auto& rrset = t.second;
          for(const auto& rr : rrset.contents) {
            response.putRR(DNSSection::Answer, lastnode+zone, t.first, rrset.ttl, rr);
          }
        }
      }
      else if(iter = node->rrsets.find(type), iter != node->rrsets.end()) {
        const auto& rrset = iter->second;
        for(const auto& rr : rrset.contents) {
          response.putRR(DNSSection::Answer, lastnode+zone, type, rrset.ttl, rr);
        }
      }
      else if(iter = node->rrsets.find(DNSType::CNAME), iter != node->rrsets.end()) {
        cout<<"We do have a CNAME!"<<endl;
        const auto& rrset = iter->second;
        for(const auto& rr : rrset.contents) {
          response.putRR(DNSSection::Answer, lastnode+zone, DNSType::CNAME, rrset.ttl, rr);
        }
        cout<<" We should actually follow this, at least within our zone"<<endl;
      }
      else {
        cout<<"Node exists, qtype doesn't, NOERROR situation, inserting SOA"<<endl;
        const auto& rrset = fnd->zone->rrsets[DNSType::SOA];
        response.putRR(DNSSection::Answer, zone, DNSType::SOA, rrset.ttl, rrset.contents[0]);
      }
    }
  }
  else {
    response.dh.rcode = (uint8_t)RCode::Refused;
  }
  return true;
}
catch(std::exception& e) {
  cout<<"Error processing query: "<<e.what()<<endl;
  return false;
}

void udpThread(ComboAddress local, Socket* sock, const DNSNode* zones)
{
  for(;;) {
    ComboAddress remote(local);
    DNSMessage dm;
    string message = SRecvfrom(*sock, sizeof(dm), remote);
    if(message.size() < sizeof(dnsheader)) {
      cerr<<"Dropping query from "<<remote.toStringWithPort()<<", too short"<<endl;
      continue;
    }
    memcpy(&dm, message.c_str(), message.size());

    if(dm.dh.qr) {
      cerr<<"Dropping non-query from "<<remote.toStringWithPort()<<endl;
      continue;
    }

    DNSMessage response;
    if(processQuestion(*zones, dm, local, remote, response)) {
      SSendto(*sock, response.serialize(), remote);
    }
  }
}

void tcpClientThread(ComboAddress local, ComboAddress remote, int s, const DNSNode* zones)
{
  Socket sock(s);
  cout<<"TCP Connection from "<<remote.toStringWithPort()<<endl;
  for(;;) {

    uint16_t len;
    
    string message = SRead(sock, 2);
    if(message.size() != 2)
      break;
    memcpy(&len, &message.at(1)-1, 2);
    len=htons(len);
    
    if(len > 512) {
      cerr<<"Remote "<<remote.toStringWithPort()<<" sent question that was too big"<<endl;
      return;
    }
    
    if(len < sizeof(dnsheader)) {
      cerr<<"Dropping query from "<<remote.toStringWithPort()<<", too short"<<endl;
      return;
    }

    cout<<"Reading "<<len<<" bytes"<<endl;
    
    message = SRead(sock, len);
    DNSMessage dm;
    memcpy(&dm, message.c_str(), message.size());

    if(dm.dh.qr) {
      cerr<<"Dropping non-query from "<<remote.toStringWithPort()<<endl;
      return;
    }

    dnsname name;
    DNSType type;
    dm.getQuestion(name, type);

    if(type == DNSType::AXFR) {
      cout<<"Should do AXFR for "<<name<<endl;
      return;
    }
    else {
      dm.payload.rewind();
      DNSMessage response;
      
      if(processQuestion(*zones, dm, local, remote, response)) {
        string ser="00"+response.serialize();
        cout<<"Should send a message of "<<ser.size()<<" bytes in response"<<endl;
        len = htons(ser.length()-2);
        ser[0] = *((char*)&len);
        ser[1] = *(((char*)&len) + 1);
        SWriten(sock, ser);
        cout<<"Sent!"<<endl;
      }
      else
        return;
    }
  }
}


void loadZones(DNSNode& zones)
{
  auto zone = zones.add({"powerdns", "org"});
  zone->zone = new DNSNode(); // XXX ICK
  zone->zone->rrsets[DNSType::SOA]={{serializeSOARecord({"ns1", "powerdns", "org"}, {"admin", "powerdns", "org"}, 1)}};
  zone->zone->rrsets[DNSType::A]={{serializeARecord("1.2.3.4")}, 300};
  zone->zone->rrsets[DNSType::AAAA]={{serializeAAAARecord("::1"), serializeAAAARecord("2001::1")}, 900};
  zone->zone->rrsets[DNSType::NS]={{serializeDNSName({"ns1", "powerdns", "org"})}, 300};

  zone->zone->add({"www"})->rrsets[DNSType::CNAME]={{serializeDNSName({"server1","powerdns","org"})}};

  zone->zone->add({"server1"})->rrsets[DNSType::A]={{serializeARecord("213.244.168.210")}};
  
  //  zone->zone->add({"*"})->rrsets[(dnstype)DNSType::A]={"\x05\x06\x07\x08"};

  zone->zone->add({"fra"})->rrsets[DNSType::NS]={{serializeDNSName({"ns1","fra","powerdns","org"})}};

  zone->zone->add({"ns1", "fra"})->rrsets[DNSType::A]={{serializeARecord("12.13.14.15")}, 86400};
}

int main(int argc, char** argv)
{
  ComboAddress local(argv[1], 53);

  Socket udplistener(local.sin4.sin_family, SOCK_DGRAM);
  SBind(udplistener, local);

  Socket tcplistener(local.sin4.sin_family, SOCK_STREAM);
  SSetsockopt(tcplistener, SOL_SOCKET, SO_REUSEPORT, 1);
  SBind(tcplistener, local);
  SListen(tcplistener, 10);
  
  DNSNode zones;

  loadZones(zones);
  
  thread udpServer(udpThread, local, &udplistener, &zones);
  


  for(;;) {
    ComboAddress remote;
    int client = SAccept(tcplistener, remote);
    thread t(tcpClientThread, local, remote, client, &zones);
    t.detach();
  }
  udpServer.join();
}