"""emotion_node.py — Facial expression recognition node (Issue #161).

Piggybacks on the face detection pipeline: subscribes to
/social/faces/detections (FaceDetectionArray), extracts face crops from the
latest camera frames, runs a TensorRT FP16 emotion CNN, applies per-person
EMA temporal smoothing, and publishes results on /social/faces/expressions.

Architecture
------------
  FaceDetectionArray → crop extraction → TRT FP16 inference (< 5 ms) →
  EMA smoothing → ExpressionArray publish

If TRT engine is not available the node falls back to the 5-point landmark
heuristic (classify_from_landmarks) which requires no GPU and adds < 0.1 ms.

ROS2 topics
-----------
  Subscribe:
    /social/faces/detections       (saltybot_social_msgs/FaceDetectionArray)
    /camera/{name}/image_raw       (sensor_msgs/Image) × 4
    /social/persons                (saltybot_social_msgs/PersonStateArray)
  Publish:
    /social/faces/expressions      (saltybot_social_msgs/ExpressionArray)
    /social/emotion/context        (std_msgs/String)  — JSON for LLM context

Parameters
----------
  engine_path          (str)   ""        — path to emotion_fp16.trt; empty = landmark only
  min_confidence       (float) 0.40      — suppress results below this threshold
  smoothing_alpha      (float) 0.30      — EMA weight (higher = faster, less stable)
  opt_out_persons      (str)   ""        — comma-separated person_ids that opted out
  face_min_size        (int)   24        — skip faces whose bbox is smaller (px side)
  landmark_fallback    (bool)  true      — use landmark heuristic when TRT unavailable
  camera_names         (str)   "front,left,rear,right"
  n_cameras            (int)   4
  publish_context      (bool)  true      — publish /social/emotion/context JSON
"""

from __future__ import annotations

import json
import threading
import time
from typing import Dict, List, Optional, Tuple

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

from std_msgs.msg import String
from sensor_msgs.msg import Image

from saltybot_social_msgs.msg import (
    FaceDetectionArray,
    PersonStateArray,
    ExpressionArray,
    Expression,
)
from saltybot_social.emotion_classifier import (
    EmotionClassifier,
    PersonEmotionTracker,
    EMOTIONS,
)

try:
    from cv_bridge import CvBridge
    _CV_BRIDGE_OK = True
except ImportError:
    _CV_BRIDGE_OK = False

try:
    import cv2
    _CV2_OK = True
except ImportError:
    _CV2_OK = False


# ── Per-camera latest-frame buffer ────────────────────────────────────────────


class _FrameBuffer:
    """Thread-safe store of the most recent image per camera."""

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._frames: Dict[int, object] = {}   # camera_id → cv2 BGR image

    def put(self, camera_id: int, img) -> None:
        with self._lock:
            self._frames[camera_id] = img

    def get(self, camera_id: int):
        with self._lock:
            return self._frames.get(camera_id)


# ── ROS2 Node ─────────────────────────────────────────────────────────────────


