"""
Test IPC Bridge v1.1 giữa Python Agent và C++ Core Engine
- Telemetry: Seqlock double-buffer (2 slot ping-pong, chống torn-read)
- Command : Hàng đợi vòng SPSC 16 slot (Agent produce, Core consume)
- Heartbeat: Agent cập nhật nhịp tim để Core watchdog giám sát
Cách chạy: khởi động bin/AutoPOE2_Core.exe trước, rồi chạy script này.
"""

import ctypes
import mmap
import sys
import time

if hasattr(sys.stdout, 'reconfigure'):
    sys.stdout.reconfigure(encoding='utf-8')

SHARED_MEM_NAME = "Local\\POE2_Auto_SharedMem_v1"
SHARED_MEM_SIZE = 4 * 1024 * 1024

# ---- Bố cục bộ nhớ v1.1 (khớp protocol.hpp) ----
SEQ_OFFSET = 0x10                  # telemetrySequence (header + 16)
CMD_WRITE_OFFSET = 0x18            # commandSequence   (header + 24)
CMD_READ_OFFSET = 0x20             # commandReadIndex  (header + 32)
AGENT_HEARTBEAT_OFFSET = 0x30      # lastAgentHeartbeat (header + 48)
TELEMETRY_SLOT0_OFFSET = 0x1000
TELEMETRY_SLOT_STRIDE = 0x10000
PLAYER_OFFSET_IN_PKT = 16          # snapshotId(8) + timestamp(8)
COMMAND_QUEUE_OFFSET = 0x300000
COMMAND_QUEUE_CAPACITY = 16
COMMAND_SLOT_STRIDE = 64

SHARED_MEM_MAGIC = 0x504F45324155544F

# OpCode khớp MacroOpCode trong protocol.hpp
OP_IDLE = 0
OP_NAVIGATE_TO_COORD = 1
OP_ENGAGE_TARGET = 2
OP_PICKUP_ITEM = 3
OP_CAST_EMERGENCY_PORTAL = 4
OP_EXECUTE_STASH_ROUTINE = 5
OP_EXECUTE_VENDOR_ROUTINE = 6
OP_TRIGGER_WORLD_MAP_TRAVEL = 96
OP_TRIGGER_AUTO_LOGIN = 97
OP_RECALIBRATE_XYZ = 98
OP_PANIC_STOP = 99


class SharedMemoryHeader(ctypes.Structure):
    _pack_ = 1
    _fields_ = [
        ("magic", ctypes.c_uint64),
        ("protocolVersion", ctypes.c_uint32),
        ("isRunning", ctypes.c_uint32),
        ("telemetrySequence", ctypes.c_uint64),
        ("commandSequence", ctypes.c_uint64),
        ("commandReadIndex", ctypes.c_uint64),
        ("lastCoreHeartbeat", ctypes.c_uint64),
        ("lastAgentHeartbeat", ctypes.c_uint64),
        ("reserved", ctypes.c_uint64),
    ]


class PlayerTelemetryData(ctypes.Structure):
    _pack_ = 1
    _fields_ = [
        ("playerAddr", ctypes.c_uint64),
        ("posX", ctypes.c_float),
        ("posY", ctypes.c_float),
        ("posZ", ctypes.c_float),
        ("yaw", ctypes.c_float),
        ("pitch", ctypes.c_float),
        ("currentHP", ctypes.c_uint32),
        ("maxHP", ctypes.c_uint32),
        ("currentMana", ctypes.c_uint32),
        ("maxMana", ctypes.c_uint32),
        ("currentES", ctypes.c_uint32),
        ("maxES", ctypes.c_uint32),
        ("currentWard", ctypes.c_uint32),
        ("maxWard", ctypes.c_uint32),
        ("currentSpirit", ctypes.c_uint32),
        ("maxSpirit", ctypes.c_uint32),
        ("activeWeaponSet", ctypes.c_uint32),
        ("movementFlags", ctypes.c_uint32),
        ("debuffMask", ctypes.c_uint64),
    ]


class EntityTelemetryData(ctypes.Structure):
    _pack_ = 1
    _fields_ = [
        ("id", ctypes.c_uint32),
        ("type", ctypes.c_uint16),
        ("rarity", ctypes.c_uint16),
        ("posX", ctypes.c_float),
        ("posY", ctypes.c_float),
        ("posZ", ctypes.c_float),
        ("distanceToPlayer", ctypes.c_float),
        ("currentHP", ctypes.c_uint32),
        ("maxHP", ctypes.c_uint32),
        ("staggerProgress", ctypes.c_uint16),
        ("currentAnimationId", ctypes.c_uint16),
        ("extraFlags", ctypes.c_uint32),
        ("name", ctypes.c_char * 48),
    ]


class AreaTelemetryData(ctypes.Structure):
    _pack_ = 1
    _fields_ = [
        ("areaLoadTimestampMs", ctypes.c_uint64),
        ("areaLevel", ctypes.c_uint32),
        ("areaSeed", ctypes.c_uint32),
        ("areaCode", ctypes.c_char * 32),
        ("areaName", ctypes.c_char * 48),
    ]


