"""
face_recognition_node.py -- ROS2 node for SCRFD face detection + ArcFace recognition.

Pipeline:
  1. Subscribe to /camera/color/image_raw (RealSense D435i color stream).
  2. Run SCRFD face detection (TensorRT FP16 or ONNX fallback).
  3. For each detected face, align and extract ArcFace embedding.
  4. Match embedding against persistent gallery.
  5. Publish FaceDetectionArray with identified faces.

Services:
  /social/enroll           -- Enroll a new person (collects N face samples).
  /social/persons/list     -- List all enrolled persons.
  /social/persons/delete   -- Delete a person from the gallery.
  /social/persons/update   -- Update a person's name.
"""

import time

import numpy as np

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

import cv2
from cv_bridge import CvBridge

from sensor_msgs.msg import Image
from builtin_interfaces.msg import Time

from saltybot_social_msgs.msg import (
    FaceDetection,
    FaceDetectionArray,
    FaceEmbedding,
    FaceEmbeddingArray,
)
from saltybot_social_msgs.srv import (
    EnrollPerson,
    ListPersons,
    DeletePerson,
    UpdatePerson,
)

from .scrfd_detector import SCRFDDetector
from .arcface_recognizer import ArcFaceRecognizer
from .face_gallery import FaceGallery


class FaceRecognitionNode(Node):
    """ROS2 node: SCRFD face detection + ArcFace gallery matching."""

    def __init__(self):
        super().__init__('face_recognizer')
        self._bridge = CvBridge()
        self._frame_count = 0
        self._fps_t0 = time.monotonic()

        # -- Parameters --------------------------------------------------------
        self.declare_parameter('scrfd_engine_path',
                               '/mnt/nvme/saltybot/models/scrfd_2.5g.engine')
        self.declare_parameter('scrfd_onnx_path',
                               '/mnt/nvme/saltybot/models/scrfd_2.5g_bnkps.onnx')
        self.declare_parameter('arcface_engine_path',
                               '/mnt/nvme/saltybot/models/arcface_r50.engine')
        self.declare_parameter('arcface_onnx_path',
                               '/mnt/nvme/saltybot/models/arcface_r50.onnx')
        self.declare_parameter('gallery_dir', '/mnt/nvme/saltybot/gallery')
        self.declare_parameter('recognition_threshold', 0.35)
        self.declare_parameter('publish_debug_image', False)
        self.declare_parameter('max_faces', 10)
        self.declare_parameter('scrfd_conf_thresh', 0.5)

        self._recognition_threshold = self.get_parameter('recognition_threshold').value
        self._pub_debug_flag = self.get_parameter('publish_debug_image').value
        self._max_faces = self.get_parameter('max_faces').value

        # -- Models ------------------------------------------------------------
        self._detector = SCRFDDetector(
            engine_path=self.get_parameter('scrfd_engine_path').value,
            onnx_path=self.get_parameter('scrfd_onnx_path').value,
            conf_thresh=self.get_parameter('scrfd_conf_thresh').value,
        )
        self._recognizer = ArcFaceRecognizer(
            engine_path=self.get_parameter('arcface_engine_path').value,
            onnx_path=self.get_parameter('arcface_onnx_path').value,
        )

        # -- Gallery -----------------------------------------------------------
        gallery_dir = self.get_parameter('gallery_dir').value
        self._gallery = FaceGallery(gallery_dir)
        self._gallery.load()
        self.get_logger().info('Gallery loaded: %d persons.', len(self._gallery))

        # -- Enrollment state --------------------------------------------------
        self._enrolling = None  # {name, samples_needed, collected: [embeddings]}

        # -- QoS profiles ------------------------------------------------------
        best_effort_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )
        reliable_qos = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            durability=DurabilityPolicy.TRANSIENT_LOCAL,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )

        # -- Subscribers -------------------------------------------------------
        self.create_subscription(
            Image,
            '/camera/color/image_raw',
            self._on_image,
            best_effort_qos,
        )

        # -- Publishers --------------------------------------------------------
        self._pub_detections = self.create_publisher(
            FaceDetectionArray, '/social/faces/detections', best_effort_qos)
        self._pub_embeddings = self.create_publisher(
            FaceEmbeddingArray, '/social/faces/embeddings', reliable_qos)
        if self._pub_debug_flag:
            self._pub_debug_img = self.create_publisher(
                Image, '/social/faces/debug_image', best_effort_qos)

        # -- Services ----------------------------------------------------------
        self.create_service(EnrollPerson, '/social/enroll', self._handle_enroll)
        self.create_service(ListPersons, '/social/persons/list', self._handle_list)
        self.create_service(DeletePerson, '/social/persons/delete', self._handle_delete)
        self.create_service(UpdatePerson, '/social/persons/update', self._handle_update)

        # Publish initial gallery state
        self._publish_gallery_embeddings()

        self.get_logger().info('FaceRecognitionNode ready.')

    # -- Image callback --------------------------------------------------------

    def _on_image(self, msg: Image):
        """Process incoming camera frame: detect, recognize, publish."""
        try:
            bgr = self._bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
        except Exception as e:
            self.get_logger().error('Image decode error: %s', str(e),
                                    throttle_duration_sec=5.0)
            return

        # Detect faces
        detections = self._detector.detect(bgr)

        # Limit face count
        if len(detections) > self._max_faces:
            detections = sorted(detections, key=lambda d: d['score'], reverse=True)
            detections = detections[:self._max_faces]

        # Build output message
        det_array = FaceDetectionArray()
        det_array.header = msg.header

        for det in detections:
            # Extract embedding and match gallery
            embedding = self._recognizer.align_and_embed(bgr, det['kps'])
            pid, pname, score = self._gallery.match(
                embedding, self._recognition_threshold)

            # Handle enrollment: collect embedding from largest face
            if self._enrolling is not None:
                self._enrollment_collect(det, embedding)

            # Build FaceDetection message
            face_msg = FaceDetection()
            face_msg.header = msg.header
            face_msg.face_id = pid
            face_msg.person_name = pname
            face_msg.confidence = det['score']
            face_msg.recognition_score = score

            bbox = det['bbox']
            face_msg.bbox_x = bbox[0]
            face_msg.bbox_y = bbox[1]
            face_msg.bbox_w = bbox[2] - bbox[0]
            face_msg.bbox_h = bbox[3] - bbox[1]

            kps = det['kps']
            for i in range(10):
                face_msg.landmarks[i] = kps[i]

            det_array.faces.append(face_msg)

        self._pub_detections.publish(det_array)

        # Debug image
        if self._pub_debug_flag and hasattr(self, '_pub_debug_img'):
            debug_img = self._draw_debug(bgr, detections, det_array.faces)
            self._pub_debug_img.publish(
                self._bridge.cv2_to_imgmsg(debug_img, encoding='bgr8'))

        # FPS logging
        self._frame_count += 1
        if self._frame_count % 30 == 0:
            elapsed = time.monotonic() - self._fps_t0
            fps = 30.0 / elapsed if elapsed > 0 else 0.0
            self._fps_t0 = time.monotonic()
            self.get_logger().info(
                'FPS: %.1f | faces: %d', fps, len(detections))

    # -- Enrollment logic ------------------------------------------------------

    def _enrollment_collect(self, det: dict, embedding: np.ndarray):
        """Collect an embedding sample during enrollment (largest face only)."""
        if self._enrolling is None:
            return

        # Only collect from the largest face (by bbox area)
        bbox = det['bbox']
        area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])

        if not hasattr(self, '_enroll_best_area'):
            self._enroll_best_area = 0.0
            self._enroll_best_embedding = None

        if area > self._enroll_best_area:
            self._enroll_best_area = area
            self._enroll_best_embedding = embedding

    def _enrollment_frame_end(self):
        """Called at end of each frame to finalize enrollment sample collection."""
        if self._enrolling is None or self._enroll_best_embedding is None:
            return

        self._enrolling['collected'].append(self._enroll_best_embedding)
        self._enroll_best_area = 0.0
        self._enroll_best_embedding = None

        collected = len(self._enrolling['collected'])
        needed = self._enrolling['samples_needed']
        self.get_logger().info('Enrollment: %d/%d samples for "%s".',
                               collected, needed, self._enrolling['name'])

        if collected >= needed:
            # Finalize enrollment
            name = self._enrolling['name']
            embeddings = self._enrolling['collected']
            mean_emb = np.mean(embeddings, axis=0).astype(np.float32)
            norm = np.linalg.norm(mean_emb)
            if norm > 0:
                mean_emb = mean_emb / norm

            pid = self._gallery.add_person(name, mean_emb, samples=len(embeddings))
            self._gallery.save()
            self._publish_gallery_embeddings()

            self.get_logger().info('Enrollment complete: person %d (%s).', pid, name)

            # Store result for the service callback
            self._enrolling['result_pid'] = pid
            self._enrolling['done'] = True
            self._enrolling = None

    # -- Service handlers ------------------------------------------------------

    def _handle_enroll(self, request, response):
        """Handle EnrollPerson service: start collecting face samples."""
        name = request.name.strip()
        if not name:
            response.success = False
            response.message = 'Name cannot be empty.'
            response.person_id = -1
            return response

        n_samples = request.n_samples if request.n_samples > 0 else 10

        self.get_logger().info('Starting enrollment for "%s" (%d samples).',
                               name, n_samples)

        # Set enrollment state — frames will collect embeddings
        self._enrolling = {
            'name': name,
            'samples_needed': n_samples,
            'collected': [],
            'done': False,
            'result_pid': -1,
        }
        self._enroll_best_area = 0.0
        self._enroll_best_embedding = None

        # Spin until enrollment is done (blocking service)
        rate = self.create_rate(10)  # 10 Hz check
        timeout_sec = n_samples * 2.0 + 10.0  # generous timeout
        t0 = time.monotonic()

        while not self._enrolling.get('done', False):
            # Finalize any pending frame collection
            self._enrollment_frame_end()

            if time.monotonic() - t0 > timeout_sec:
                self._enrolling = None
                response.success = False
                response.message = f'Enrollment timed out after {timeout_sec:.0f}s.'
                response.person_id = -1
                return response

            rclpy.spin_once(self, timeout_sec=0.1)

        response.success = True
        response.message = f'Enrolled "{name}" with {n_samples} samples.'
        response.person_id = self._enrolling.get('result_pid', -1) if self._enrolling else -1

        # Clean up (already set to None in _enrollment_frame_end on success)
        return response

    def _handle_list(self, request, response):
        """Handle ListPersons service: return all gallery entries."""
        entries = self._gallery.get_all()
        for entry in entries:
            emb_msg = FaceEmbedding()
            emb_msg.person_id = entry['person_id']
            emb_msg.person_name = entry['name']
            emb_msg.embedding = entry['embedding'].tolist()
            emb_msg.sample_count = entry['samples']

            secs = int(entry['enrolled_at'])
            nsecs = int((entry['enrolled_at'] - secs) * 1e9)
            emb_msg.enrolled_at = Time(sec=secs, nanosec=nsecs)

            response.persons.append(emb_msg)
        return response

    def _handle_delete(self, request, response):
        """Handle DeletePerson service: remove a person from the gallery."""
        if self._gallery.delete_person(request.person_id):
            self._gallery.save()
            self._publish_gallery_embeddings()
            response.success = True
            response.message = f'Deleted person {request.person_id}.'
        else:
            response.success = False
            response.message = f'Person {request.person_id} not found.'
        return response

    def _handle_update(self, request, response):
        """Handle UpdatePerson service: rename a person."""
        new_name = request.new_name.strip()
        if not new_name:
            response.success = False
            response.message = 'New name cannot be empty.'
            return response

        if self._gallery.update_name(request.person_id, new_name):
            self._gallery.save()
            self._publish_gallery_embeddings()
            response.success = True
            response.message = f'Updated person {request.person_id} to "{new_name}".'
        else:
            response.success = False
            response.message = f'Person {request.person_id} not found.'
        return response

    # -- Gallery publishing ----------------------------------------------------

    def _publish_gallery_embeddings(self):
        """Publish current gallery as FaceEmbeddingArray (latched-like)."""
        entries = self._gallery.get_all()
        msg = FaceEmbeddingArray()
        msg.header.stamp = self.get_clock().now().to_msg()

        for entry in entries:
            emb_msg = FaceEmbedding()
            emb_msg.person_id = entry['person_id']
            emb_msg.person_name = entry['name']
            emb_msg.embedding = entry['embedding'].tolist()
            emb_msg.sample_count = entry['samples']

            secs = int(entry['enrolled_at'])
            nsecs = int((entry['enrolled_at'] - secs) * 1e9)
            emb_msg.enrolled_at = Time(sec=secs, nanosec=nsecs)

            msg.embeddings.append(emb_msg)

        self._pub_embeddings.publish(msg)

    # -- Debug image -----------------------------------------------------------

    def _draw_debug(self, bgr: np.ndarray, detections: list[dict],
                    face_msgs: list) -> np.ndarray:
        """Draw bounding boxes, landmarks, and names on the image."""
        vis = bgr.copy()
        for det, face_msg in zip(detections, face_msgs):
            bbox = det['bbox']
            x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])

            # Color: green if recognized, yellow if unknown
            if face_msg.face_id >= 0:
                color = (0, 255, 0)
                label = f'{face_msg.person_name} ({face_msg.recognition_score:.2f})'
            else:
                color = (0, 255, 255)
                label = f'unknown ({face_msg.confidence:.2f})'

            cv2.rectangle(vis, (x1, y1), (x2, y2), color, 2)
            cv2.putText(vis, label, (x1, y1 - 8),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

            # Draw landmarks
            kps = det['kps']
            for k in range(5):
                px, py = int(kps[k * 2]), int(kps[k * 2 + 1])
                cv2.circle(vis, (px, py), 2, (0, 0, 255), -1)

        return vis


# -- Entry point ---------------------------------------------------------------

def main(args=None):
    """ROS2 entry point for face_recognition node."""
    rclpy.init(args=args)
    node = FaceRecognitionNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
