#include "ModbusTCPServer.h"
#include <iostream>
#include <cstring>
#include <stdexcept>

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")
#else
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/select.h>
#include <errno.h>
#endif

ModbusTCPServer::ModbusTCPServer()
    : m_modbusCtx(nullptr)
    , m_mapping(nullptr)
    , m_isRunning(false)
    , m_shouldStop(false)
    , m_serverSocket(-1)
    , m_port(502)
    , m_maxConnections(10)
{
}

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

bool ModbusTCPServer::start(int port, int maxConnections)
{
    if (m_isRunning.load()) {
        setLastError("服务器已在运行");
        return true;
    }

    m_port = port;
    m_maxConnections = maxConnections;

    // 创建modbus TCP context
    m_modbusCtx = modbus_new_tcp("0.0.0.0", port);
    if (m_modbusCtx == nullptr) {
        setLastError("创建Modbus TCP上下文失败");
        return false;
    }

    // 创建数据映射 - 分配足够的空间
    // 线圈: 0-9999, 离散输入: 0-9999, 保持寄存器: 0-9999, 输入寄存器: 0-9999
    m_mapping = modbus_mapping_new(10000, 10000, 10000, 10000);
    if (m_mapping == nullptr) {
        setLastError("创建数据映射失败");
        modbus_free(m_modbusCtx);
        m_modbusCtx = nullptr;
        return false;
    }

    // 监听连接
    m_serverSocket = modbus_tcp_listen(m_modbusCtx, maxConnections);
    if (m_serverSocket == -1) {
        setLastError(std::string("监听失败: ") + modbus_strerror(errno));
        modbus_mapping_free(m_mapping);
        m_mapping = nullptr;
        modbus_free(m_modbusCtx);
        m_modbusCtx = nullptr;
        return false;
    }

    m_isRunning.store(true);
    m_shouldStop.store(false);
    
    // 启动服务器线程
    m_serverThread = std::make_unique<std::thread>(&ModbusTCPServer::serverLoop, this);
    
    return true;
}

void ModbusTCPServer::stop()
{
    if (!m_isRunning.load()) {
        return;
    }

    m_shouldStop.store(true);
    
    // 关闭服务器socket以中断accept调用
    if (m_serverSocket != -1) {
#ifdef _WIN32
        closesocket(m_serverSocket);
#else
        close(m_serverSocket);
#endif
        m_serverSocket = -1;
    }

    // 等待服务器线程退出
    if (m_serverThread && m_serverThread->joinable()) {
        m_serverThread->join();
    }

    // 清理所有客户端连接
    {
        std::lock_guard<std::mutex> lock(m_clientsMutex);
        m_clients.clear();
    }

    // 清理资源
    if (m_mapping) {
        modbus_mapping_free(m_mapping);
        m_mapping = nullptr;
    }

    if (m_modbusCtx) {
        modbus_free(m_modbusCtx);
        m_modbusCtx = nullptr;
    }

    m_isRunning.store(false);
}

void ModbusTCPServer::serverLoop()
{
    fd_set readfds;
    int maxfd = m_serverSocket;
    
    while (!m_shouldStop.load()) {
        FD_ZERO(&readfds);
        FD_SET(m_serverSocket, &readfds);
        
        maxfd = m_serverSocket;
        
        // 添加所有客户端socket到监听集合
        {
            std::lock_guard<std::mutex> lock(m_clientsMutex);
            for (const auto& pair : m_clients) {
                int clientSocket = pair.first;
                FD_SET(clientSocket, &readfds);
                if (clientSocket > maxfd) {
                    maxfd = clientSocket;
                }
            }
        }
        
        // 设置超时时间（1秒），避免阻塞太久
        struct timeval timeout;
        timeout.tv_sec = 1;
        timeout.tv_usec = 0;
        
        int selectResult = select(maxfd + 1, &readfds, nullptr, nullptr, &timeout);
        
        if (selectResult == -1) {
            if (!m_shouldStop.load()) {
                std::cerr << "select错误: " << strerror(errno) << std::endl;
            }
            break;
        }
        
        if (selectResult == 0) {
            // 超时，继续循环
            continue;
        }
        
        // 检查服务器socket是否有新连接
        if (FD_ISSET(m_serverSocket, &readfds)) {
            handleNewConnection();
        }
        
        // 检查客户端socket是否有数据
        std::vector<int> socketsToRemove;
        {
            std::lock_guard<std::mutex> lock(m_clientsMutex);
            for (const auto& pair : m_clients) {
                int clientSocket = pair.first;
                if (FD_ISSET(clientSocket, &readfds)) {
                    try {
                        handleClientData(pair.second);
                    } catch (...) {
                        // 客户端处理出错，标记删除
                        socketsToRemove.push_back(clientSocket);
                    }
                }
            }
        }
        
        // 删除有问题的客户端连接
        for (int socket : socketsToRemove) {
            removeClient(socket);
        }
    }
}

