#pragma once

#include "ASIO_Utils.h"
#include "Core/Statistics/Statistics.h"

#include <array>
#include <map>
#include <memory>
#include <vector>

namespace Psc::asio_socket {

class TCP_Connect {
public:
    Accept_Info info;
    std::shared_ptr<asio::ip::tcp::socket> socket;
    ASIO_StreamRingBuffer send_buffer;
    ASIO_StreamRingBuffer recv_buffer;
    Speed_Statistics push_speed;
    Speed_Statistics lose_speed;
    Value_Statistics send_num;
    std::array<char, 16 * 1024> recv_storage{};
    std::vector<char> send_storage;
    bool read_pending = false;
    bool write_pending = false;
    bool closing = false;

    [[nodiscard]] std::string to_string() const {
        return info.to_string();
    }
};

class TCP_Server : public Socket_Base {
public:
    ~TCP_Server() override;
    void tick();
    TCP_Server& set_tcp_no_delay(bool no_delay);
    void close() override;
    TCP_Server& create();
    TCP_Server& set_connect_system_buffer_size(size_t size);
    TCP_Server& set_connect_user_buffer_size(size_t size);
    TCP_Server& set_recv_system_buffer_size(size_t size);
    bool listen(const std::string& ip, std::uint32_t port);
    bool listen(const Sockaddr_In& address);
    std::string to_string() override;
    void write_to_all_clients(const std::string& data);

    struct Read_Info {
        Accept_Info info;
        std::string data;
    };

    enum State {
        Not_Created,
        Not_Set_Field,
        Not_bind,
        Not_Listen,
        Working
    };

    State state = Not_Created;

    std::vector<Read_Info> read_from_all_clients();
    std::vector<Socket_FD> client_fds();
    std::vector<std::shared_ptr<TCP_Connect>> get_all_clients();
    void flush_clients();
    [[nodiscard]] std::shared_ptr<TCP_Connect> accept() const;

protected:
    mutable asio::io_context io_context;
    std::unique_ptr<asio::ip::tcp::acceptor> acceptor;
    std::map<Socket_FD, std::shared_ptr<TCP_Connect>> tcp_clients;
    size_t connect_user_buffer_size = 1024 * 1024;
    size_t connect_system_buffer_size = 4096 * 10;
    size_t max_pending_send_bytes = 1024 * 1024;
    size_t recv_system_buffer_size = 1024 * 1024;
    bool no_delay = true;
    bool accept_pending = false;
    std::shared_ptr<asio::ip::tcp::socket> pending_accept_socket;
    mutable std::vector<std::weak_ptr<TCP_Connect>> accepted_clients;

    void start_accept();
    void start_read(const std::shared_ptr<TCP_Connect>& conn);
    void start_write(const std::shared_ptr<TCP_Connect>& conn);
    void close_client(const std::shared_ptr<TCP_Connect>& conn);
    void cleanup_closed_clients();
};

} // namespace Psc::asio_socket