#!/usr/bin/env python3
"""Serial VLM Client — 通过串口调用 VLM 推理。
用在客户端 (192.168.175.65)，模拟 OpenAI client.chat.completions.create 接口。
"""

import json, struct, time, serial

MAGIC = b'\xAB\xCD'
FLAG_CMD  = ord('C')
FLAG_IMG  = ord('I')
FLAG_RESP = ord('R')


class SerialVLMError(Exception):
    pass


class SerialVLMClient:
    """串口 VLM 客户端，兼容 OpenAI client.chat.completions.create 调用模式。"""

    def __init__(self, port='/dev/ttyUSB0', baud=921600, timeout=95):
        self.port = port
        self.baud = baud
        self.timeout = timeout
        self._ser = None

    def _connect(self):
        if self._ser is None or not self._ser.is_open:
            self._ser = serial.Serial(self.port, self.baud, timeout=0.1)
            self._ser.reset_input_buffer()
            self._ser.reset_output_buffer()

    def _send_packet(self, flag, data):
        if isinstance(data, str):
            data = data.encode('utf-8')
        self._ser.write(MAGIC)
        self._ser.write(bytes([flag]))
        self._ser.write(struct.pack('>I', len(data)))
        self._ser.write(data)
        self._ser.flush()

    def _recv_exact(self, n, timeout=15):
        deadline = time.time() + timeout
        buf = b''
        while len(buf) < n:
            remain = n - len(buf)
            chunk = self._ser.read(remain)
            if chunk:
                buf += chunk
            if time.time() > deadline:
                raise SerialVLMError(f"recv timeout: got {len(buf)}/{n}")
            if not chunk:
                time.sleep(0.005)
        return buf

    def _recv_packet(self):
        while True:
            b = self._ser.read(1)
            if not b:
                continue
            if b == b'\xAB':
                b2 = self._ser.read(1)
                if b2 == b'\xCD':
                    break
        flag = self._recv_exact(1)[0]
        length = struct.unpack('>I', self._recv_exact(4))[0]
        data = self._recv_exact(length)
        return flag, data

    def infer(self, image_bytes: bytes, prompt: str) -> dict:
        """发送图片+提示词，返回 {"success": bool, "answer": str, "elapsed_sec": float}"""
        self._connect()
        try:
            # Reset buffers
            self._ser.reset_input_buffer()
            self._ser.reset_output_buffer()

            # Send command
            cmd = json.dumps({"prompt": prompt, "image_len": len(image_bytes)})
            self._send_packet(FLAG_CMD, cmd)

            # Send image
            self._send_packet(FLAG_IMG, image_bytes)

            # Receive response
            t0 = time.time()
            flag, data = self._recv_packet()
            elapsed = time.time() - t0

            if flag != FLAG_RESP:
                return {"success": False, "error": f"unexpected flag: {flag}"}

            result = json.loads(data.decode('utf-8'))
            return result

        except SerialVLMError:
            return {"success": False, "error": "serial timeout"}
        except Exception as e:
            return {"success": False, "error": str(e)}

    def close(self):
        if self._ser and self._ser.is_open:
            self._ser.close()
            self._ser = None

    # ---- OpenAI-compatible interface ----
    class _FakeChoices:
        class _Msg:
            def __init__(self, content): self.content = content
        def __init__(self, content):
            self.message = self._Msg(content)

    class _FakeResponse:
        def __init__(self, content):
            self.choices = [SerialVLMClient._FakeChoices(content)]

    class _FakeCompletions:
        def __init__(self, client): self._client = client
        def create(self, *, model=None, messages=None, max_tokens=None, temperature=None, timeout=None):
            # Extract image and text from OpenAI-format messages
            import base64
            image_bytes = None
            prompt = ""
            for msg in (messages or []):
                content = msg.get("content", "")
                if isinstance(content, list):
                    for item in content:
                        t = item.get("type", "")
                        if t == "text":
                            prompt = item.get("text", "")
                        elif t == "image_url":
                            url = item.get("image_url", {}).get("url", "")
                            if url.startswith("data:") and "," in url:
                                image_bytes = base64.b64decode(url.split(",", 1)[1])
            if image_bytes is None:
                raise SerialVLMError("no image in messages")

            result = self._client.infer(image_bytes, prompt)
            if not result.get("success"):
                raise SerialVLMError(result.get("error", "unknown"))
            return SerialVLMClient._FakeResponse(result["answer"])

    class _FakeChat:
        def __init__(self, client): self.completions = SerialVLMClient._FakeCompletions(client)

    @property
    def chat(self):
        return self._FakeChat(self)


# Convenience function
def serial_infer(image_bytes, prompt, port='/dev/ttyUSB0', baud=921600):
    client = SerialVLMClient(port=port, baud=baud)
    try:
        return client.infer(image_bytes, prompt)
    finally:
        client.close()
