#pragma once
#include <stdint.h>
#include <netinet/in.h>
#include <vector>
#include <string>

/* --------------------------------------- */

#define STUN_IP4 0x01
#define STUN_IP6 0x02

#define STUN_MSG_TYPE_NONE                     0x0000
#define STUN_MSG_TYPE_BINDING_REQUEST          0x0001
#define STUN_MSG_TYPE_BINDING_RESPONSE_SUCCESS 0x0101
#define STUN_MSG_TYPE_BINDING_RESPONSE_ERROR   0x0111
#define STUN_MSG_TYPE_BINDING_INDICATION       0x0011
                                               
#define STUN_ATTR_TYPE_NONE                    0x0000
#define STUN_ATTR_TYPE_MAPPED_ADDR             0x0001
#define STUN_ATTR_TYPE_CHANGE_REQ              0x0003
#define STUN_ATTR_TYPE_USERNAME                0x0006
#define STUN_ATTR_TYPE_MESSAGE_INTEGRITY       0x0008
#define STUN_ATTR_TYPE_ERR_CODE                0x0009
#define STUN_ATTR_TYPE_UNKNOWN_ATTRIBUTES      0x000a
#define STUN_ATTR_TYPE_CHANNEL_NUMBER          0x000c
#define STUN_ATTR_TYPE_LIFETIME                0x000d
#define STUN_ATTR_TYPE_XOR_PEER_ADDR           0x0012
#define STUN_ATTR_TYPE_DATA                    0x0013
#define STUN_ATTR_TYPE_REALM                   0x0014
#define STUN_ATTR_TYPE_NONCE                   0x0015
#define STUN_ATTR_TYPE_XOR_RELAY_ADDRESS       0x0016
#define STUN_ATTR_TYPE_REQ_ADDRESS_FAMILY      0x0017
#define STUN_ATTR_TYPE_EVEN_PORT               0x0018
#define STUN_ATTR_TYPE_REQUESTED_TRANSPORT     0x0019
#define STUN_ATTR_TYPE_DONT_FRAGMENT           0x001a
#define STUN_ATTR_TYPE_XOR_MAPPED_ADDRESS      0x0020
#define STUN_ATTR_TYPE_RESERVATION_TOKEN       0x0022
#define STUN_ATTR_TYPE_PRIORITY                0x0024
#define STUN_ATTR_TYPE_USE_CANDIDATE           0x0025
#define STUN_ATTR_TYPE_PADDING                 0x0026
#define STUN_ATTR_TYPE_RESPONSE_PORT           0x0027
#define STUN_ATTR_TYPE_SOFTWARE                0x8022
#define STUN_ATTR_TYPE_ALTERNATE_SERVER        0x8023
#define STUN_ATTR_TYPE_FINGERPRINT             0x8028
#define STUN_ATTR_TYPE_ICE_CONTROLLED          0x8029
#define STUN_ATTR_TYPE_ICE_CONTROLLING         0x802a
#define STUN_ATTR_TYPE_RESPONSE_ORIGIN         0x802b
#define STUN_ATTR_TYPE_OTHER_ADDRESS           0x802c

/* --------------------------------------- */

std::string stun_message_type_to_string(uint16_t type);
std::string stun_attribute_type_to_string(uint16_t type);
std::string stun_family_type_to_string(uint8_t type);
 
/* 
   Compute the hmac-sha1 over message.
   uint8_t* message:  the data over which we compute the hmac sha
   uint32_t nbytes:   the number of bytse in message
   std::string key:   key to use for hmac 
   uint8_t* output:   we write the sha1 into this buffer.
*/
int stun_compute_hmac_sha1(uint8_t* message, uint32_t nbytes, std::string key, uint8_t* output);

/* 
   Compute the Message-Integrity of a stun message. 
   This will not change the given buffer.
     
   std::vector<uint8_t>& buffer: the buffer that contains a valid stun message 
   std::string key:              key to use for hmac 
   uint8_t* output:              will be filled with the correct hmac-sha1 of that represents the integrity message value. 
*/
int stun_compute_message_integrity(std::vector<uint8_t>& buffer, std::string key, uint8_t* output);

/* 
   Compute the fingerprint value for the stun message.
   This will not change the given buffer.
   std::vector<uint8_t>& buffer:    the buffer that contains a valid stun message.
   uint32_t& result:                will be set to the calculated crc value.
*/
int stun_compute_fingerprint(std::vector<uint8_t>& buffer, uint32_t& result);

/* --------------------------------------- */

/* https://tools.ietf.org/html/rfc5389#section-15.10 */
class StunAttribSoftware {
public:
  char* value;
};
  
