#ifndef SOCKS_H__
#define SOCKS_H__

#include <memory>
#include <vector>
#include <boost/asio.hpp>
#include "Identity.h"
#include "Streaming.h"
#include "I2PTunnel.h"

namespace i2p
{
namespace proxy
{

	const size_t socks_buffer_size = 8192;
	const size_t max_socks_hostname_size = 255; // Limit for socks5 and bad idea to traverse

	class SOCKSServer;
	class SOCKSHandler {

		private:
			enum state {
				GET_VERSION,
				SOCKS4A,
				SOCKS5_S1, //Authentication negotiation
				SOCKS5_S2, //Authentication
				SOCKS5_S3, //Request
				READY
			};
			enum parseState {
				GET4A_COMMAND,
				GET4A_PORT1,
				GET4A_PORT2,
				GET4A_IP1,
				GET4A_IP2,
				GET4A_IP3,
				GET4A_IP4,
				GET4A_IDENT,
				GET4A_HOST,
				GET5_AUTHNUM,
				GET5_AUTH,
				GET5_REQUESTV,
				GET5_COMMAND,
				GET5_GETRSV,
				GET5_GETADDRTYPE,
				GET5_IPV4_1,
				GET5_IPV4_2,
				GET5_IPV4_3,
				GET5_IPV4_4,
				GET5_IPV6_1,
				GET5_IPV6_2,
				GET5_IPV6_3,
				GET5_IPV6_4,
				GET5_IPV6_5,
				GET5_IPV6_6,
				GET5_IPV6_7,
				GET5_IPV6_8,
				GET5_IPV6_9,
				GET5_IPV6_10,
				GET5_IPV6_11,
				GET5_IPV6_12,
				GET5_IPV6_13,
				GET5_IPV6_14,
				GET5_IPV6_15,
				GET5_IPV6_16,
				GET5_HOST_SIZE,
				GET5_HOST,
				GET5_PORT1,
				GET5_PORT2,
				DONE
			};
			enum authMethods {
				AUTH_NONE = 0, //No authentication, skip to next step
				AUTH_GSSAPI = 1, //GSSAPI authentication
				AUTH_USERPASSWD = 2, //Username and password
				AUTH_UNACCEPTABLE = 0xff //No acceptable method found
			};
			enum addrTypes {
				ADDR_IPV4 = 1, //IPv4 address (4 octets)
				ADDR_DNS = 3, // DNS name (up to 255 octets)
				ADDR_IPV6 = 4 //IPV6 address (16 octets)
			};
			enum errTypes {
				SOCKS5_OK = 0, // No error for SOCKS5
				SOCKS5_GEN_FAIL = 1, // General server failure
				SOCKS5_RULE_DENIED = 2, // Connection disallowed by ruleset
				SOCKS5_NET_UNREACH = 3, // Network unreachable
				SOCKS5_HOST_UNREACH = 4, // Host unreachable
				SOCKS5_CONN_REFUSED = 5, // Connection refused by the peer
				SOCKS5_TTL_EXPIRED = 6, // TTL Expired
				SOCKS5_CMD_UNSUP = 7, // Command unsuported
				SOCKS5_ADDR_UNSUP = 8, // Address type unsuported
				SOCKS4_OK = 90, // No error for SOCKS4
				SOCKS4_FAIL = 91, // Failed establishing connecting or not allowed
				SOCKS4_IDENTD_MISSING = 92, // Couldn't connect to the identd server
				SOCKS4_IDENTD_DIFFER = 93 // The ID reported by the application and by identd differ
			};
			enum cmdTypes {
				CMD_CONNECT = 1, // TCP Connect
				CMD_BIND = 2, // TCP Bind
				CMD_UDP = 3 // UDP associate
			};
			enum socksVersions {
				SOCKS4 = 4, // SOCKS4
				SOCKS5 = 5 // SOCKS5
			};


			void GotClientRequest(boost::system::error_code & ecode, std::string & host, uint16_t port);
			std::size_t HandleData(uint8_t *sock_buff, std::size_t len);
			std::size_t HandleVersion(uint8_t *sock_buff);
			std::size_t HandleSOCKS4A(uint8_t *sock_buff, std::size_t len);
			std::size_t HandleSOCKS5Step1(uint8_t *sock_buff, std::size_t len);
			std::size_t HandleSOCKS5Step3(uint8_t *sock_buff, std::size_t len);
			void HandleSockRecv(const boost::system::error_code & ecode, std::size_t bytes_transfered);
			void Terminate();
			void CloseSock();
			void CloseStream();
			void AsyncSockRead();
			void Socks5AuthNegoFailed();
			void Socks5ChooseAuth();
			void SocksRequestFailed();
			void SocksRequestSuccess();
			void SentSocksFailed(const boost::system::error_code & ecode);
			//HACK: we need to pass the shared_ptr to ensure the buffer will live enough
			void SentSocksResponse(const boost::system::error_code & ecode, std::shared_ptr<std::vector<uint8_t>> response);
			void HandleStreamRequestComplete (std::shared_ptr<i2p::stream::Stream> stream);

			uint8_t m_sock_buff[socks_buffer_size];
			SOCKSServer * m_parent;
			boost::asio::ip::tcp::socket * m_sock;
			std::shared_ptr<i2p::stream::Stream> m_stream;
			state m_state;
			parseState m_pstate;
			uint8_t m_command;
			uint16_t m_port;
			uint32_t m_ip;
			std::string m_destination;
			uint8_t m_authleft; //Authentication methods left
			//TODO: this will probably be more elegant as enums
			authMethods m_authchosen; //Authentication chosen
			addrTypes m_addrtype; //Address type chosen
			uint8_t m_addrleft; //Octets of DNS address left
			errTypes m_error; //Error cause
			socksVersions m_socksv; //Socks version
			bool m_need_more; //The parser still needs to receive more data

		public:
			SOCKSHandler(SOCKSServer * parent, boost::asio::ip::tcp::socket * sock) : 
				m_parent(parent), m_sock(sock), m_stream(nullptr), m_state(GET_VERSION),
				m_authchosen(AUTH_UNACCEPTABLE), m_addrtype(ADDR_IPV4), m_error(SOCKS5_GEN_FAIL)
				{ AsyncSockRead(); m_destination.reserve(max_socks_hostname_size+1); }
			~SOCKSHandler() { CloseSock(); CloseStream(); }
	};

	class SOCKSServer: public i2p::client::I2PTunnel
	{
		public:
			SOCKSServer(int port) : I2PTunnel(nullptr),
				m_Acceptor (GetService (), boost::asio::ip::tcp::endpoint (boost::asio::ip::tcp::v4(), port)),
				m_Timer (GetService ()) {};
			~SOCKSServer() { Stop(); }

			void Start ();
			void Stop ();

		private:

			void Accept();
			void HandleAccept(const boost::system::error_code& ecode, boost::asio::ip::tcp::socket * socket);

		private:

			boost::asio::ip::tcp::acceptor m_Acceptor;
			boost::asio::deadline_timer m_Timer;
	};	

	typedef SOCKSServer SOCKSProxy;
}
}


#endif