#include "mqtt_client.hpp"
#include <iostream>
#include <thread>
#include <chrono>
#include <future>

extern std::atomic<bool> g_running;

MQTTClient::MQTTClient(const MQTTConfig &config)
    : config_(config), connected_(false)
{
    initializeClient();
}

MQTTClient::~MQTTClient()
{
    force_disconnect(); // 析构时直接释放资源，不阻塞
}

void MQTTClient::initializeClient()
{
    std::string address = "tcp://" + config_.server_ip + ":" + std::to_string(config_.server_port);
    client_ = std::make_shared<mqtt::async_client>(address, config_.client_id);
    client_->set_callback(*this);
}

void MQTTClient::connect()
{
    std::lock_guard<std::mutex> lock(mutex_);
    if (connected_ || !g_running) // 添加g_running检查
        return;

    try
    {
        auto connOpts = mqtt::connect_options_builder()
                            .clean_session(config_.clean_session)
                            .automatic_reconnect(true)
                            .keep_alive_interval(std::chrono::seconds(config_.keep_alive))
                            .user_name(config_.username)
                            .password(config_.password)
                            .finalize();

        // 异步连接，不阻塞
        auto tok = client_->connect(connOpts); // 返回 token
    }
    catch (const mqtt::exception &e)
    {
        connected_ = false;
        LOG_ERROR("[MQTT] Connect failed: " + std::string(e.what()));
    }
}

void MQTTClient::disconnect()
{
    std::lock_guard<std::mutex> lock(mutex_);
    if (!connected_)
        return;

    try
    {
        auto disc_tok = client_->disconnect(); // 异步返回 token

        // 等待短时间，检查 g_running
        for (int i = 0; i < 10 && g_running; ++i)
        {
            if (disc_tok->is_complete())
                break;
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }

        if (!disc_tok->is_complete())
        {
            LOG_WARN("[MQTT] Disconnect timed out, forcing disconnection");
            client_->stop_consuming();
            client_.reset();
        }
    }
    catch (const mqtt::exception &e)
    {
        LOG_ERROR("[MQTT] Disconnect failed: " + std::string(e.what()));
    }

    connected_ = false;
    if (on_disconnect_)
        on_disconnect_();
}

void MQTTClient::force_disconnect()
{
    std::lock_guard<std::mutex> lock(mutex_);
    if (!connected_)
        return;

    try
    {
        // 直接停止客户端，不等待正常断开
        client_->stop_consuming();
        client_.reset();
    }
    catch (const mqtt::exception &e)
    {
        LOG_ERROR("[MQTT] Force disconnect failed: " + std::string(e.what()));
    }

    connected_ = false;
    if (on_disconnect_)
        on_disconnect_();
}

void MQTTClient::publish(const std::string &topic, const std::string &payload, int qos)
{
    if (qos == -1)
        qos = config_.qos;

    try
    {
        client_->publish(topic, payload.data(), payload.size(), qos, false); // 异步，不阻塞
        // 如果 topic 不包含 "heartbeat"，才打印日志
        if (topic.find("heartbeat") == std::string::npos)
        {
            LOG_INFO("[MQTT] Published message to topic: " + topic);
        }
    }
    catch (const mqtt::exception &e)
    {
        LOG_ERROR("[MQTT] Publish failed: " + std::string(e.what()));
        connected_ = false;
        if (on_disconnect_)
            on_disconnect_();
    }
}

void MQTTClient::subscribe(const std::string &topic, int qos)
{
    if (qos == -1)
        qos = config_.qos;

    try
    {
        client_->subscribe(topic, qos); // 异步，不阻塞
        LOG_INFO("[MQTT] Subscribed to topic: " + topic);
    }
    catch (const mqtt::exception &e)
    {
        LOG_ERROR("[MQTT] Subscribe failed: " + std::string(e.what()));
        connected_ = false;
        if (on_disconnect_)
            on_disconnect_();
    }
}

void MQTTClient::switchServer(const MQTTConfig &newConfig)
{
    std::lock_guard<std::mutex> lock(mutex_);
    disconnect();
    config_ = newConfig;
    initializeClient();
    connect();
}

bool MQTTClient::isConnected() const
{
    return connected_;
}

void MQTTClient::setConnectCallback(ConnectCallback cb) { on_connect_ = cb; }
void MQTTClient::setDisconnectCallback(DisconnectCallback cb) { on_disconnect_ = cb; }
void MQTTClient::setMessageCallback(MessageCallback cb) { on_message_ = cb; }

void MQTTClient::connection_lost(const std::string &cause)
{
    connected_ = false;
    LOG_WARN("[MQTT] Connection lost: " + cause);
    if (on_disconnect_)
        on_disconnect_();
}

void MQTTClient::connected(const std::string &cause)
{
    connected_ = true;
    LOG_INFO("[MQTT] Reconnected: " + cause);
    if (on_connect_)
        on_connect_();
}

void MQTTClient::message_arrived(mqtt::const_message_ptr msg)
{
    if (on_message_)
        on_message_(msg->get_topic(), msg->get_payload_str());
}