WSCommModifyService/WebSocketHandler.cpp

#include "WebSocketHandler.h"
#include "WebSocketSession.h"
#include "rapidjson/document.h"
#include "rapidjson/writer.h"
#include "rapidjson/stringbuffer.h"
#include "Logger.h"
#include <algorithm>
#include <chrono>

WebSocketHandler::WebSocketHandler(bool use_dll) : use_dll_(use_dll), timeout_thread_running_(false)
{
    if (use_dll_)
    {
        comm_wrapper_ = std::make_shared<DllWrapper>();
    }
    else
    {
        comm_wrapper_ = std::make_shared<ActiveXWrapper>();
    }
}

WebSocketHandler::~WebSocketHandler()
{
    Cleanup();
}

bool WebSocketHandler::Initialize()
{
    if (!comm_wrapper_->Initialize())
    {
        LOG_ERROR("Failed to initialize communication wrapper");
        return false;
    }

    // 设置事件回调
    comm_wrapper_->SetEventCallback([this](const CommEventData &event_data)
                                    { OnCommEvent(event_data); });

    // 初始化监控
    if (!comm_wrapper_->InitMonitor())
    {
        LOG_ERROR("Failed to initialize communication monitor");
        return false;
    }

    // 启动超时检测线程
    //timeout_thread_running_ = true;
    //timeout_check_thread_ = std::thread(&WebSocketHandler::TimeoutCheckThread, this);

    LOG_INFO("WebSocketHandler initialized successfully");
    return true;
}

void WebSocketHandler::AddClient(std::shared_ptr<WebSocketSession> session)
{
    std::lock_guard<std::mutex> lock(clients_mutex_);
    connected_clients_.insert(session);
    LOG_INFO("Client connected, total clients: " + std::to_string(connected_clients_.size()));
}

void WebSocketHandler::RemoveClient(std::shared_ptr<WebSocketSession> session)
{
    std::lock_guard<std::mutex> lock(clients_mutex_);
    // 从连接客户端集合中移除
    connected_clients_.erase(session);

    // 关闭WebSocket连接
    session->close();

    LOG_INFO("Client disconnected, total clients: " + std::to_string(connected_clients_.size()));
}

void WebSocketHandler::HandleMessage(std::shared_ptr<WebSocketSession> session, const std::string &message)
{
    rapidjson::Document doc;
    if (doc.Parse(message.c_str()).HasParseError())
    {
        LOG_ERROR("Failed to parse JSON message: " + message);
        return;
    }

    if (!doc.HasMember("cmd") || !doc["cmd"].IsString())
    {
        LOG_ERROR("Invalid message format: missing 'cmd' field");
        return;
    }

    std::string cmd = doc["cmd"].GetString();

    if (cmd == "InitMonitor")
    {
        if (!doc.HasMember("port") || !doc["port"].IsInt())
        {
            LOG_ERROR("Invalid InitMonitor message: missing 'port' field");
            return;
        }
        int port = doc["port"].GetInt();
        HandleInitMonitor(session, port);
    }
    else if (cmd == "StopPort")
    {
        if (!doc.HasMember("port") || !doc["port"].IsInt())
        {
            LOG_ERROR("Invalid StopPort message: missing 'port' field");
            return;
        }
        int port = doc["port"].GetInt();
        HandleStopPort(session, port);
    }
    else if (cmd == "ReadData")
    {
        if (!doc.HasMember("port") || !doc["port"].IsInt() ||
            !doc.HasMember("data") || !doc["data"].IsString())
        {
            LOG_ERROR("Invalid ReadData message: missing required fields");
            return;
        }
        int port = doc["port"].GetInt();
        std::string data = doc["data"].GetString();
        HandleReadData(session, port, data);
    }
    else if (cmd == "WriteData")
    {
        if (!doc.HasMember("port") || !doc["port"].IsInt() ||
            !doc.HasMember("data") || !doc["data"].IsString())
        {
            LOG_ERROR("Invalid WriteData message: missing required fields");
            return;
        }
        int port = doc["port"].GetInt();
        std::string data = doc["data"].GetString();
        HandleWriteData(session, port, data);
    }
    else
    {
        LOG_ERROR("Unknown command: " + cmd);
    }
}

void WebSocketHandler::HandleInitMonitor(std::shared_ptr<WebSocketSession> session, int port)
{
    bool result = comm_wrapper_->MonitorPort(port);
    SendResponse(session, "InitMonitor", result, port);

    if (!result)
    {
        int error_code = comm_wrapper_->GetLastError();
        std::string error_msg = comm_wrapper_->GetErrorMessage(error_code);
        LOG_ERROR("InitMonitor failed for port " + std::to_string(port) + ": " + error_msg);
    }
}

void WebSocketHandler::HandleStopPort(std::shared_ptr<WebSocketSession> session, int port)
{
    bool result = comm_wrapper_->StopPort(port);
    SendResponse(session, "StopPort", result, port);
}

void WebSocketHandler::HandleReadData(std::shared_ptr<WebSocketSession> session, int port, const std::string &data)
{
    // 移除端口活动时间
    RemovePortActivity(port);
    
    bool result = comm_wrapper_->ReadData(port, data);
    SendResponse(session, "ReadData", result, port);
}

void WebSocketHandler::HandleWriteData(std::shared_ptr<WebSocketSession> session, int port, const std::string &data)
{
	// 移除端口活动时间
	RemovePortActivity(port);
    
    bool result = comm_wrapper_->WriteData(port, data);
    SendResponse(session, "WriteData", result, port);
}

