#include "TCP_Server.h"

#include <array>

namespace Psc::asio_socket {

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

void TCP_Server::tick() {
    if (state == Working) {
        start_accept();
        flush_clients();
    }

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

    if (state == Working) {
        flush_clients();
        cleanup_closed_clients();
    }
}

TCP_Server & TCP_Server::set_tcp_no_delay(bool value) {
    no_delay = value;
    return *this;
}

TCP_Server& TCP_Server::set_connect_system_buffer_size(size_t size) {
    connect_system_buffer_size = size;
    return *this;
}

TCP_Server& TCP_Server::set_connect_user_buffer_size(size_t size) {
    connect_user_buffer_size = size;
    return *this;
}

TCP_Server& TCP_Server::set_recv_system_buffer_size(size_t size) {
    recv_system_buffer_size = size;
    return *this;
}

TCP_Server& TCP_Server::create() {
    acceptor = std::make_unique<asio::ip::tcp::acceptor>(io_context);
    socket_fd = socket_id(acceptor.get());
    state = Not_bind;
    return *this;
}

bool TCP_Server::listen(const std::string& ip, std::uint32_t port) {
    return listen(Sockaddr_In(ip, port));
}

bool TCP_Server::listen(const Sockaddr_In& addr) {
    if (!acceptor) create();
    if (addr.ip.empty() || addr.port == 0) {
        state = Not_Set_Field;
        return false;
    }

    asio::error_code ec;
    const auto address_value = asio::ip::make_address(addr.ip, ec);
    if (ec) {
        state = Not_bind;
        return false;
    }

    asio::ip::tcp::endpoint endpoint(address_value, static_cast<unsigned short>(addr.port));
    acceptor->open(endpoint.protocol(), ec);
    if (ec) return false;
    acceptor->set_option(asio::socket_base::reuse_address(true), ec);
    acceptor->bind(endpoint, ec);
    if (ec) {
        state = Not_bind;
        return false;
    }
    acceptor->listen(asio::socket_base::max_listen_connections, ec);
    if (ec) {
        state = Not_Listen;
        return false;
    }
    address = addr;
    state = Working;
    start_accept();
    return true;
}

void TCP_Server::close() {
    accept_pending = false;
    if (pending_accept_socket) {
        asio::error_code ec;
        pending_accept_socket->cancel(ec);
        pending_accept_socket->close(ec);
    }
    pending_accept_socket.reset();

    for (auto& [_, conn] : tcp_clients) {
        close_client(conn);
    }
    tcp_clients.clear();
    accepted_clients.clear();

    if (acceptor) {
        asio::error_code ec;
        acceptor->cancel(ec);
        acceptor->close(ec);
    }
    socket_fd = static_cast<Socket_FD>(-1);
    state = Not_Created;
}

std::shared_ptr<TCP_Connect> TCP_Server::accept() const {
    while (!accepted_clients.empty()) {
        auto conn = accepted_clients.front().lock();
        accepted_clients.erase(accepted_clients.begin());
        if (conn && !conn->closing) return conn;
    }
    return nullptr;
}

void TCP_Server::start_accept() {
    if (accept_pending || !acceptor || state != Working || !acceptor->is_open()) return;

    pending_accept_socket = std::make_shared<asio::ip::tcp::socket>(io_context);
    accept_pending = true;
    acceptor->async_accept(*pending_accept_socket,
        [this](const asio::error_code& ec) {
            accept_pending = false;
            if (state != Working) return;
            if (ec == asio::error::operation_aborted) return;
            if (!ec && pending_accept_socket) {
                asio::error_code option_ec;
                pending_accept_socket->set_option(asio::ip::tcp::no_delay(no_delay), option_ec);
                pending_accept_socket->set_option(
                    asio::socket_base::send_buffer_size(static_cast<int>(connect_system_buffer_size)),
                    option_ec);
                pending_accept_socket->set_option(
                    asio::socket_base::receive_buffer_size(static_cast<int>(recv_system_buffer_size)),
                    option_ec);

                auto conn = std::make_shared<TCP_Connect>();
                conn->socket = pending_accept_socket;
                conn->send_buffer.init(connect_user_buffer_size);
                conn->recv_buffer.init(connect_user_buffer_size);
                conn->info.fd = socket_id(conn->socket.get());
                conn->info.sockaddr = endpoint_to_sockaddr(conn->socket->remote_endpoint(option_ec));
                tcp_clients[conn->info.fd] = conn;
                accepted_clients.push_back(conn);
                start_read(conn);
                start_write(conn);
            }
            pending_accept_socket.reset();
            start_accept();
        });
}

void TCP_Server::start_read(const std::shared_ptr<TCP_Connect>& conn) {
    if (!conn || conn->closing || conn->read_pending || !conn->socket || !conn->socket->is_open()) return;

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

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

void TCP_Server::start_write(const std::shared_ptr<TCP_Connect>& conn) {
    if (!conn || conn->closing || conn->write_pending || !conn->socket || !conn->socket->is_open()) return;

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

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

            conn->send_buffer.skip(sent);
            conn->push_speed.update(static_cast<unsigned int>(sent));
            conn->send_num.update(static_cast<double>(sent));
            if (sent != 0) start_write(conn);
        });
}

