#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
VLM 图生文节点 —— 收到触发信号后拍图发给 VLM 服务，结果调用 TTS 服务播报
"""
import rclpy
from rclpy.node import Node
from std_msgs.msg import Int32, String
from sensor_msgs.msg import CompressedImage
from cv_bridge import CvBridge
import cv2
import base64
import threading
from openai import OpenAI
import os
import time
import numpy as np

from origincar_msg.srv import Speak


class VLMProcessor(Node):
    def __init__(self):
        super().__init__('vlm_detect')

        # 声明 ROS2 参数
        self.declare_parameter('vlm_host', 'http://192.168.10.189:8000')
        self.declare_parameter('vlm_model', './OpenGVLab/InternVL3-1B/')
        self.declare_parameter('image_topic', '/image_mjpeg')
        self.declare_parameter('trigger_topic', '/sign4return')
        self.declare_parameter('trigger_sign', 9)
        self.declare_parameter('result_topic', '/vlm_result')
        self.declare_parameter('prompt_text', '请描述这张图片的内容，用一句简短的话概括，不超过20个字。')
        self.declare_parameter('max_tokens', 100)

        vlm_host = self.get_parameter('vlm_host').value
        vlm_model = self.get_parameter('vlm_model').value
        image_topic = self.get_parameter('image_topic').value
        trigger_topic = self.get_parameter('trigger_topic').value
        self.trigger_sign = self.get_parameter('trigger_sign').value
        result_topic = self.get_parameter('result_topic').value
        self.prompt_text = self.get_parameter('prompt_text').value
        self.max_tokens = self.get_parameter('max_tokens').value

        # 初始化 OpenAI 客户端
        self.client = OpenAI(
            base_url=f"{vlm_host}/v1",
            api_key="EMPTY",
        )
        self.vlm_model = vlm_model

        # ROS2 通信
        self.bridge = CvBridge()
        self.latest_image = None
        self.image_lock = threading.Lock()

        self.image_sub = self.create_subscription(
            CompressedImage, image_topic, self.image_callback, 10
        )
        self.sign_sub = self.create_subscription(
            Int32, trigger_topic, self.sign_callback, 10
        )
        self.result_pub = self.create_publisher(String, result_topic, 10)

        # TTS 服务客户端
        self.tts_client = self.create_client(Speak, '/tts/speak')
        while not self.tts_client.wait_for_service(timeout_sec=5.0):
            self.get_logger().info('Waiting for TTS service...')

        self.get_logger().info(
            f"VLM Processor 就绪 | host={vlm_host} | model={vlm_model} | "
            f"image={image_topic} | trigger={trigger_topic}(sign={self.trigger_sign})"
        )

    def image_callback(self, msg):
        with self.image_lock:
            try:
                np_arr = np.frombuffer(msg.data, np.uint8)
                cv_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
                self.latest_image = cv_image
                self.get_logger().debug("图片已接收")
            except Exception as e:
                self.get_logger().error(f"图片接收错误: {e}")

    def sign_callback(self, msg):
        if msg.data == self.trigger_sign:
            self.get_logger().info(f"收到触发信号 ({msg.data}), 开始推理...")
            with self.image_lock:
                if self.latest_image is None:
                    self.get_logger().warning("无可用图片")
                    return
                temp_path = "/tmp/vlm_temp_image.jpg"
                cv2.imwrite(temp_path, self.latest_image)
                self.get_logger().info(f"临时图片已保存: {temp_path}")

            try:
                description = self.process_image(temp_path)
                self.get_logger().info(f"图像描述: {description}")

                # 发布结果到话题
                result_msg = String()
                result_msg.data = description
                self.result_pub.publish(result_msg)

                # 调用 TTS 服务播报
                if self.tts_client.service_is_ready():
                    req = Speak.Request()
                    req.text = description
                    future = self.tts_client.call_async(req)
                    future.add_done_callback(self._tts_done_callback)
                else:
                    self.get_logger().warning('TTS service not available')

                os.remove(temp_path)
            except Exception as e:
                self.get_logger().error(f"图像推理失败: {e}")

    def _tts_done_callback(self, future):
        try:
            resp = future.result()
            if resp.success:
                self.get_logger().debug(f'TTS OK: {resp.message}')
            else:
                self.get_logger().warning(f'TTS failed: {resp.message}')
        except Exception as e:
            self.get_logger().error(f'TTS call error: {e}')

    def process_image(self, image_path):
        with open(image_path, "rb") as image_file:
            base64_image = base64.b64encode(image_file.read()).decode('utf-8')

        start_time = time.time()
        response = self.client.chat.completions.create(
            model=self.vlm_model,
            messages=[{
                "role": "user",
                "content": [
                    {"type": "text", "text": self.prompt_text},
                    {"type": "image_url", "image_url": {
                        "url": f"data:image/jpeg;base64,{base64_image}"
                    }},
                ]
            }],
            max_tokens=self.max_tokens,
            timeout=30,
        )
        self.get_logger().info(f"VLM 推理耗时 {time.time() - start_time:.1f}s")
        return response.choices[0].message.content


def main(args=None):
    rclpy.init(args=args)
    node = VLMProcessor()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
