#ifndef _TCP_SERVER_HPP_
#define _TCP_SERVER_HPP_

#ifndef ASIO_STANDALONE
#define ASIO_STANDALONE
#endif
#include <asio.hpp>

#include <type_traits>

#include "buffers.hpp"

using asio::ip::tcp;

template<typename SessionType>
using tcp_cb_connected = std::function<std::shared_ptr<SessionType>(tcp::socket&)>;
template<typename SessionType>
using tcp_cb_received = std::function<void(SessionType*, const unsigned char* data, size_t length)>;
template<typename SessionType>
using tcp_cb_disconnected = std::function<void(SessionType*)>;

template<typename SessionType>
class TCPSession : public std::enable_shared_from_this<TCPSession<SessionType>> {
public:
    TCPSession(tcp::socket sock) : socket(std::move(sock)) {
        static_assert(std::is_base_of<TCPSession<SessionType>, SessionType>::value , "SessionType should derive from TCPSession<SessionType>");
    }

    void ReadData() {
        socket.async_read_some(asio::buffer(read_buffer, 1536), [this, self = this->shared_from_this()](std::error_code ec, std::size_t length) {
            if (!ec) {
                if(received)
                    received(static_cast<SessionType*>(this), read_buffer, length);
                ReadData();
            } else {
                if(disconnected)
                    disconnected(static_cast<SessionType*>(this));
            }
        });
    }

    void SendData(StaticBuffer&& buffer, size_t length) {
        auto abuf = asio::buffer(buffer.Data(), length);
        socket.async_write_some(abuf, [buf = std::move(buffer)](std::error_code ec, size_t length) {});
    }

    void Close() {
        socket.shutdown(tcp::socket::shutdown_both);
        socket.close();
    }

    tcp::socket& UnderlyingSocket() { return socket; }
    
    SessionType& SetReceivedCallback(const tcp_cb_received<SessionType>& cb) { received = cb; return static_cast<SessionType&>(*this);}
    SessionType& SetDisconnectedCallback(const tcp_cb_disconnected<SessionType>& cb) { disconnected = cb; return static_cast<SessionType&>(*this); }

private:
    tcp::socket socket;
    tcp_cb_received<SessionType> received;
    tcp_cb_disconnected<SessionType> disconnected;
    unsigned char read_buffer[1536];
};

template<typename SessionType>
class TCPServer : public std::enable_shared_from_this<TCPServer<SessionType>> {
public:
    TCPServer(asio::io_context& ctx, short port) : acceptor(ctx, tcp::endpoint(tcp::v4(), port))  {}

    void Start() {
        if(started)
            return;
        started = true;
        _Accept();
    }

    void Stop() {
        if(!started)
            return;
        started = false;
        acceptor.close();
    }

    TCPServer& SetConnectedCallback(const tcp_cb_connected<SessionType>& cb) { connected = cb; return *this; }

protected:
    void _Accept() {
        acceptor.async_accept([this, self = this->shared_from_this()](std::error_code ec, tcp::socket socket) {
        if (!ec) {
            if(connected) {
                auto session = connected(socket);
                if(session != nullptr)
                    session->ReadData();
            }
        }
        if(started)
            _Accept();
        });
    }

    bool started = false;    
    tcp_cb_connected<SessionType> connected;
    tcp::acceptor acceptor;
};

using connect_fail_cb = std::function<void(const std::string&)>;

template<typename SessionType>
class TCPClient : public std::enable_shared_from_this<TCPClient<SessionType>> {
public:
    TCPClient(asio::io_context& ctx) : socket(ctx) {}
    void Connect(const char* ipaddr, short port, uint32_t timeout_time, connect_fail_cb cb = nullptr) {
        tcp::endpoint endpoint(asio::ip::make_address_v4(ipaddr), port);
        auto timeout = std::make_unique<asio::steady_timer>(socket.get_io_context(), asio::chrono::milliseconds(timeout_time));
        timeout->async_wait([this, cb](std::error_code ec) {
            if(!ec) {
                if(cb)
                    cb(ec.message());
                socket.cancel();
            }
        });
        socket.async_connect(endpoint, [this, cb, timeout = std::move(timeout), self = this->shared_from_this()](std::error_code ec) {
            if (!ec) {
                if(connected) {
                    auto session = connected(socket);
                    if(session != nullptr)
                        session->ReadData();
                }
            } else {
                if(cb)
                    cb(ec.message());
            }
            timeout->cancel();
        });
    }

    TCPClient& SetConnectedCallback(tcp_cb_connected<SessionType> cb) { connected = cb; return *this; }

protected:
    tcp::socket socket;
    tcp_cb_connected<SessionType> connected;
};

#endif