#!/usr/bin/env python3
"""测试推理服务 - 发送一张测试图片"""
import socket
import struct
import json
import base64
import numpy as np
import cv2
import os

SOCK_PATH = "/tmp/edge-infer.sock"

def create_test_image():
    """创建一张 640x480 的测试图片，画几个几何图形"""
    img = np.zeros((480, 640, 3), dtype=np.uint8)
    # 画一些简单的形状模拟检测目标
    cv2.rectangle(img, (50, 50), (200, 200), (255, 255, 255), -1)
    cv2.circle(img, (400, 300), 80, (255, 255, 255), -1)
    # 编码为 JPEG
    _, buf = cv2.imencode('.jpg', img)
    return buf.tobytes()

def send_msg(conn, payload):
    data = json.dumps(payload).encode("utf-8")
    conn.sendall(struct.pack(">I", len(data)) + data)

def recv_msg(conn):
    hdr = conn.recv(4)
    if not hdr:
        return None
    length = struct.unpack(">I", hdr)[0]
    data = b""
    while len(data) < length:
        chunk = conn.recv(length - len(data))
        if not chunk:
            return None
        data += chunk
    return json.loads(data.decode("utf-8"))

def main():
    if not os.path.exists(SOCK_PATH):
        print(f"错误: socket {SOCK_PATH} 不存在")
        return
    
    print("连接推理服务...")
    conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    conn.connect(SOCK_PATH)
    
    # 创建测试图片
    jpeg_bytes = create_test_image()
    msg = {
        "stream_id": 0,
        "device_id": "test-001",
        "url": "test://local",
        "ts": 1234567890.0,
        "jpeg_b64": base64.b64encode(jpeg_bytes).decode("utf-8")
    }
    
    print("发送测试图片...")
    send_msg(conn, msg)
    
    print("等待推理结果...")
    result = recv_msg(conn)
    
    if result:
        print(f"\n推理结果:")
        print(f"  stream_id: {result['stream_id']}")
        print(f"  device_id: {result['device_id']}")
        print(f"  detections: {len(result['detections'])} 个目标")
        for d in result['detections']:
            print(f"    - {d['class']}: {d['conf']:.3f} bbox={d['bbox']}")
    else:
        print("未收到结果")
    
    conn.close()
    print("\n测试完成!")

if __name__ == "__main__":
    main()