class EmotionNode(Node):
    """Facial expression recognition — TRT FP16 + landmark fallback."""

    def __init__(self) -> None:
        super().__init__("emotion_node")

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("engine_path",      "")
        self.declare_parameter("min_confidence",   0.40)
        self.declare_parameter("smoothing_alpha",  0.30)
        self.declare_parameter("opt_out_persons",  "")
        self.declare_parameter("face_min_size",    24)
        self.declare_parameter("landmark_fallback", True)
        self.declare_parameter("camera_names",     "front,left,rear,right")
        self.declare_parameter("n_cameras",        4)
        self.declare_parameter("publish_context",  True)

        engine_path   = self.get_parameter("engine_path").value
        self._min_conf = self.get_parameter("min_confidence").value
        alpha          = self.get_parameter("smoothing_alpha").value
        opt_out_str    = self.get_parameter("opt_out_persons").value
        self._face_min = self.get_parameter("face_min_size").value
        self._lm_fallback = self.get_parameter("landmark_fallback").value
        cam_names_str  = self.get_parameter("camera_names").value
        n_cameras      = self.get_parameter("n_cameras").value
        self._pub_ctx  = self.get_parameter("publish_context").value

        # ── Classifier + tracker ───────────────────────────────────────────
        self._classifier = EmotionClassifier(engine_path=engine_path, alpha=alpha)
        self._tracker    = PersonEmotionTracker(alpha=alpha)

        # Parse opt-out list
        for pid_str in opt_out_str.split(","):
            pid_str = pid_str.strip()
            if pid_str.isdigit():
                self._tracker.set_opt_out(int(pid_str), True)

        # ── Camera frame buffer + cv_bridge ───────────────────────────────
        self._frame_buf = _FrameBuffer()
        self._bridge    = CvBridge() if _CV_BRIDGE_OK else None
        self._cam_names = [n.strip() for n in cam_names_str.split(",")]
        self._cam_id_map: Dict[str, int] = {
            name: idx for idx, name in enumerate(self._cam_names)
        }

        # ── Latest persons for person_id ↔ face_id correlation ────────────
        self._persons_lock = threading.Lock()
        self._face_to_person: Dict[int, int] = {}   # face_id → person_id

        # ── Context for LLM (latest emotion per person) ───────────────────
        self._emotion_context: Dict[int, str] = {}  # person_id → emotion label

        # ── QoS profiles ──────────────────────────────────────────────────
        be_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=4,
        )

        # ── Subscriptions ──────────────────────────────────────────────────
        self.create_subscription(
            FaceDetectionArray,
            "/social/faces/detections",
            self._on_faces,
            be_qos,
        )
        self.create_subscription(
            PersonStateArray,
            "/social/persons",
            self._on_persons,
            be_qos,
        )
        for name in self._cam_names[:n_cameras]:
            cam_id = self._cam_id_map.get(name, 0)
            self.create_subscription(
                Image,
                f"/camera/{name}/image_raw",
                lambda msg, cid=cam_id: self._on_image(msg, cid),
                be_qos,
            )

        # ── Publishers ─────────────────────────────────────────────────────
        self._expr_pub = self.create_publisher(
            ExpressionArray, "/social/faces/expressions", be_qos
        )
        if self._pub_ctx:
            self._ctx_pub = self.create_publisher(
                String, "/social/emotion/context", 10
            )
        else:
            self._ctx_pub = None

        # ── Load TRT engine in background ──────────────────────────────────
        threading.Thread(target=self._load_engine, daemon=True).start()

        self.get_logger().info(
            f"EmotionNode ready — engine='{engine_path}', "
            f"alpha={alpha:.2f}, min_conf={self._min_conf:.2f}, "
            f"cameras={self._cam_names[:n_cameras]}"
        )

    # ── Engine loading ────────────────────────────────────────────────────────

    def _load_engine(self) -> None:
        if not self._classifier._engine_path:
            if self._lm_fallback:
                self.get_logger().info(
                    "No TRT engine path — using landmark heuristic fallback"
                )
            else:
                self.get_logger().warn(
                    "No TRT engine and landmark_fallback=false — "
                    "emotion_node will not classify"
                )
            return

        ok = self._classifier.load()
        if ok:
            self.get_logger().info("TRT emotion engine loaded ✓")
        else:
            self.get_logger().warn(
                "TRT engine load failed — falling back to landmark heuristic"
            )

    # ── Camera frame ingestion ────────────────────────────────────────────────

    def _on_image(self, msg: Image, camera_id: int) -> None:
        if not _CV_BRIDGE_OK or self._bridge is None:
            return
        try:
            bgr = self._bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
            self._frame_buf.put(camera_id, bgr)
        except Exception as exc:
            self.get_logger().debug(f"cv_bridge error cam{camera_id}: {exc}")

    # ── Person-state for face_id → person_id mapping ─────────────────────────

    def _on_persons(self, msg: PersonStateArray) -> None:
        mapping: Dict[int, int] = {}
        for ps in msg.persons:
            if ps.face_id >= 0:
                mapping[ps.face_id] = ps.person_id
        with self._persons_lock:
            self._face_to_person = mapping

    def _get_person_id(self, face_id: int) -> int:
        with self._persons_lock:
            return self._face_to_person.get(face_id, -1)

    # ── Main face-detection callback ──────────────────────────────────────────

    def _on_faces(self, msg: FaceDetectionArray) -> None:
        if not msg.faces:
            return

        expressions: List[Expression] = []

        for face in msg.faces:
            person_id = self._get_person_id(face.face_id)

            # Opt-out check
            if person_id >= 0 and self._tracker.is_opt_out(person_id):
                expr = Expression()
                expr.header = msg.header
                expr.person_id = person_id
                expr.face_id   = face.face_id
                expr.emotion   = ""
                expr.confidence = 0.0
                expr.scores    = [0.0] * 7
                expr.is_opt_out = True
                expr.source     = "opt_out"
                expressions.append(expr)
                continue

            # Try TRT crop classification first
            result = None
            if self._classifier.ready and _CV2_OK:
                result = self._classify_crop(face, person_id)

            # Fallback to landmark heuristic
            if result is None and self._lm_fallback:
                result = self._classifier.classify_from_landmarks(
                    list(face.landmarks),
                    person_id=person_id,
                    tracker=self._tracker,
                )

            if result is None:
                continue

            if result.confidence < self._min_conf:
                continue

            # Build ROS message
            expr = Expression()
            expr.header     = msg.header
            expr.person_id  = person_id
            expr.face_id    = face.face_id
            expr.emotion    = result.emotion
            expr.confidence = result.confidence
            # Pad/trim scores to exactly 7
            sc = result.scores
            expr.scores = (sc + [0.0] * 7)[:7]
            expr.is_opt_out = False
            expr.source     = result.source
            expressions.append(expr)

            # Update LLM context cache
            if person_id >= 0:
                self._emotion_context[person_id] = result.emotion

        if not expressions:
            return

        # Publish ExpressionArray
        arr = ExpressionArray()
        arr.header = msg.header
        arr.expressions = expressions
        self._expr_pub.publish(arr)

        # Publish JSON context for conversation node
        if self._ctx_pub is not None:
            self._publish_context()

        self._tracker.tick()

    def _classify_crop(self, face, person_id: int):
        """Extract crop from frame buffer and run TRT inference."""
        # Resolve camera_id from face header frame_id
        frame_id = face.header.frame_id if face.header.frame_id else "front"
        cam_name = frame_id.split("/")[-1].split("_")[0]   # "front", "left", etc.
        camera_id = self._cam_id_map.get(cam_name, 0)

        frame = self._frame_buf.get(camera_id)
        if frame is None:
            return None

        h, w = frame.shape[:2]
        x1 = max(0, int(face.bbox_x * w))
        y1 = max(0, int(face.bbox_y * h))
        x2 = min(w, int((face.bbox_x + face.bbox_w) * w))
        y2 = min(h, int((face.bbox_y + face.bbox_h) * h))

        if (x2 - x1) < self._face_min or (y2 - y1) < self._face_min:
            return None

        crop = frame[y1:y2, x1:x2]
        if crop.size == 0:
            return None

        return self._classifier.classify_crop(
            crop, person_id=person_id, tracker=self._tracker
        )

    # ── LLM context publisher ─────────────────────────────────────────────────

    def _publish_context(self) -> None:
        """Publish latest-emotions dict as JSON for conversation_node."""
        ctx = {str(pid): emo for pid, emo in self._emotion_context.items()}
        msg = String()
        msg.data = json.dumps({"emotions": ctx, "ts": time.time()})
        self._ctx_pub.publish(msg)


# ── Entry point ───────────────────────────────────────────────────────────────


def main(args=None) -> None:
    rclpy.init(args=args)
    node = EmotionNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == "__main__":
    main()
