#include "TCP_Client.h"

#include <array>

namespace Psc::asio_socket {

TCP_Client::~TCP_Client() {
    close();
}

void TCP_Client::set_state(State s) {
    const auto old = state.exchange(s);
    if (old != s && on_state_change) on_state_change(old, s);
}

void TCP_Client::create() {
    send_buffer.init(max_buffer_size);
    recv_buffer.init(max_buffer_size);
    connect_pending = false;
    read_pending = false;
    write_pending = false;
    send_storage.clear();
    resolver = std::make_unique<asio::ip::tcp::resolver>(io_context);
    socket = std::make_unique<asio::ip::tcp::socket>(io_context);
    socket_fd = socket_id(socket.get());
    set_state(Reconnecting);
}

void TCP_Client::close() {
    connect_pending = false;
    read_pending = false;
    write_pending = false;
    if (resolver) {
        asio::error_code ec;
        resolver->cancel();
    }
    if (socket) {
        asio::error_code ec;
        socket->cancel(ec);
        socket->shutdown(asio::ip::tcp::socket::shutdown_both, ec);
        socket->close(ec);
    }
    socket_fd = static_cast<Socket_FD>(-1);
    set_state(Not_Created);
}

void TCP_Client::schedule_recreate() {
    connect_pending = false;
    read_pending = false;
    write_pending = false;
    if (resolver) resolver->cancel();
    if (socket) {
        asio::error_code ec;
        socket->cancel(ec);
        socket->close(ec);
    }
    next_reconnect_tp = std::chrono::steady_clock::now() + std::chrono::milliseconds(reconnect_backoff_ms);
    set_state(Reconnecting);
}

bool TCP_Client::start_connect() {
    if (connect_pending || state == Connecting || state == Connected) return true;
    if (dest_address.ip.empty() || dest_address.port == 0) {
        set_state(Not_Set_Field);
        return false;
    }

    io_context.restart();
    resolver = std::make_unique<asio::ip::tcp::resolver>(io_context);
    socket = std::make_unique<asio::ip::tcp::socket>(io_context);
    connect_pending = true;
    set_state(Connecting);
    resolver->async_resolve(dest_address.ip, std::to_string(dest_address.port),
        [this](const asio::error_code& ec, asio::ip::tcp::resolver::results_type endpoints) {
            if (state != Connecting || !socket) return;
            if (ec) {
                schedule_recreate();
                return;
            }

            asio::async_connect(*socket, endpoints,
                [this](const asio::error_code& connect_ec, const asio::ip::tcp::endpoint&) {
                    connect_pending = false;
                    if (state != Connecting || !socket) return;
                    if (connect_ec) {
                        schedule_recreate();
                        return;
                    }

                    socket_fd = socket_id(socket.get());
                    address = dest_address;
                    set_state(Connected);
                    start_read();
                    start_write();
                });
        });
    return true;
}

void TCP_Client::tick() {
    if (state == Reconnecting && std::chrono::steady_clock::now() >= next_reconnect_tp) {
        start_connect();
    }
    if (state == Connected) {
        start_read();
        start_write();
    }

    io_context.restart();
    io_context.poll();

    if (state == Connected) {
        start_read();
        start_write();
    }
}

void TCP_Client::start_read() {
    if (read_pending || state != Connected || !socket || !socket->is_open()) return;

    read_pending = true;
    socket->async_read_some(asio::buffer(recv_storage),
        [this](const asio::error_code& ec, std::size_t n) {
            read_pending = false;
            if (state != Connected) return;
            if (ec == asio::error::operation_aborted) return;
            if (ec == asio::error::eof || ec == asio::error::connection_reset || ec) {
                schedule_recreate();
                return;
            }

            recv_buffer.write_best_effort(recv_storage.data(), n);
            start_read();
        });
}

void TCP_Client::start_write() {
    if (write_pending || state != Connected || !socket || !socket->is_open()) return;

    send_storage.resize(16 * 1024);
    auto size = send_buffer.peek_best_effort(send_storage.data(), send_storage.size());
    if (size == 0) return;
    send_storage.resize(size);

    write_pending = true;
    socket->async_write_some(asio::buffer(send_storage.data(), send_storage.size()),
        [this](const asio::error_code& ec, std::size_t sent) {
            write_pending = false;
            if (state != Connected) return;
            if (ec == asio::error::operation_aborted) return;
            if (ec) {
                schedule_recreate();
                return;
            }

            send_buffer.skip(sent);
            if (sent != 0) start_write();
        });
}

std::string TCP_Client::read() {
    std::string ret;
    std::array<char, 16 * 1024> buffer{};
    for (;;) {
        auto n = recv_buffer.read_best_effort(buffer.data(), buffer.size());
        if (n == 0) break;
        ret.append(buffer.data(), n);
    }
    return ret;
}

void TCP_Client::send(const std::string& data) {
    if (data.empty()) return;
    send_buffer.write_best_effort(data.data(), data.size());
    if (state == Connected) start_write();
}

std::string TCP_Client::to_string() {
    return "TCP_Client:[" + dest_address.to_string() + "]";
}

} // namespace Psc::asio_socket