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

#include "VrError.h"
#include "VrLog.h"

#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();
}

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

    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 ERR_CODE(NET_ERR_CREAT_INIT);
    }

    // 创建数据映射 - 分配足够的空间
    // 线圈: 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 ERR_CODE(NET_ERR_CONFIG);
    }

    // 监听连接
    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 ERR_CODE(NET_ERR_CREAT_LISTEN);
    }

    m_isRunning.store(true);
    m_shouldStop.store(false);
    
    // 启动服务器线程
    m_serverThread = std::make_unique<std::thread>(&ModbusTCPServer::serverLoop, this);
    
    // 初始状态：服务器启动但没有客户端连接
    if (m_connectionStatusCallback) {
        m_connectionStatusCallback(false);
    }
    
    return SUCCESS;
}

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()) {
                LOG_WARNING("select error: %s\n", strerror(errno));
            }
            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()) {
            LOG_ERRO("Accept connection failed: %s\n", modbus_strerror(errno));
        }
        return;
    }
    
    // 创建新的modbus context用于此客户端
    modbus_t* clientCtx = modbus_new_tcp("0.0.0.0", m_port);
    if (clientCtx == nullptr) {
        LOG_ERRO("Create Modbus context for client failed\n");
#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;
    }
    
    LOG_INFO("New client connected: socket=%d\n", clientSocket);
    
    // 触发连接状态回调
    if (m_connectionStatusCallback) {
        m_connectionStatusCallback(true);
    }
}

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) {
            LOG_VERBOSE("Client disconnected: socket=%d\n", client->socket);
        } else {
            LOG_ERRO("Receive data error: %s\n", strerror(errno));
        }
        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()) {
        LOG_INFO("Remove client connection: socket=%d\n", socket);
        m_clients.erase(it);
        
        // 如果没有客户端连接了，触发断开回调
        if (m_clients.empty() && m_connectionStatusCallback) {
            m_connectionStatusCallback(false);
        }
    }
}

