#!/usr/bin/env python3

import math
import time
from enum import Enum
from dataclasses import dataclass, field
from typing import Optional, Dict, List
from collections import deque

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
from std_msgs.msg import String, Float32, Bool
from geometry_msgs.msg import Pose
from saltybot_social_msgs.msg import (
    VoiceCommand,
    PersonState,
    PersonStateArray,
    Expression,
    PersonalityState,
)


class EmotionType(Enum):
    """Core emotion types for facial expression selection."""
    NEUTRAL = "neutral"
    HAPPY = "happy"
    CURIOUS = "curious"
    CONCERNED = "concerned"
    EXCITED = "excited"
    CONFUSED = "confused"
    TIRED = "tired"
    PLAYFUL = "playful"


@dataclass
class EmotionState:
    """Internal state of robot emotion system."""
    primary_emotion: EmotionType = EmotionType.NEUTRAL
    intensity: float = 0.5  # 0.0 = minimal, 1.0 = extreme
    confidence: float = 0.5  # 0.0 = uncertain, 1.0 = certain
    context: str = ""  # e.g., "person_interacting", "low_battery", "idle"
    triggered_by: str = ""  # e.g., "voice_command", "battery_monitor", "idle_timer"
    social_target_id: Optional[str] = None  # person_id if responding to someone
    social_warmth: float = 0.5  # 0.0 = cold, 1.0 = warm (familiarity with target)
    last_update_time: float = 0.0
    transition_start_time: float = 0.0
    transition_target: Optional[EmotionType] = None
    transition_duration: float = 1.0  # seconds for smooth transition
    idle_flags: Dict[str, bool] = field(default_factory=dict)  # blink, look_around, breathing
    expression_name: str = "neutral"  # actual face expression name


class IdleBehaviorManager:
    """Manages idle animations and subtle behaviors."""

    def __init__(self, node_logger):
        self.logger = node_logger
        self.blink_interval = 3.0  # seconds
        self.look_around_interval = 8.0
        self.breathing_phase = 0.0
        self.last_blink_time = time.time()
        self.last_look_around_time = time.time()

    def update(self, dt: float) -> Dict[str, bool]:
        """Update idle behaviors, return active flags."""
        current_time = time.time()
        flags = {}

        # Blink behavior: ~30% of the time
        if current_time - self.last_blink_time > self.blink_interval:
            flags["blink"] = True
            self.last_blink_time = current_time
            self.blink_interval = 3.0 + (hash(current_time) % 100) / 100.0
        else:
            flags["blink"] = False

        # Look around: softer transitions every 8 seconds
        if current_time - self.last_look_around_time > self.look_around_interval:
            flags["look_around"] = True
            self.last_look_around_time = current_time
            self.look_around_interval = 8.0 + (hash(current_time) % 200) / 100.0
        else:
            flags["look_around"] = False

        # Breathing: continuous gentle oscillation
        self.breathing_phase = (self.breathing_phase + dt * 0.5) % (2 * math.pi)
        breathing_intensity = (math.sin(self.breathing_phase) + 1.0) / 2.0
        flags["breathing"] = breathing_intensity > 0.3

        return flags


class SocialMemoryManager:
    """Tracks interaction history and familiarity with people."""

    def __init__(self):
        self.interactions: Dict[str, Dict] = {}
        self.max_history = 100
        self.recent_interactions = deque(maxlen=self.max_history)

    def record_interaction(self, person_id: str, interaction_type: str, intensity: float = 1.0):
        """Record an interaction with a person."""
        if person_id not in self.interactions:
            self.interactions[person_id] = {
                "count": 0,
                "warmth": 0.5,
                "last_interaction": 0.0,
                "positive_interactions": 0,
            }

        entry = self.interactions[person_id]
        entry["count"] += 1
        entry["last_interaction"] = time.time()

        if intensity > 0.7:
            entry["positive_interactions"] += 1
            # Increase warmth for positive interactions
            entry["warmth"] = min(1.0, entry["warmth"] + 0.05)
        else:
            # Slight decrease for negative interactions
            entry["warmth"] = max(0.3, entry["warmth"] - 0.02)

        self.recent_interactions.append((person_id, interaction_type, time.time()))

    def get_warmth_modifier(self, person_id: Optional[str]) -> float:
        """Get warmth multiplier for a person (0.5 = neutral, 1.0 = familiar, 0.3 = stranger)."""
        if not person_id or person_id not in self.interactions:
            return 0.5

        return self.interactions[person_id].get("warmth", 0.5)

    def get_familiarity_score(self, person_id: Optional[str]) -> float:
        """Get interaction count-based familiarity (normalized)."""
        if not person_id or person_id not in self.interactions:
            return 0.0

        count = self.interactions[person_id].get("count", 0)
        return min(1.0, count / 10.0)  # Saturate at 10 interactions


