import asyncio
from pymodbus.datastore import ModbusSlaveContext, ModbusServerContext
from pymodbus.server import StartAsyncTcpServer
from aiohttp import web
import logging

# 禁用 pymodbus 冗余日志（可选）
logging.getLogger("pymodbus").setLevel(logging.WARNING)

# ===== 配置区 =====
SERVER_IP = "0.0.0.0"
PORTS = [5020, 5021, 5022]          # 每个端口一个从站
REGISTER_SIZE = 100                 # 寄存器数量 (0 ～ 99)
COMMAND_REGISTERS = {2, 3}          # 哪些地址是“命令寄存器”，读取后自动清零

# 场景定义：每个场景是一个字典 {address: value}
SCENARIOS = {
    "reset": {},                    # 全零
    "call_car": {2: 1},             # 地址2=1 表示叫车
    "call_material": {2: 2},        # 地址2=2 表示叫料
    "custom_test": {2: 1, 3: 5, 10: 100},
}

# ===== 自定义 Modbus Slave Context（核心：实现读取后清零）=====
class AutoResetModbusSlaveContext(ModbusSlaveContext):
    def __init__(self, command_registers, size=100, **kwargs):
        super().__init__(**kwargs)
        self.command_registers = command_registers
        self.register_size = size
        # 初始化所有寄存器为 0
        self.setValues(3, 0, [0] * size)

    def getValues(self, fc_as_hex, address, count=1):
        """重写读取方法：若读取的是命令寄存器，则读完后清零"""
        values = super().getValues(fc_as_hex, address, count)
        
        # 仅处理 Holding Registers (功能码 3)
        if fc_as_hex == 3:
            for i in range(count):
                reg_addr = address + i
                if reg_addr in self.command_registers and reg_addr < self.register_size:
                    # 读取后立即清零
                    self.setValues(3, reg_addr, [0])
        return values

# ===== 全局状态 =====
slave_states = {}  # port -> {context, scenarios, current_profile}
modbus_server_tasks = []

# ===== 启动 Modbus 从站 =====
async def run_modbus_slave(port):
    print(f"[Modbus] Starting slave on port {port}...")

    # 创建自定义上下文（支持自动清零）
    store = AutoResetModbusSlaveContext(
        command_registers=COMMAND_REGISTERS,
        size=REGISTER_SIZE,
        zero_mode=True
    )

    # 构建完整寄存器列表（用于场景切换）
    def build_full_registers(scenario_dict):
        regs = [0] * REGISTER_SIZE
        for addr, val in scenario_dict.items():
            if 0 <= addr < REGISTER_SIZE and 0 <= val <= 65535:
                regs[addr] = val
        return regs

    # 预计算所有场景的完整寄存器值
    scenario_data = {}
    for name, sparse in SCENARIOS.items():
        scenario_data[name] = build_full_registers(sparse)

    # 初始状态：全零
    store.setValues(3, 0, scenario_data["reset"])

    slave_states[port] = {
        "store": store,
        "scenarios": scenario_data,
        "current_profile": "reset"
    }

    # 启动服务器（pymodbus 3.6+）
    server_context = ModbusServerContext(slaves=store, single=True)
    server, task = await StartAsyncTcpServer(
        context=server_context,
        address=(SERVER_IP, port)
    )
    modbus_server_tasks.append(task)
    print(f"[Modbus] Port {port} ready. Command registers: {sorted(COMMAND_REGISTERS)}")

# ===== HTTP API =====
async def handle_list_ports(request):
    return web.json_response({"ports": list(slave_states.keys())})


async def handle_list_profiles(request):
    port = int(request.query.get("port", PORTS[0]))
    if port not in slave_states:
        return web.json_response({"error": "Port not found"}, status=404)
    profiles = list(SCENARIOS.keys())
    current = slave_states[port]["current_profile"]
    return web.json_response({
        "port": port,
        "profiles": profiles,
        "current_profile": current
    })


async def handle_switch_profile(request):
    try:
        port = int(request.query.get("port"))
        profile = request.query.get("profile")

        if port not in slave_states:
            return web.json_response({"error": "Port not found"}, status=404)
        if profile not in SCENARIOS:
            return web.json_response({"error": f"Profile '{profile}' not defined"}, status=400)

        state = slave_states[port]
        full_regs = state["scenarios"][profile]
        state["store"].setValues(3, 0, full_regs)
        state["current_profile"] = profile

        return web.json_response({
            "success": True,
            "port": port,
            "profile": profile,
            "message": f"Switched to profile '{profile}'"
        })
    except Exception as e:
        return web.json_response({"error": str(e)}, status=400)


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_states:
            return web.json_response({"error": "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_states[port]["store"].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 start_http_server():
    app = web.Application()
    app.router.add_get("/ports", handle_list_ports)
    app.router.add_get("/profiles", handle_list_profiles)
    app.router.add_post("/switch_profile", handle_switch_profile)
    app.router.add_post("/write", handle_write_register)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "0.0.0.0", 8080)
    await site.start()
    print("[HTTP] API Server running on http://0.0.0.0:8080")
    print("  GET  /ports")
    print("  GET  /profiles?port=5020")
    print("  POST /switch_profile?port=5020&profile=call_car")
    print("  POST /write {\"port\":5020, \"address\":2, \"value\":1}")


# ===== Main =====
async def main():
    print("🚀 Starting Multi-Slave Modbus Tester (pymodbus 3.6+)...")
    print(f"Command registers (auto-reset on read): {sorted(COMMAND_REGISTERS)}")
    
    for port in PORTS:
        asyncio.create_task(run_modbus_slave(port))
    
    await start_http_server()
    
    print("✅ Ready for CECC testing. Press Ctrl+C to stop.")
    try:
        await asyncio.Event().wait()
    except KeyboardInterrupt:
        print("\n🛑 Shutting down...")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n🛑 Exited.")