void ModbusTCPServer::handleNewConnection()
{
    int clientSocket = modbus_tcp_accept(m_modbusCtx, &m_serverSocket);
    if (clientSocket == -1) {
        if (!m_shouldStop.load()) {
            std::cerr << "接受连接失败: " << modbus_strerror(errno) << std::endl;
        }
        return;
    }
    
    // 创建新的modbus context用于此客户端
    modbus_t* clientCtx = modbus_new_tcp("0.0.0.0", m_port);
    if (clientCtx == nullptr) {
        std::cerr << "为客户端创建Modbus上下文失败" << std::endl;
#ifdef _WIN32
        closesocket(clientSocket);
#else
        close(clientSocket);
#endif
        return;
    }
    
    // 设置socket为非阻塞模式
    modbus_set_socket(clientCtx, clientSocket);
    
    // 创建客户端连接对象
    auto client = std::make_shared<ClientConnection>(clientSocket, clientCtx);
    
    // 添加到客户端列表
    {
        std::lock_guard<std::mutex> lock(m_clientsMutex);
        m_clients[clientSocket] = client;
    }
    
    std::cout << "新客户端连接: socket=" << clientSocket << std::endl;
}

void ModbusTCPServer::handleClientData(std::shared_ptr<ClientConnection> client)
{
    uint8_t buffer[MODBUS_TCP_MAX_ADU_LENGTH];
    
    // 接收数据
    int bytesReceived = recv(client->socket, reinterpret_cast<char*>(buffer), sizeof(buffer), 0);
    
    if (bytesReceived <= 0) {
        // 连接断开或错误
        if (bytesReceived == 0) {
            std::cout << "客户端断开连接: socket=" << client->socket << std::endl;
        } else {
            std::cerr << "接收数据错误: " << strerror(errno) << std::endl;
        }
        throw std::runtime_error("客户端连接错误");
    }
    
    // 处理Modbus请求
    processModbusRequest(client, buffer, bytesReceived);
}

void ModbusTCPServer::removeClient(int socket)
{
    std::lock_guard<std::mutex> lock(m_clientsMutex);
    auto it = m_clients.find(socket);
    if (it != m_clients.end()) {
        std::cout << "移除客户端连接: socket=" << socket << std::endl;
        m_clients.erase(it);
    }
}

void ModbusTCPServer::processModbusRequest(std::shared_ptr<ClientConnection> client, const uint8_t* query, int queryLength)
{
    // 处理请求并发送响应
    std::lock_guard<std::mutex> lock(m_dataMutex);
    int responseRc = modbus_reply(client->modbusCtx, query, queryLength, m_mapping);
    
    if (responseRc == -1) {
        std::cerr << "发送响应失败: " << modbus_strerror(errno) << std::endl;
        throw std::runtime_error("发送响应失败");
    }

    // 检查是否是写操作，如果是则调用回调
    if (queryLength >= 8) { // 确保有足够的字节
        uint8_t function = query[7]; // 功能码在TCP ADU的第8字节
        uint16_t address = (query[8] << 8) | query[9];
        uint16_t quantity = (query[10] << 8) | query[11];
        uint8_t unitId = query[6];

        switch (function) {
            case MODBUS_FC_WRITE_SINGLE_COIL:
            case MODBUS_FC_WRITE_MULTIPLE_COILS:
                if (m_writeCoilsCallback) {
                    const uint8_t* values = &query[12]; // 数据从第13字节开始
                    ErrorCode result = m_writeCoilsCallback(unitId, address, quantity, values);
                    if (result != ErrorCode::SUCCESS) {
                        // 发送异常响应
                        modbus_reply_exception(client->modbusCtx, query, static_cast<unsigned int>(result));
                    }
                }
                break;

            case MODBUS_FC_WRITE_SINGLE_REGISTER:
            case MODBUS_FC_WRITE_MULTIPLE_REGISTERS:
                if (m_writeRegistersCallback) {
                    const uint16_t* values = reinterpret_cast<const uint16_t*>(&query[12]);
                    ErrorCode result = m_writeRegistersCallback(unitId, address, quantity, values);
                    if (result != ErrorCode::SUCCESS) {
                        // 发送异常响应
                        modbus_reply_exception(client->modbusCtx, query, static_cast<unsigned int>(result));
                    }
                }
                break;
        }
    }
}