void TCP_Server::close_client(const std::shared_ptr<TCP_Connect>& conn) {
    if (!conn || conn->closing) return;
    conn->closing = true;
    if (conn->socket) {
        asio::error_code ec;
        conn->socket->cancel(ec);
        conn->socket->shutdown(asio::ip::tcp::socket::shutdown_both, ec);
        conn->socket->close(ec);
    }
}

void TCP_Server::cleanup_closed_clients() {
    for (auto it = tcp_clients.begin(); it != tcp_clients.end();) {
        if (!it->second || it->second->closing || !it->second->socket || !it->second->socket->is_open()) {
            it = tcp_clients.erase(it);
        } else {
            ++it;
        }
    }

    for (auto it = accepted_clients.begin(); it != accepted_clients.end();) {
        auto conn = it->lock();
        if (!conn || conn->closing) {
            it = accepted_clients.erase(it);
        } else {
            ++it;
        }
    }
}

void TCP_Server::flush_clients() {
    start_accept();
    for (auto& [_, conn] : tcp_clients) {
        start_read(conn);
        start_write(conn);
    }
}

void TCP_Server::write_to_all_clients(const std::string& data) {
    if (data.empty()) return;
    for (auto& [_, conn] : tcp_clients) {
        if (!conn || conn->closing) continue;
        auto written = conn->send_buffer.write_best_effort(data.data(), data.size());
        if (written < data.size()) conn->lose_speed.update(static_cast<unsigned int>(data.size() - written));
        start_write(conn);
    }
}

std::vector<TCP_Server::Read_Info> TCP_Server::read_from_all_clients() {
    std::vector<Read_Info> ret;
    std::array<char, 16 * 1024> buffer{};

    for (auto& [_, conn] : tcp_clients) {
        if (!conn || conn->closing) continue;

        std::string data;
        for (;;) {
            auto n = conn->recv_buffer.read_best_effort(buffer.data(), buffer.size());
            if (n == 0) break;
            data.append(buffer.data(), n);
        }
        if (!data.empty()) ret.push_back({conn->info, std::move(data)});
    }

    cleanup_closed_clients();
    return ret;
}

std::vector<Socket_FD> TCP_Server::client_fds() {
    std::vector<Socket_FD> ret;
    for (const auto& [fd, conn] : tcp_clients) {
        if (conn && !conn->closing) ret.push_back(fd);
    }
    return ret;
}

std::vector<std::shared_ptr<TCP_Connect>> TCP_Server::get_all_clients() {
    std::vector<std::shared_ptr<TCP_Connect>> ret;
    for (const auto& [_, conn] : tcp_clients) {
        if (conn && !conn->closing) ret.push_back(conn);
    }
    return ret;
}

std::string TCP_Server::to_string() {
    return "TCP_Server:[" + (address ? address->to_string() : std::string{}) + "]";
}

} // namespace Psc::asio_socket