#include "tcp_client.hpp"
#include "logger.hpp"

#include <unistd.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <cstring>
#include <chrono>
#include <thread>
#include <fcntl.h>
#include <sys/select.h>
#include <sys/time.h>

TcpClient::TcpClient(const std::string &id, const std::string &ip, int port)
    : id_(id), ip_(ip), port_(port)
{
}

TcpClient::~TcpClient()
{
    stop();
}

void TcpClient::start()
{
    if (running_)
        return;
    running_ = true;
    worker_ = std::thread(&TcpClient::client_loop, this);
}

void TcpClient::stop()
{
    running_ = false;
    close_socket();
    if (worker_.joinable())
        worker_.join();
}

bool TcpClient::is_connected() const
{
    return connected_;
}

void TcpClient::client_loop()
{
    while (running_)
    {
        auto start = std::chrono::steady_clock::now();
        bool success = try_connect();
        auto end = std::chrono::steady_clock::now();

        if (!success)
        {
            LOG_WARN("[" + id_ + "] Connection failed, retrying in 5s...");
            auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
            int remaining = 5000 - static_cast<int>(elapsed);
            if (remaining > 0)
                std::this_thread::sleep_for(std::chrono::milliseconds(remaining));
            continue;
        }

        connected_ = true;
        if (status_callback_)
            status_callback_(true);

        LOG_INFO("[" + id_ + "] Connected to " + ip_ + ":" + std::to_string(port_));

        char buffer[1024];

        while (running_ && connected_)
        {
            ssize_t n = recv(sock_fd_, buffer, sizeof(buffer), 0);
            if (n > 0)
            {
                std::string data(buffer, n);
                if (receive_callback_)
                    receive_callback_(data);
            }
            else if (n == 0)
            {
                LOG_WARN("[" + id_ + "] Connection closed by remote");
                break;
            }
            else
            {
                LOG_ERROR("[" + id_ + "] recv() failed: " + std::string(strerror(errno)));
                break;
            }
        }

        connected_ = false;
        close_socket();

        if (status_callback_)
            status_callback_(false);

        std::this_thread::sleep_for(std::chrono::seconds(2)); // 防止频繁重连
    }
}

bool TcpClient::try_connect()
{
    close_socket();

    sock_fd_ = socket(AF_INET, SOCK_STREAM, 0);
    if (sock_fd_ < 0)
    {
        LOG_ERROR("[" + id_ + "] socket() failed: " + std::string(strerror(errno)));
        return false;
    }

    // 设置非阻塞
    int flags = fcntl(sock_fd_, F_GETFL, 0);
    if (flags < 0)
    {
        LOG_ERROR("[" + id_ + "] fcntl(F_GETFL) failed: " + std::string(strerror(errno)));
        close_socket();
        return false;
    }

    if (fcntl(sock_fd_, F_SETFL, flags | O_NONBLOCK) < 0)
    {
        LOG_ERROR("[" + id_ + "] fcntl(F_SETFL) failed: " + std::string(strerror(errno)));
        close_socket();
        return false;
    }

    sockaddr_in server_addr{};
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port_);

    if (inet_pton(AF_INET, ip_.c_str(), &server_addr.sin_addr) <= 0)
    {
        LOG_ERROR("[" + id_ + "] Invalid IP: " + ip_);
        close_socket();
        return false;
    }

    int ret = connect(sock_fd_, (struct sockaddr *)&server_addr, sizeof(server_addr));
    if (ret == 0)
    {
        // 立刻成功
        // 恢复阻塞
        fcntl(sock_fd_, F_SETFL, flags);
        return true;
    }
    else if (ret < 0 && errno != EINPROGRESS)
    {
        LOG_ERROR("[" + id_ + "] connect() failed: " + std::string(strerror(errno)));
        close_socket();
        return false;
    }

    // 等待写事件
    fd_set writefds;
    FD_ZERO(&writefds);
    FD_SET(sock_fd_, &writefds);

    struct timeval tv;
    tv.tv_sec = 3; // 超时时间3秒
    tv.tv_usec = 0;

    ret = select(sock_fd_ + 1, nullptr, &writefds, nullptr, &tv);
    if (ret <= 0)
    {
        LOG_ERROR("[" + id_ + "] connect() timeout or select() error");
        close_socket();
        return false;
    }

    // 检查连接状态
    int err = 0;
    socklen_t len = sizeof(err);
    if (getsockopt(sock_fd_, SOL_SOCKET, SO_ERROR, &err, &len) < 0)
    {
        LOG_ERROR("[" + id_ + "] getsockopt() failed: " + std::string(strerror(errno)));
        close_socket();
        return false;
    }

    if (err != 0)
    {
        LOG_ERROR("[" + id_ + "] connect() failed after select: " + std::string(strerror(err)));
        close_socket();
        return false;
    }

    // 恢复阻塞
    fcntl(sock_fd_, F_SETFL, flags);

    return true;
}

void TcpClient::close_socket()
{
    if (sock_fd_ >= 0)
    {
        ::close(sock_fd_);
        sock_fd_ = -1;
    }
}

bool TcpClient::send_data(const std::string &data)
{
    std::lock_guard<std::mutex> lock(send_mutex_);

    if (sock_fd_ < 0)
        return false;

    ssize_t sent = send(sock_fd_, data.data(), data.size(), 0);
    if (sent < 0)
    {
        LOG_ERROR("[" + id_ + "] send() failed: " + std::string(strerror(errno)));
        return false;
    }

    return true;
}

void TcpClient::set_receive_callback(ReceiveCallback cb)
{
    receive_callback_ = cb;
}

void TcpClient::set_status_callback(StatusCallback cb)
{
    status_callback_ = cb;
}