void ModbusTCPServer::updateCoil(uint16_t address, bool value)
{
    std::lock_guard<std::mutex> lock(m_dataMutex);
    if (m_mapping && address < m_mapping->nb_bits) {
        m_mapping->tab_bits[address] = value ? 1 : 0;
    }
}

void ModbusTCPServer::updateCoils(uint16_t startAddress, const std::vector<bool>& values)
{
    std::lock_guard<std::mutex> lock(m_dataMutex);
    if (!m_mapping) return;
    
    for (size_t i = 0; i < values.size() && (startAddress + i) < m_mapping->nb_bits; ++i) {
        m_mapping->tab_bits[startAddress + i] = values[i] ? 1 : 0;
    }
}

void ModbusTCPServer::updateDiscreteInput(uint16_t address, bool value)
{
    std::lock_guard<std::mutex> lock(m_dataMutex);
    if (m_mapping && address < m_mapping->nb_input_bits) {
        m_mapping->tab_input_bits[address] = value ? 1 : 0;
    }
}

void ModbusTCPServer::updateDiscreteInputs(uint16_t startAddress, const std::vector<bool>& values)
{
    std::lock_guard<std::mutex> lock(m_dataMutex);
    if (!m_mapping) return;
    
    for (size_t i = 0; i < values.size() && (startAddress + i) < m_mapping->nb_input_bits; ++i) {
        m_mapping->tab_input_bits[startAddress + i] = values[i] ? 1 : 0;
    }
}

void ModbusTCPServer::updateHoldingRegister(uint16_t address, uint16_t value)
{
    std::lock_guard<std::mutex> lock(m_dataMutex);
    if (m_mapping && address < m_mapping->nb_registers) {
        m_mapping->tab_registers[address] = value;
    }
}

void ModbusTCPServer::updateHoldingRegisters(uint16_t startAddress, const std::vector<uint16_t>& values)
{
    std::lock_guard<std::mutex> lock(m_dataMutex);
    if (!m_mapping) return;
    
    for (size_t i = 0; i < values.size() && (startAddress + i) < m_mapping->nb_registers; ++i) {
        m_mapping->tab_registers[startAddress + i] = values[i];
    }
}

void ModbusTCPServer::updateInputRegister(uint16_t address, uint16_t value)
{
    std::lock_guard<std::mutex> lock(m_dataMutex);
    if (m_mapping && address < m_mapping->nb_input_registers) {
        m_mapping->tab_input_registers[address] = value;
    }
}

void ModbusTCPServer::updateInputRegisters(uint16_t startAddress, const std::vector<uint16_t>& values)
{
    std::lock_guard<std::mutex> lock(m_dataMutex);
    if (!m_mapping) return;
    
    for (size_t i = 0; i < values.size() && (startAddress + i) < m_mapping->nb_input_registers; ++i) {
        m_mapping->tab_input_registers[startAddress + i] = values[i];
    }
}

std::string ModbusTCPServer::getLastError() const
{
    std::lock_guard<std::mutex> lock(m_errorMutex);
    return m_lastError;
}

void ModbusTCPServer::setLastError(const std::string& error)
{
    std::lock_guard<std::mutex> lock(m_errorMutex);
    m_lastError = error;
} 

bool IYModbusTCPServer::CreateInstance(IYModbusTCPServer** ppModbusTCPServer)
{
    if (ppModbusTCPServer == nullptr) {
        return false;
    }

    *ppModbusTCPServer = new ModbusTCPServer();
    return true;
}