class StunAttribFingerprint {
public:
  uint32_t value;
};
  
/* https://tools.ietf.org/html/rfc5389#section-15.4 */
class StunAttribMessageIntegrity {
public:
  uint8_t* sha1;
};

/* https://tools.ietf.org/html/rfc5245#section-19.1 */
class StunAttribPriority {
public:
  uint32_t value;
};

/* https://tools.ietf.org/html/rfc5245#section-19.1 */
class StunAttribIceControllling {
public:
  uint64_t tie_breaker;
};

/* https://tools.ietf.org/html/rfc3489#section-11.2.6 */
class StunAttribUsername {
public:
  char* value;     /* Must use `length` member of attribute that indicates the number of valid bytes in the username. */
};

/* https://tools.ietf.org/html/rfc5389#section-15.2 */
class StunAttribXorMappedAddress {
public:
  uint8_t family;
  uint16_t port;
  uint8_t ip[16];
};
  
/* --------------------------------------- */

class StunAttribute {
public:
  StunAttribute();
  void print();
    
public:
  uint16_t type;
  uint16_t length;
  union {
    StunAttribXorMappedAddress xor_address;
    StunAttribUsername username;
    StunAttribIceControllling ice_controlling;
    StunAttribPriority priority;
    StunAttribSoftware software;
    StunAttribMessageIntegrity message_integrity;
    StunAttribFingerprint fingerprint;
  };
};

/* --------------------------------------- */

class StunMessage {
public:
  StunMessage();
  void setType(uint16_t type);
  void setTransactionId(uint32_t a, uint32_t b, uint32_t c);
  void print();
  void addAttribute(StunAttribute& attr);
  void removeAttributes();
  StunAttribute* getAttributeByType(uint16_t type);
    
public:
  uint16_t type;
  uint16_t length;
  uint32_t cookie;
  uint32_t transaction_id[3];
  std::vector<StunAttribute> attributes;
};

/* --------------------------------------- */

class StunReader {
public:
  StunReader();
  int parse(uint8_t* data, size_t nbytes, size_t& nparsed, StunMessage& msg); /* `nparsed` and `msg` are filled. */

private:
  int parseXorMappedAddress(StunAttribute& attr);
  int parseUsername(StunAttribute& attr);
  int parseIceControlling(StunAttribute& attr);
  int parsePriority(StunAttribute& attr);
  int parseSoftware(StunAttribute& attr);
  int parseMessageIntegrity(StunAttribute& attr);
  int parseFingerprint(StunAttribute& attr);
    
  uint8_t readU8();
  uint16_t readU16();
  uint32_t readU32();
  uint64_t readU64();
    
private:
  uint8_t* buffer_data;
  size_t buffer_size;
  size_t read_dx;
};

/* --------------------------------------- */

class StunWriter {
public:
  StunWriter();
  
  /* write header and finalize. call for each stun message */
  int begin(StunMessage& msg, uint8_t paddingByte = 0x00);  /* I've added the padding byte here so that we can use the examples that can be found here https://tools.ietf.org/html/rfc5769#section-2.2 as they use 0x20 or 0x00 as the padding byte which is correct as you are free to use w/e padding byte you want. */
  int end();

  /* write attributes */
  int writeXorMappedAddress(sockaddr_in addr);
  int writeXorMappedAddress(uint8_t family, uint16_t port, uint32_t ip);
  int writeXorMappedAddress(uint8_t family, uint16_t port, const std::string& ip);
  int writeUsername(const std::string& username);
  int writeSoftware(const std::string& software);
  int writeMessageIntegrity(const std::string& password); /* When using WebRtc this is the ice-upwd of the other agent. */
  int writeFingerprint(); /* Must be the last attribute in the message. When adding a fingerprint, make sure that it is added after the message-integrity (when you also use a message-integrity). */

  /* get buffer */
  uint8_t* getBufferPtr();
  size_t getBufferSize();
    
private:
  void writeU8(uint8_t v);
  void writeU16(uint16_t v);
  void writeU32(uint32_t v);
  void writeU64(uint64_t v);
  void rewriteU16(size_t dx, uint16_t v);
  void rewriteU32(size_t dx, uint32_t v);
  void writeString(const std::string& str);
  void writePadding();
  int convertIp4StringToInt(const std::string& ip, uint32_t& result);
    
private:
  std::vector<uint8_t> buffer;
  uint8_t padding_byte;
};
  
/* --------------------------------------- */

inline uint8_t* StunWriter::getBufferPtr() {

  if (0 == buffer.size()) {
    return NULL;
  }

  return &buffer[0];
}

inline size_t StunWriter::getBufferSize() {
  return buffer.size();
}

/* --------------------------------------- */