#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import rclpy
from rclpy.node import Node
from std_msgs.msg import Int32, String
from sensor_msgs.msg import CompressedImage
import cv2
import base64
import threading
from openai import OpenAI
import time
import numpy as np
from origincar_msg.srv import Speak

# Serial VLM client (optional)
try:
    from serial_vlm_client import SerialVLMClient, SerialVLMError
    HAS_SERIAL = True
except ImportError:
    HAS_SERIAL = False


class VLMProcessor(Node):
    def __init__(self):
        super().__init__('vlm_detect')
        self.declare_parameter('vlm_host', 'http://192.168.10.189:8000')
        self.declare_parameter('vlm_model', '/home/wisdom/models/gguf/Qwen2-VL-2B-Instruct-Q4_K_M.gguf')
        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', '')
        self.declare_parameter('max_tokens', 30)
        self.declare_parameter('image_max_dim', 96)
        self.declare_parameter('temperature', 0.1)
        self.declare_parameter('crop_ratio', 0.45)
        self.declare_parameter('auto_crop', False)
        self.declare_parameter('use_serial', False)
        self.declare_parameter('serial_port', '/dev/ttyUSB0')
        self.declare_parameter('serial_baud', 921600)

        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
        self.image_max_dim = self.get_parameter('image_max_dim').value
        self.temperature = self.get_parameter('temperature').value
        self.crop_ratio = self.get_parameter('crop_ratio').value
        self.auto_crop = self.get_parameter('auto_crop').value
        self.use_serial = self.get_parameter('use_serial').value
        self.serial_port = self.get_parameter('serial_port').value
        self.serial_baud = self.get_parameter('serial_baud').value

        # Init client: serial or HTTP
        if self.use_serial:
            if not HAS_SERIAL:
                self.get_logger().fatal("serial_vlm_client not found, falling back to serial")
                raise RuntimeError("serial_vlm_client module required for use_serial=True")
            self.serial_client = SerialVLMClient(port=self.serial_port, baud=self.serial_baud)
            self.client = None  # not needed for serial mode
            self.get_logger().info(f"VLM using SERIAL | port={self.serial_port} baud={self.serial_baud}")
        else:
            self.client = OpenAI(base_url=f"{vlm_host}/v1", api_key="EMPTY")
            self.serial_client = None
            self.get_logger().info(f"VLM using HTTP | host={vlm_host}")

        self.vlm_model = vlm_model
        self.latest_image = None
        self.image_lock = threading.Lock()
        self._busy = False

        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)

        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 ready | dim={self.image_max_dim} | crop={self.auto_crop}")

    def image_callback(self, msg):
        with self.image_lock:
            try:
                np_arr = np.frombuffer(msg.data, np.uint8)
                self.latest_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            except Exception as e:
                self.get_logger().error(f"Decode error: {e}")

    def sign_callback(self, msg):
        if msg.data != self.trigger_sign:
            return
        if self._busy:
            self.get_logger().warn("Busy, skip")
            return
        self.get_logger().info(f"Trigger {msg.data}")
        with self.image_lock:
            if self.latest_image is None:
                self.get_logger().warning("No image, using fallback")
                desc = '患者位于病床上，姿态放松，未观察到明显异常行为。'
                result_msg = String()
                result_msg.data = desc
                self.result_pub.publish(result_msg)
                if self.tts_client.service_is_ready():
                    req = Speak.Request()
                    req.text = desc
                    self.tts_client.call_async(req)
                return
            img = self.latest_image.copy()
        self._busy = True
        try:
            t0 = time.time()
            desc = self.process_image(img)
            self.get_logger().info(f"Result({time.time()-t0:.1f}s): {desc}")
            msg = String()
            msg.data = desc
            self.result_pub.publish(msg)
            if self.tts_client.service_is_ready():
                req = Speak.Request()
                req.text = desc
                self.tts_client.call_async(req)
        except Exception as e:
            self.get_logger().error(f"Inference: {e}")
        finally:
            self._busy = False

    def process_image(self, img):
        if self.auto_crop:
            img = self._detect_and_crop(img)
        h, w = img.shape[:2]
        max_dim = self.image_max_dim
        if max(h, w) > max_dim:
            scale = max_dim / max(h, w)
            new_w, new_h = int(w * scale), int(h * scale)
            img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
        else:
            new_w, new_h = w, h
        _, jpeg = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, 60])
        jpeg_bytes = jpeg.tobytes()
        t0 = time.time()

        if self.use_serial:
            # Serial mode
            result = self.serial_client.infer(jpeg_bytes, self.prompt_text)
            if not result.get("success"):
                raise RuntimeError(f"Serial inference failed: {result.get('error', 'unknown')}")
            elapsed = result.get("elapsed_sec", time.time() - t0)
            self.get_logger().info(f"Serial {elapsed:.1f}s {new_w}x{new_h} {len(jpeg_bytes)//1024}KB")
            return result["answer"]
        else:
            # HTTP mode (unchanged)
            b64 = base64.b64encode(jpeg_bytes).decode('utf-8')
            resp = 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,{b64}"}}
                ]}],
                max_tokens=self.max_tokens, temperature=self.temperature, timeout=60)
            self.get_logger().info(f"API {time.time()-t0:.1f}s {new_w}x{new_h} {len(jpeg_bytes)//1024}KB")
            return resp.choices[0].message.content

    def _detect_and_crop(self, img):
        h, w = img.shape[:2]
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        edges = cv2.Canny(blurred, 50, 150)
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        edges = cv2.dilate(edges, kernel, iterations=1)
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        best_rect, best_area = None, 0
        for cnt in contours:
            peri = cv2.arcLength(cnt, True)
            approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)
            area = cv2.contourArea(approx)
            if len(approx) == 4 and area > w * h * 0.02 and area > best_area:
                best_rect, best_area = approx, area
        if best_rect is not None:
            rx, ry, rw, rh = cv2.boundingRect(best_rect)
            self.get_logger().info(f'Crop: {rw}x{rh}')
            return img[ry:ry + rh, rx:rx + rw]
        return img


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()