void WebSocketHandler::SendResponse(std::shared_ptr<WebSocketSession> session, const std::string &cmd, bool result, int port)
{
    rapidjson::StringBuffer buffer;
    rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);

    writer.StartObject();
    writer.Key("cmd");
    writer.String(cmd.c_str());
    writer.Key("result");
    writer.Bool(result);
    if (port > 0)
    {
        writer.Key("port");
        writer.Int(port);
    }
    writer.EndObject();

    std::string response = buffer.GetString();
    auto message = std::make_shared<std::string>(response);
    session->safe_send(message);

    LOG_DEBUG("Sent response: " + response);
}

void WebSocketHandler::BroadcastEvent(const std::string &cmd, int port, const std::string &data)
{
    rapidjson::StringBuffer buffer;
    rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);

    writer.StartObject();
    writer.Key("cmd");
    writer.String(cmd.c_str());
    writer.Key("port");
    writer.Int(port);
    if (!data.empty())
    {
        writer.Key("data");
        writer.String(data.c_str());
    }
    writer.EndObject();

    std::string event_message = buffer.GetString();
    auto message = std::make_shared<std::string>(event_message);

    std::lock_guard<std::mutex> lock(clients_mutex_);
    for (auto &client : connected_clients_)
    {
        client->safe_send(message);
    }

    LOG_DEBUG(std::to_string(connected_clients_.size()) + " Broadcasted event: " + event_message);
}

void WebSocketHandler::Cleanup()
{
    // 停止超时检测线程
    if (timeout_thread_running_)
    {
        timeout_thread_running_ = false;
        if (timeout_check_thread_.joinable())
        {
            timeout_check_thread_.join();
        }
    }

    if (comm_wrapper_)
    {
        comm_wrapper_->Cleanup();
    }

    std::lock_guard<std::mutex> lock(clients_mutex_);
    connected_clients_.clear();

    LOG_INFO("WebSocketHandler cleaned up");
}

void WebSocketHandler::OnCommEvent(const CommEventData &event_data)
{
    std::string cmd;
    std::string data_str;

    switch (event_data.event_type)
    {
    case CommEventType::OP_OPEN:
        cmd = "OP_OPEN";
        break;
    case CommEventType::OP_CLOSE:
        cmd = "OP_CLOSE";
        break;
    case CommEventType::OP_READ:
        cmd = "OP_READ";
        // 将字节数组转换为十六进制字符串
        for (size_t i = 0; i < event_data.data.size(); i++)
        {
            if (i > 0)
                data_str += " ";
            char hex[3];
            sprintf_s(hex, "%02X", event_data.data[i]);
            data_str += hex;
        }
        break;
    case CommEventType::OP_WRITE:
        cmd = "OP_WRITE";
        // 将字节数组转换为十六进制字符串
        for (size_t i = 0; i < event_data.data.size(); i++)
        {
            if (i > 0)
                data_str += " ";
            char hex[3];
            sprintf_s(hex, "%02X", event_data.data[i]);
            data_str += hex;
        }
        break;
    default:
        return; // 忽略未知事件
	}

	BroadcastEvent(cmd, event_data.port, data_str);

	switch (event_data.event_type)
	{
	case CommEventType::OP_READ:
	case CommEventType::OP_WRITE:
		// 更新端口活动时间
		UpdatePortActivity(event_data.port);
		break;
	default:
        break;
	}
}

void WebSocketHandler::UpdatePortActivity(int port)
{
    std::lock_guard<std::mutex> lock(activity_mutex_);
    last_activity_time_[port] = std::chrono::steady_clock::now();
}

void WebSocketHandler::RemovePortActivity(int port)
{
	std::lock_guard<std::mutex> lock(activity_mutex_);
    last_activity_time_.erase(port);
}

void WebSocketHandler::TimeoutCheckThread()
{
    //while (timeout_thread_running_)
    //{
    //    // 等待一段时间再检查
    //    std::this_thread::sleep_for(std::chrono::seconds(5));

    //    if (!timeout_thread_running_)
    //        break;

    //    // 获取当前时间
    //    auto now = std::chrono::steady_clock::now();
    //    
    //    // 创建一个副本以避免在遍历时锁定
    //    std::map<int, std::chrono::steady_clock::time_point> activity_copy;
    //    {
    //        std::lock_guard<std::mutex> lock(activity_mutex_);
    //        activity_copy = last_activity_time_;
    //    }

    //    // 检查每个端口的活动时间
    //    for (const auto &entry : activity_copy)
    //    {
    //        int port = entry.first;
    //        auto last_activity = entry.second;
    //        
    //        // 计算自上次活动以来的时间
    //        auto duration = std::chrono::duration_cast<std::chrono::seconds>(now - last_activity).count();
    //        
    //        // 如果超过30秒没有活动，则停止监听该端口
    //        if (duration >= TIMEOUT_SECONDS)
    //        {
    //            LOG_INFO("Port " + std::to_string(port) + " timeout, stopping monitoring");
    //            
    //            // 停止监听端口
    //            if (comm_wrapper_)
    //            {
    //                comm_wrapper_->StopPort(port);
    //            }
    //            
    //            // 从活动时间映射中移除该端口
    //            {
    //                std::lock_guard<std::mutex> lock(activity_mutex_);
    //                last_activity_time_.erase(port);
    //            }
    //        }
    //    }
    //}
}