class EmotionTransitioner:
    """Handles smooth transitions between emotions."""

    def __init__(self):
        self.transition_curves = {
            (EmotionType.NEUTRAL, EmotionType.EXCITED): "ease_out",
            (EmotionType.NEUTRAL, EmotionType.CONCERNED): "ease_in",
            (EmotionType.EXCITED, EmotionType.NEUTRAL): "ease_in",
            (EmotionType.CONCERNED, EmotionType.NEUTRAL): "ease_out",
        }

    def get_transition_progress(self, state: EmotionState) -> float:
        """Get interpolation progress [0.0, 1.0]."""
        if not state.transition_target:
            return 1.0

        elapsed = time.time() - state.transition_start_time
        progress = min(1.0, elapsed / state.transition_duration)
        return progress

    def should_transition(self, state: EmotionState) -> bool:
        """Check if transition is complete."""
        if not state.transition_target:
            return False
        return self.get_transition_progress(state) >= 1.0

    def apply_transition(self, state: EmotionState) -> EmotionState:
        """Apply transition logic if in progress."""
        if not self.should_transition(state):
            return state

        # Transition complete
        state.primary_emotion = state.transition_target
        state.transition_target = None
        state.confidence = min(1.0, state.confidence + 0.1)
        return state

    def initiate_transition(
        self,
        state: EmotionState,
        target_emotion: EmotionType,
        duration: float = 0.8,
    ) -> EmotionState:
        """Start a smooth transition to new emotion."""
        if state.primary_emotion == target_emotion:
            return state

        state.transition_target = target_emotion
        state.transition_start_time = time.time()
        state.transition_duration = duration
        state.confidence = max(0.3, state.confidence - 0.2)
        return state


