import asyncio
from pymodbus.datastore import ModbusSlaveContext, ModbusServerContext
from pymodbus.server import StartAsyncTcpServer
from pymodbus.constants import Endian
from pymodbus.payload import BinaryPayloadBuilder
from aiohttp import web
# ===== 配置 =====
SERVER_IP = "0.0.0.0"
PORTS = [5020, 5021, 5022]
REGISTER_SIZE = 100
slave_contexts = {}
modbus_server_tasks = []  # ← 保存 server task 引用
def create_initial_holding_registers(size=100):
    builder = BinaryPayloadBuilder(byteorder=Endian.BIG, wordorder=Endian.BIG)
    for i in range(size):
        builder.add_16bit_uint(i * 10 + 100)
    return builder.to_registers()
async def run_modbus_slave(port):
    print(f"[Modbus] Starting slave on port {port}...")
    # 创建单个从站的数据存储
    slave_store = ModbusSlaveContext(zero_mode=True)
    initial_regs = create_initial_holding_registers(REGISTER_SIZE)
    slave_store.setValues(3, 0, initial_regs)  # 功能码 3 (Holding Registers)
    # ✅ 关键：用 ModbusServerContext 包装，指定 unit ID（默认 0）
    server_context = ModbusServerContext(slaves=slave_store, single=True)
    # 保存到全局（供 HTTP API 使用）
    slave_contexts[port] = slave_store  # 注意：保存的是 slave_store，不是 server_context
    # 启动服务器
    server, task = await StartAsyncTcpServer(
        context=server_context,  # ← 传入包装后的 context
        address=(SERVER_IP, port)
    )
    # 保存 task 引用，防止被 GC 回收
    modbus_server_tasks.append(task)
    print(f"[Modbus] Slave on port {port} is running.")
# ===== HTTP API =====
async def handle_write_register(request):
    try:
        data = await request.json()
        port = int(data["port"])
        address = int(data["address"])
        value = int(data["value"])
        if port not in slave_contexts:
            return web.json_response({"error": f"Port {port} not found"}, status=404)
        if not (0 <= address < REGISTER_SIZE):
            return web.json_response({"error": f"Address must be 0-{REGISTER_SIZE - 1}"}, status=400)
        if not (0 <= value <= 65535):
            return web.json_response({"error": "Value must be 0-65535"}, status=400)
        slave_contexts[port].setValues(3, address, [value])
        return web.json_response({"success": True})
    except Exception as e:
        return web.json_response({"error": str(e)}, status=400)
async def handle_list_ports(request):
    return web.json_response({"ports": list(slave_contexts.keys())})
async def start_http_server():
    app = web.Application()
    app.router.add_post("/write", handle_write_register)
    app.router.add_get("/ports", handle_list_ports)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "0.0.0.0", 8080)
    await site.start()
    print("[HTTP] Server running on http://0.0.0.0:8080")
# ===== Main =====
async def main():
    print("🚀 Starting Modbus Slaves and HTTP API...")
    # 启动所有 Modbus 从站
    for port in PORTS:
        asyncio.create_task(run_modbus_slave(port))
    # 启动 HTTP 服务器
    await start_http_server()
    print("✅ All services started. Press Ctrl+C to stop.")
    try:
        # 永久等待（因为 Modbus server task 已在后台运行）
        await asyncio.Event().wait()
    except KeyboardInterrupt:
        print("\n🛑 Shutting down...")
if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n🛑 Exited.")