class TelemetryPacket(ctypes.Structure):
    _pack_ = 1
    _fields_ = [
        ("snapshotId", ctypes.c_uint64),
        ("timestamp", ctypes.c_uint64),
        ("player", PlayerTelemetryData),
        ("entityCount", ctypes.c_uint32),
        ("entities", EntityTelemetryData * 256),
        ("terrainWalkability", ctypes.c_uint8 * (128 * 128)),
        ("area", AreaTelemetryData),
    ]


class CommandPacket(ctypes.Structure):
    _pack_ = 1
    _fields_ = [
        ("commandId", ctypes.c_uint64),
        ("opCode", ctypes.c_uint32),
        ("targetX", ctypes.c_float),
        ("targetY", ctypes.c_float),
        ("targetZ", ctypes.c_float),
        ("targetEntityId", ctypes.c_uint32),
        ("priority", ctypes.c_uint32),
        ("timeoutMs", ctypes.c_uint32),
    ]


# ==========================================================
# Khóa kiểm chứng hợp đồng nhị phân ngay khi nạp module
# ==========================================================
assert ctypes.sizeof(SharedMemoryHeader) == 64, f"SharedMemoryHeader != 64B ({ctypes.sizeof(SharedMemoryHeader)})"
assert ctypes.sizeof(PlayerTelemetryData) == 84, f"PlayerTelemetryData != 84B ({ctypes.sizeof(PlayerTelemetryData)})"
assert ctypes.sizeof(EntityTelemetryData) == 88, f"EntityTelemetryData != 88B ({ctypes.sizeof(EntityTelemetryData)})"
assert ctypes.sizeof(AreaTelemetryData) == 96, f"AreaTelemetryData != 96B ({ctypes.sizeof(AreaTelemetryData)})"
assert ctypes.sizeof(CommandPacket) == 36, f"CommandPacket != 36B ({ctypes.sizeof(CommandPacket)})"
assert ctypes.sizeof(TelemetryPacket) == 39112, f"TelemetryPacket != 39112B ({ctypes.sizeof(TelemetryPacket)})"

assert SharedMemoryHeader.magic.offset == 0
assert SharedMemoryHeader.protocolVersion.offset == 8
assert SharedMemoryHeader.isRunning.offset == 12
assert SharedMemoryHeader.telemetrySequence.offset == 16
assert SharedMemoryHeader.commandSequence.offset == 24
assert SharedMemoryHeader.commandReadIndex.offset == 32
assert SharedMemoryHeader.lastCoreHeartbeat.offset == 40
assert SharedMemoryHeader.lastAgentHeartbeat.offset == 48
assert SharedMemoryHeader.reserved.offset == 56


def read_u64(shm, offset: int) -> int:
    shm.seek(offset)
    return int.from_bytes(shm.read(8), 'little')


def seqlock_read(shm, retries: int = 16):
    """Đọc snapshot telemetry chống rách theo giao thức seqlock.
    Trả về (sequence, PlayerTelemetryData) hoặc None."""
    for _ in range(retries):
        s1 = read_u64(shm, SEQ_OFFSET)
        if s1 == 0:
            return None  # Chưa có snapshot nào được publish
        slot = s1 & 1
        shm.seek(TELEMETRY_SLOT0_OFFSET + slot * TELEMETRY_SLOT_STRIDE + PLAYER_OFFSET_IN_PKT)
        player_bytes = shm.read(ctypes.sizeof(PlayerTelemetryData))
        s2 = read_u64(shm, SEQ_OFFSET)
        if s1 == s2:  # Writer không publish lần nào trong lúc đọc -> dữ liệu nguyên vẹn
            return s1, PlayerTelemetryData.from_buffer_copy(player_bytes)
    return None  # Quá số lần retry


def publish_command(shm, op_code: int, x: float = 0.0, y: float = 0.0,
                    z: float = 0.0, entity_id: int = 0):
    """Đẩy lệnh vào SPSC Command Ring (producer phía Agent)."""
    w = read_u64(shm, CMD_WRITE_OFFSET)
    r = read_u64(shm, CMD_READ_OFFSET)
    if w - r >= COMMAND_QUEUE_CAPACITY:
        print("[Command] Hàng đợi đầy - bỏ qua lần gửi này")
        return None

    cmd = CommandPacket()
    cmd.commandId = w + 1
    cmd.opCode = op_code
    cmd.targetX = x
    cmd.targetY = y
    cmd.targetZ = z
    cmd.targetEntityId = entity_id
    cmd.priority = 1
    cmd.timeoutMs = 3000

    shm.seek(COMMAND_QUEUE_OFFSET + (w % COMMAND_QUEUE_CAPACITY) * COMMAND_SLOT_STRIDE)
    shm.write(bytes(cmd))
    shm.seek(CMD_WRITE_OFFSET)
    shm.write((w + 1).to_bytes(8, 'little'))  # Publish (release)
    return cmd.commandId


