import asyncio
import json
import time
from collections import Counter, defaultdict

import numpy as np
import zmq
import zmq.asyncio as azmq
from pydantic_core import ValidationError

from control_backend.agents import BaseAgent
from control_backend.agents.perception.visual_emotion_recognition_agent.visual_emotion_recognizer import (  # noqa
    DeepFaceEmotionRecognizer,
)
from control_backend.core.agent_system import InternalMessage
from control_backend.core.config import settings
from control_backend.schemas.belief_message import Belief


class VisualEmotionRecognitionAgent(BaseAgent):
    def __init__(
        self,
        name: str,
        socket_address: str,
        bind: bool = False,
        timeout_ms: int = 1000,
        window_duration: int = settings.behaviour_settings.visual_emotion_recognition_window_duration_s,  # noqa
        min_frames_required: int = settings.behaviour_settings.visual_emotion_recognition_min_frames_per_face,  # noqa
    ):
        """
        Initialize the Visual Emotion Recognition Agent.

        :param name: Name of the agent
        :param socket_address: Address of the socket to connect or bind to
        :param bind: Whether to bind to the socket address (True) or connect (False)
        :param timeout_ms: Timeout for socket receive operations in milliseconds
        :param window_duration: Duration in seconds over which to aggregate emotions
        :param min_frames_required: Minimum number of frames per face required to consider a face
        valid
        """
        super().__init__(name)
        self.socket_address = socket_address
        self.socket_bind = bind
        self.timeout_ms = timeout_ms
        self.window_duration = window_duration
        self.min_frames_required = min_frames_required

        # Pause functionality
        # NOTE: flag is set when running, cleared when paused
        self._paused = asyncio.Event()
        self._paused.set()

    async def setup(self):
        """
        Initialize the agent resources.
        1. Initializes the :class:`VisualEmotionRecognizer`.
        2. Connects to the video input ZMQ socket.
        3. Starts the background emotion recognition loop.
        """
        self.logger.info("Setting up %s.", self.name)

        self.emotion_recognizer = DeepFaceEmotionRecognizer()

        self.video_in_socket = azmq.Context.instance().socket(zmq.SUB)

        self.video_in_socket.setsockopt(zmq.RCVHWM, 3)

        if self.socket_bind:
            self.video_in_socket.bind(self.socket_address)
        else:
            self.video_in_socket.connect(self.socket_address)

        self.video_in_socket.setsockopt_string(zmq.SUBSCRIBE, "")
        self.video_in_socket.setsockopt(zmq.RCVTIMEO, self.timeout_ms)

        self.add_behavior(self.emotion_update_loop())

    async def emotion_update_loop(self):
        """
        Background loop to receive video frames, recognize emotions, and update beliefs.
        1. Receives video frames from the ZMQ socket.
        2. Uses the :class:`VisualEmotionRecognizer` to detect emotions.
        3. Aggregates emotions over a time window.
        4. Sends updates to the BDI Core Agent about detected emotions.
        """
        # Next time to process the window and update emotions
        next_window_time = time.time() + self.window_duration

        # Tracks counts of detected emotions per face index
        face_stats = defaultdict(Counter)

        prev_dominant_emotions = set()

        while self._running:
            try:
                await self._paused.wait()

                width, height, image_bytes = await self.video_in_socket.recv_multipart()

                width = int.from_bytes(width, 'little')
                height = int.from_bytes(height, 'little')

                # Convert bytes to a numpy buffer
                image_array = np.frombuffer(image_bytes, np.uint8)

                frame = image_array.reshape((height, width, 3))

                # Get the dominant emotion from each face
                current_emotions = self.emotion_recognizer.sorted_dominant_emotions(frame)
                # Update emotion counts for each detected face
                for i, emotion in enumerate(current_emotions):
                    face_stats[i][emotion] += 1

                # If window duration has passed, process the collected stats
                if time.time() >= next_window_time:
                    window_dominant_emotions = set()
                    # Determine dominant emotion for each face in the window
                    for _, counter in face_stats.items():
                        total_detections = sum(counter.values())

                        if total_detections >= self.min_frames_required:
                            dominant_emotion = counter.most_common(1)[0][0]
                            window_dominant_emotions.add(dominant_emotion)

                    await self.update_emotions(prev_dominant_emotions, window_dominant_emotions)
                    prev_dominant_emotions = window_dominant_emotions
                    face_stats.clear()
                    next_window_time = time.time() + self.window_duration

            except zmq.Again:
                pass

            except Exception as e:
                self.logger.error(f"Error in emotion recognition loop: {e}")

    
    async def update_emotions(self, prev_emotions: set[str], emotions: set[str]):
        """
        Compare emotions from previous window and current emotions,
        send updates to BDI Core Agent.
        """
        emotions_to_remove = prev_emotions - emotions
        emotions_to_add = emotions - prev_emotions

        if not emotions_to_add and not emotions_to_remove:
            return

        emotion_beliefs_remove = []
        for emotion in emotions_to_remove:
            self.logger.info(f"Emotion '{emotion}' has disappeared.")
            try:
                emotion_beliefs_remove.append(
                    Belief(name="emotion_detected", arguments=[emotion], remove=True)
                )
            except ValidationError:
                self.logger.warning("Invalid belief for emotion removal: %s", emotion)

        emotion_beliefs_add = []
        for emotion in emotions_to_add:
            self.logger.info(f"New emotion detected: '{emotion}'")
            try:
                emotion_beliefs_add.append(Belief(name="emotion_detected", arguments=[emotion]))
            except ValidationError:
                self.logger.warning("Invalid belief for new emotion: %s", emotion)

        beliefs_list_add = [b.model_dump() for b in emotion_beliefs_add]
        beliefs_list_remove = [b.model_dump() for b in emotion_beliefs_remove]
        payload = {"create": beliefs_list_add, "delete": beliefs_list_remove}

        message = InternalMessage(
            to=settings.agent_settings.bdi_core_name,
            sender=self.name,
            body=json.dumps(payload),
            thread="beliefs",
        )
        await self.send(message)

    async def handle_message(self, msg: InternalMessage):
        """
        Handle incoming messages.

        Expects messages to pause or resume the Visual Emotion Recognition 
        processing from User Interrupt Agent.

        :param msg: The received internal message.
        """
        sender = msg.sender

        if sender == settings.agent_settings.user_interrupt_name:
            if msg.body == "PAUSE":
                self.logger.info("Pausing Visual Emotion Recognition processing.")
                self._paused.clear()
            elif msg.body == "RESUME":
                self.logger.info("Resuming Visual Emotion Recognition processing.")
                self._paused.set()
            else:
                self.logger.warning(f"Unknown command from User Interrupt Agent: {msg.body}")
        else:
            self.logger.debug(f"Ignoring message from unknown sender: {sender}")

    async def stop(self):
        """
        Clean up resources used by the agent.
        """
        self.video_in_socket.close()
        await super().stop()
        