class EmotionEngineNode(Node):
    """
    Context-aware emotion engine for SaltyBot.

    Subscribes to:
    - /social/voice_command (reactive to speech)
    - /social/person_state (person tracking for social context)
    - /social/personality/state (personality/mood context)
    - /saltybot/battery (low battery detection)
    - /saltybot/balance (balance/stability concerns)
    - /diagnostics (health monitoring)

    Publishes:
    - /saltybot/emotion_state (current emotion + metadata)
    """

    def __init__(self):
        super().__init__("saltybot_emotion_engine")

        # Configuration parameters
        self.declare_parameter("personality.extroversion", 0.6)
        self.declare_parameter("personality.playfulness", 0.5)
        self.declare_parameter("personality.responsiveness", 0.8)
        self.declare_parameter("personality.anxiety", 0.3)
        self.declare_parameter("battery_warning_threshold", 0.25)
        self.declare_parameter("battery_critical_threshold", 0.10)
        self.declare_parameter("update_rate_hz", 10.0)

        # QoS for reliable topic communication
        qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=5,
        )

        # State tracking
        self.emotion_state = EmotionState()
        self.last_emotion_state = EmotionState()
        self.battery_level = 0.5
        self.balance_stable = True
        self.people_present: Dict[str, PersonState] = {}
        self.voice_command_cooldown = 0.0
        self.idle_timer = 0.0
        self.system_health = "nominal"

        # Managers
        self.idle_manager = IdleBehaviorManager(self.get_logger())
        self.social_memory = SocialMemoryManager()
        self.transitioner = EmotionTransitioner()

        # Subscriptions
        self.voice_sub = self.create_subscription(
            VoiceCommand,
            "/social/voice_command",
            self.voice_command_callback,
            qos,
        )

        self.person_state_sub = self.create_subscription(
            PersonStateArray,
            "/social/person_state",
            self.person_state_callback,
            qos,
        )

        self.personality_sub = self.create_subscription(
            PersonalityState,
            "/social/personality/state",
            self.personality_callback,
            qos,
        )

        self.battery_sub = self.create_subscription(
            Float32,
            "/saltybot/battery",
            self.battery_callback,
            qos,
        )

        self.balance_sub = self.create_subscription(
            Bool,
            "/saltybot/balance_stable",
            self.balance_callback,
            qos,
        )

        self.health_sub = self.create_subscription(
            String,
            "/saltybot/system_health",
            self.health_callback,
            qos,
        )

        # Publisher
        self.emotion_pub = self.create_publisher(
            String,
            "/saltybot/emotion_state",
            qos,
        )

        # Main update loop
        update_rate = self.get_parameter("update_rate_hz").value
        self.update_timer = self.create_timer(1.0 / update_rate, self.update_callback)

        self.get_logger().info(
            "Emotion engine initialized: "
            f"extroversion={self.get_parameter('personality.extroversion').value}, "
            f"playfulness={self.get_parameter('personality.playfulness').value}"
        )

    def voice_command_callback(self, msg: VoiceCommand):
        """React to voice commands with emotional responses."""
        self.voice_command_cooldown = 0.5  # Cooldown to prevent rapid re-triggering

        intent = msg.intent
        confidence = msg.confidence

        # Map command intents to emotions
        if intent.startswith("nav."):
            # Navigation commands -> excitement
            self.emotion_state = self.transitioner.initiate_transition(
                self.emotion_state,
                EmotionType.EXCITED,
                duration=0.6,
            )
            self.emotion_state.context = "navigation_command"
            self.emotion_state.intensity = min(0.9, confidence * 0.8 + 0.3)

        elif intent.startswith("social."):
            # Social commands -> happy/curious
            if "remember" in intent or "forget" in intent:
                self.emotion_state = self.transitioner.initiate_transition(
                    self.emotion_state,
                    EmotionType.CURIOUS,
                    duration=0.8,
                )
                self.emotion_state.intensity = 0.6
            else:
                self.emotion_state = self.transitioner.initiate_transition(
                    self.emotion_state,
                    EmotionType.HAPPY,
                    duration=0.7,
                )
                self.emotion_state.intensity = 0.7

        elif intent == "fallback":
            # Unrecognized command -> confused
            self.emotion_state = self.transitioner.initiate_transition(
                self.emotion_state,
                EmotionType.CONFUSED,
                duration=0.5,
            )
            self.emotion_state.intensity = min(0.5, confidence)

        self.emotion_state.triggered_by = "voice_command"
        self.emotion_state.social_target_id = msg.speaker_id

    def person_state_callback(self, msg: PersonStateArray):
        """Update state based on person tracking and engagement."""
        self.people_present.clear()
        for person_state in msg.person_states:
            person_id = str(person_state.person_id)
            self.people_present[person_id] = person_state

            # Record interaction based on engagement state
            if person_state.state == PersonState.STATE_ENGAGED:
                self.social_memory.record_interaction(person_id, "engaged", 0.8)
            elif person_state.state == PersonState.STATE_TALKING:
                self.social_memory.record_interaction(person_id, "talking", 0.9)
            elif person_state.state == PersonState.STATE_APPROACHING:
                self.social_memory.record_interaction(person_id, "approaching", 0.5)

        # If people present and engaged -> be happier
        engaged_count = sum(
            1 for p in self.people_present.values()
            if p.state == PersonState.STATE_ENGAGED
        )

        if engaged_count > 0:
            # Boost happiness when with familiar people
            playfulness = self.get_parameter("personality.playfulness").value
            if playfulness > 0.6:
                target_emotion = EmotionType.PLAYFUL
            else:
                target_emotion = EmotionType.HAPPY

            if self.emotion_state.primary_emotion != target_emotion:
                self.emotion_state = self.transitioner.initiate_transition(
                    self.emotion_state,
                    target_emotion,
                    duration=0.9,
                )
            self.emotion_state.intensity = 0.7 + (0.3 * playfulness)
            self.emotion_state.context = f"engaged_with_{engaged_count}_people"

    def personality_callback(self, msg: PersonalityState):
        """Update emotion context based on personality state."""
        # Mood from personality system influences intensity
        if msg.mood == "playful":
            self.emotion_state.intensity = min(1.0, self.emotion_state.intensity + 0.1)
        elif msg.mood == "annoyed":
            self.emotion_state.intensity = max(0.0, self.emotion_state.intensity - 0.1)

    def battery_callback(self, msg: Float32):
        """React to low battery with concern."""
        self.battery_level = msg.data

        battery_critical = self.get_parameter("battery_critical_threshold").value
        battery_warning = self.get_parameter("battery_warning_threshold").value

        if self.battery_level < battery_critical:
            # Critical: very concerned
            self.emotion_state = self.transitioner.initiate_transition(
                self.emotion_state,
                EmotionType.CONCERNED,
                duration=0.5,
            )
            self.emotion_state.intensity = 0.9
            self.emotion_state.context = "critical_battery"
            self.emotion_state.triggered_by = "battery_monitor"

        elif self.battery_level < battery_warning:
            # Warning: mildly concerned
            if self.emotion_state.primary_emotion == EmotionType.NEUTRAL:
                self.emotion_state = self.transitioner.initiate_transition(
                    self.emotion_state,
                    EmotionType.CONCERNED,
                    duration=0.8,
                )
            self.emotion_state.intensity = max(self.emotion_state.intensity, 0.5)
            self.emotion_state.context = "low_battery"

    def balance_callback(self, msg: Bool):
        """React to balance/traction issues."""
        self.balance_stable = msg.data

        if not self.balance_stable:
            # Balance concern
            self.emotion_state = self.transitioner.initiate_transition(
                self.emotion_state,
                EmotionType.CONCERNED,
                duration=0.4,
            )
            self.emotion_state.intensity = 0.8
            self.emotion_state.context = "balance_unstable"
            self.emotion_state.triggered_by = "balance_monitor"

    def health_callback(self, msg: String):
        """React to system health status."""
        self.system_health = msg.data

        if msg.data == "degraded":
            self.emotion_state = self.transitioner.initiate_transition(
                self.emotion_state,
                EmotionType.CONCERNED,
                duration=0.7,
            )
            self.emotion_state.intensity = 0.6
            self.emotion_state.context = "system_degraded"

    def update_callback(self):
        """Main update loop."""
        current_time = time.time()
        dt = 1.0 / self.get_parameter("update_rate_hz").value

        # Update transitions
        self.emotion_state = self.transitioner.apply_transition(self.emotion_state)

        # Update idle behaviors
        self.emotion_state.idle_flags = self.idle_manager.update(dt)

        # Cooldown voice commands
        self.voice_command_cooldown = max(0.0, self.voice_command_cooldown - dt)

        # Idle detection: return to neutral if no interaction for 10+ seconds
        self.idle_timer += dt
        if (
            self.voice_command_cooldown <= 0
            and not self.people_present
            and self.idle_timer > 10.0
            and self.emotion_state.primary_emotion != EmotionType.NEUTRAL
        ):
            self.emotion_state = self.transitioner.initiate_transition(
                self.emotion_state,
                EmotionType.NEUTRAL,
                duration=1.2,
            )
            self.emotion_state.context = "idle"
            self.idle_timer = 0.0

        if self.voice_command_cooldown > 0:
            self.idle_timer = 0.0  # Reset idle timer on activity

        # Apply social memory warmth modifier
        if self.emotion_state.social_target_id and self.emotion_state.primary_emotion == EmotionType.HAPPY:
            warmth = self.social_memory.get_warmth_modifier(self.emotion_state.social_target_id)
            self.emotion_state.social_warmth = warmth
            self.emotion_state.intensity = self.emotion_state.intensity * (0.8 + warmth * 0.2)

        # Update last timestamp
        self.emotion_state.last_update_time = current_time

        # Map emotion to expression name
        self.emotion_state.expression_name = self._map_emotion_to_expression()

        # Publish emotion state
        self._publish_emotion_state()

    def _map_emotion_to_expression(self) -> str:
        """Map internal emotion state to face expression name."""
        emotion = self.emotion_state.primary_emotion
        intensity = self.emotion_state.intensity

        # Intensity-based modulation
        intensity_suffix = ""
        if intensity > 0.7:
            intensity_suffix = "_intense"
        elif intensity < 0.3:
            intensity_suffix = "_subtle"

        base_mapping = {
            EmotionType.NEUTRAL: "neutral",
            EmotionType.HAPPY: "happy",
            EmotionType.CURIOUS: "curious",
            EmotionType.EXCITED: "surprised",  # Use surprised for excitement
            EmotionType.CONCERNED: "sad",  # Concern maps to sad expression
            EmotionType.CONFUSED: "confused",
            EmotionType.TIRED: "sad",
            EmotionType.PLAYFUL: "happy",
        }

        base = base_mapping.get(emotion, "neutral")
        return base + intensity_suffix

    def _publish_emotion_state(self) -> None:
        """Publish current emotion state as structured JSON string."""
        import json

        state_dict = {
            "emotion": self.emotion_state.primary_emotion.value,
            "intensity": float(self.emotion_state.intensity),
            "confidence": float(self.emotion_state.confidence),
            "expression": self.emotion_state.expression_name,
            "context": self.emotion_state.context,
            "triggered_by": self.emotion_state.triggered_by,
            "social_target_id": self.emotion_state.social_target_id,
            "social_warmth": float(self.emotion_state.social_warmth),
            "idle_flags": self.emotion_state.idle_flags,
            "timestamp": self.emotion_state.last_update_time,
            "battery_level": float(self.battery_level),
            "balance_stable": self.balance_stable,
            "system_health": self.system_health,
        }

        msg = String()
        msg.data = json.dumps(state_dict)
        self.emotion_pub.publish(msg)


def main(args=None):
    rclpy.init(args=args)
    node = EmotionEngineNode()

    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