def update_agent_heartbeat(shm) -> None:
    shm.seek(AGENT_HEARTBEAT_OFFSET)
    shm.write(int(time.time() * 1000).to_bytes(8, 'little'))


def main() -> bool:
    print("==================================================")
    print("   AutoPOE2 IPC Bridge Test (Protocol v1.1)       ")
    print("==================================================")

    try:
        shm = mmap.mmap(0, SHARED_MEM_SIZE, tagname=SHARED_MEM_NAME, access=mmap.ACCESS_WRITE)
    except Exception as e:
        print(f"[Error] Không tìm thấy Shared Memory. Hãy chắc chắn CoreEngine đang chạy! Lỗi: {e}")
        return False

    # 0. Chờ Core khởi tạo header hợp lệ
    header = None
    deadline = time.time() + 10
    while time.time() < deadline:
        shm.seek(0)
        h = SharedMemoryHeader.from_buffer_copy(shm.read(ctypes.sizeof(SharedMemoryHeader)))
        if h.magic == SHARED_MEM_MAGIC and h.isRunning == 1:
            header = h
            break
        time.sleep(0.2)

    if header is None:
        print("[FAIL] Core không phản hồi hoặc magic header không khớp!")
        return False

    print(f"  - Magic Header: 0x{header.magic:X} (Chuẩn: 0x{SHARED_MEM_MAGIC:X})")
    print(f"  - Phiên bản Giao thức: 0x{header.protocolVersion:X}")
    print(f"  - Trạng thái Core: {'HOẠT ĐỘNG' if header.isRunning == 1 else 'TẮT'}")

    # 1. Đọc telemetry qua Seqlock double-buffer
    print("\n[1] Đọc telemetry qua Seqlock double-buffer...")
    result = None
    deadline = time.time() + 5
    while time.time() < deadline:
        result = seqlock_read(shm)
        if result:
            break
        time.sleep(0.05)
    if result is None:
        print("[FAIL] Không đọc được snapshot telemetry nào trong 5 giây!")
        return False

    seq1, player = result
    print(f"  - Snapshot seq={seq1} | Player HP: {player.currentHP}/{player.maxHP}"
          f" | Spirit: {player.currentSpirit}/{player.maxSpirit}")

    # 2. Xác nhận sequence tăng liên tục (Core đang publish 120Hz)
    time.sleep(0.2)
    update_agent_heartbeat(shm)
    result2 = seqlock_read(shm)
    if result2 is None or result2[0] <= seq1:
        print("[FAIL] Telemetry sequence không tăng - Core có thể đã dừng!")
        return False
    print(f"  - Sequence tăng liên tục: {seq1} -> {result2[0]} (Core publish 120Hz OK)")

    # 3. Gửi 3 lệnh vào SPSC Command Ring
    print("\n[2] Gửi 3 lệnh vào SPSC Command Ring...")
    sent = []
    sent.append(publish_command(shm, OP_NAVIGATE_TO_COORD, 1280.0, 720.0))
    sent.append(publish_command(shm, OP_ENGAGE_TARGET, entity_id=12345))
    sent.append(publish_command(shm, OP_PICKUP_ITEM))
    if any(c is None for c in sent):
        print("[FAIL] Gửi lệnh thất bại!")
        return False
    print(f"  - Đã gửi lệnh #{sent[0]} (NAVIGATE 1280,720), #{sent[1]} (ENGAGE), #{sent[2]} (PICKUP)")

    # 4. Chờ Core tiêu thụ cả 3 lệnh (commandReadIndex tăng >= 3)
    print("\n[3] Chờ Core rút lệnh (SPSC FIFO)...")
    start_read = read_u64(shm, CMD_READ_OFFSET)
    consumed = False
    deadline = time.time() + 5
    while time.time() < deadline:
        update_agent_heartbeat(shm)  # Giữ nhịp tim để watchdog không ngắt giữa chừng
        if read_u64(shm, CMD_READ_OFFSET) - start_read >= 3:
            consumed = True
            break
        time.sleep(0.05)

    if not consumed:
        print("[FAIL] Core không tiêu thụ lệnh trong 5 giây!")
        return False
    print("  - Core đã tiêu thụ cả 3 lệnh theo đúng FIFO (readIndex tăng >= 3)")

    # 5. Báo cáo nhịp tim 2 chiều
    core_hb = read_u64(shm, 0x28)
    agent_hb = read_u64(shm, AGENT_HEARTBEAT_OFFSET)
    now_ms = int(time.time() * 1000)
    print(f"\n[4] Nhịp tim: Core lệch {now_ms - core_hb}ms | Agent lệch {now_ms - agent_hb}ms")
    print("[SUCCESS] Thử nghiệm IPC v1.1 (Seqlock + SPSC Ring + Heartbeat) THÀNH CÔNG!")
    print("[Info] Script kết thúc -> Core sẽ tự ngắt sau 5s mất nhịp tim Agent (watchdog demo).")
    return True


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)