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

import json, struct, time, socket


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


class TCPVLMError(Exception):
    pass


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

    def __init__(self, host='192.168.127.1', port=9216, timeout=95):
        self.host = host
        self.port = port
        self.timeout = timeout

    def _recv_exact(self, sock, n, timeout=15):
        deadline = time.time() + timeout
        buf = b''
        while len(buf) < n:
            remain = n - len(buf)
            sock.settimeout(max(0.1, deadline - time.time()))
            try:
                chunk = sock.recv(remain)
            except socket.timeout:
                if time.time() > deadline:
                    raise TCPVLMError(f"recv timeout: got {len(buf)}/{n}")
                continue
            if not chunk:
                raise TCPVLMError("connection closed by server")
            buf += chunk
        return buf

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

    def _recv_packet(self, sock):
        while True:
            b = self._recv_exact(sock, 1, timeout=60)
            if b == b'\xAB':
                b2 = self._recv_exact(sock, 1, timeout=5)
                if b2 == b'\xCD':
                    break
        flag = self._recv_exact(sock, 1)[0]
        length = struct.unpack('>I', self._recv_exact(sock, 4))[0]
        if length > 1024 * 1024:
            raise TCPVLMError(f"packet too large: {length}")
        data = self._recv_exact(sock, length, timeout=30)
        return flag, data

    def infer(self, image_bytes: bytes, prompt: str) -> dict:
        """发送图片+提示词，返回 {"success": bool, "answer": str, "elapsed_sec": float}"""
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(self.timeout)
        try:
            sock.connect((self.host, self.port))

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

            # Send image
            self._send_packet(sock, FLAG_IMG, image_bytes)

            # Receive response
            t0 = time.time()
            flag, data = self._recv_packet(sock)
            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 TCPVLMError:
            return {"success": False, "error": "TCP timeout"}
        except Exception as e:
            return {"success": False, "error": str(e)}
        finally:
            try:
                sock.close()
            except Exception:
                pass

    # ---- 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 = [TCPVLMClient._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):
            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 TCPVLMError("no image in messages")
            result = self._client.infer(image_bytes, prompt)
            if not result.get("success"):
                raise TCPVLMError(result.get("error", "unknown"))
            return TCPVLMClient._FakeResponse(result["answer"])

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

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


def tcp_infer(image_bytes, prompt, host='192.168.127.1', port=9216):
    client = TCPVLMClient(host=host, port=port)
    return client.infer(image_bytes, prompt)