void ModbusTCPServer::processModbusRequest(std::shared_ptr<ClientConnection> client, const uint8_t* query, int queryLength)
{
    // 使用modbus标准异常码，0表示成功
    uint8_t exceptionCode = 0;
    
    // 解析ModbusTCP ADU的基本字段
    uint16_t transactionId = 0;
    uint16_t protocolId = 0;
    uint16_t length = 0;
    uint8_t unitId = 0;
    uint8_t function = 0;
    uint16_t address = 0;
    uint16_t value = 0;
    uint16_t quantity = 0;
    uint8_t byteCount = 0;
    std::vector<uint16_t> registerValues;
    std::vector<uint8_t> coilValues;
    
    // 基本长度检查
    if (queryLength < 8) {
        exceptionCode = MODBUS_EXCEPTION_ILLEGAL_DATA_VALUE;
        modbus_reply_exception(client->modbusCtx, query, exceptionCode);
        return;
    } 
    

    // 解析ADU头部
    transactionId = (query[0] << 8) | query[1];
    protocolId = (query[2] << 8) | query[3];
    length = (query[4] << 8) | query[5];
    unitId = query[6];
    function = query[7];
    
    // 验证协议ID
    if (protocolId != 0) {
        exceptionCode = MODBUS_EXCEPTION_ILLEGAL_DATA_VALUE;
    }
    // 验证长度字段
    else if (length != (queryLength - 6)) {
        exceptionCode = MODBUS_EXCEPTION_ILLEGAL_DATA_VALUE;
    }

    if(exceptionCode != 0) {
        modbus_reply_exception(client->modbusCtx, query, exceptionCode);
        return;
    }

    // 根据功能码验证数据格式并执行回调
    LOG_DEBUG("Modbus %02X - Trans:%d Unit:%d\n", function, transactionId, unitId);
    
    switch (function) {
        case MODBUS_FC_WRITE_SINGLE_COIL:
            {
                if (queryLength < 12) {
                    exceptionCode = MODBUS_EXCEPTION_ILLEGAL_DATA_VALUE;
                } else {
                    address = (query[8] << 8) | query[9];
                    value = (query[10] << 8) | query[11];
                    
                    if (m_writeCoilsCallback) {
                        uint8_t coilValue = (value == 0xFF00) ? 1 : 0;
                        IYModbusTCPServer::ErrorCode result = m_writeCoilsCallback(unitId, address, 1, &coilValue);
                        exceptionCode = (uint8_t)result;
                    }
                }
            }
            break;
            
        case MODBUS_FC_WRITE_MULTIPLE_COILS:
            {
                if (queryLength < 13) {
                    exceptionCode = MODBUS_EXCEPTION_ILLEGAL_DATA_VALUE;
                } else {
                    address = (query[8] << 8) | query[9];
                    quantity = (query[10] << 8) | query[11];
                    byteCount = query[12];
                    
                    if (queryLength < (13 + byteCount)) {
                        exceptionCode = MODBUS_EXCEPTION_ILLEGAL_DATA_VALUE;
                    } else {
                        coilValues.assign(&query[13], &query[13 + byteCount]);
                        
                        if (m_writeCoilsCallback) {
                            IYModbusTCPServer::ErrorCode result = m_writeCoilsCallback(unitId, address, quantity, coilValues.data());
                            exceptionCode = (uint8_t)result;
                        }
                    }
                }
            }
            break;

        case MODBUS_FC_WRITE_SINGLE_REGISTER:
            {
                if (queryLength < 12) {
                    exceptionCode = MODBUS_EXCEPTION_ILLEGAL_DATA_VALUE;
                } else {
                    address = (query[8] << 8) | query[9];
                    value = (query[10] << 8) | query[11];
                    
                    if (m_writeRegistersCallback) {
                        IYModbusTCPServer::ErrorCode result = m_writeRegistersCallback(unitId, address, 1, &value);
                        exceptionCode = (uint8_t)result;
                    }
                }
            }
            break;
            
        case MODBUS_FC_WRITE_MULTIPLE_REGISTERS:
            {
                if (queryLength < 13) {
                    exceptionCode = MODBUS_EXCEPTION_ILLEGAL_DATA_VALUE;
                } else {
                    address = (query[8] << 8) | query[9];
                    quantity = (query[10] << 8) | query[11];
                    byteCount = query[12];
                    
                    if (queryLength < (13 + byteCount) || byteCount != (quantity * 2)) {
                        exceptionCode = MODBUS_EXCEPTION_ILLEGAL_DATA_VALUE;
                    } else {
                        registerValues.reserve(quantity);
                        for (int i = 0; i < quantity; i++) {
                            uint16_t regValue = (query[13 + i*2] << 8) | query[13 + i*2 + 1];
                            registerValues.push_back(regValue);
                        }
                        
                        if (m_writeRegistersCallback) {
                            IYModbusTCPServer::ErrorCode result = m_writeRegistersCallback(unitId, address, quantity, registerValues.data());
                            exceptionCode = (uint8_t)result;
                        }
                    }
                }
            }
            break;
            
        default:
            // exceptionCode = MODBUS_EXCEPTION_ILLEGAL_FUNCTION;
            break;
    }
    
    LOG_DEBUG("exceptionCode: %d\n", exceptionCode);
    
    // 统一处理响应
    if (exceptionCode != 0) {
        modbus_reply_exception(client->modbusCtx, query, exceptionCode);
        return;
    }
    
    // 发送正常响应
    {
        std::lock_guard<std::mutex> lock(m_dataMutex);
        int responseRc = modbus_reply(client->modbusCtx, query, queryLength, m_mapping);
        
        if (responseRc == -1) {
            modbus_reply_exception(client->modbusCtx, query, MODBUS_EXCEPTION_SLAVE_OR_SERVER_FAILURE);
            return;
        }
    }
}

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;
}