#!/usr/bin/env python3
"""NeoPixel LED Controller for SaltyBot.

Subscribes to system state topics (battery, balance, social, emotion, health)
and renders animated LED patterns with smooth transitions. Supports configurable
LED count, GPIO pin, and auto-dim brightness.

Published topics:
  /saltybot/led_state (std_msgs/String) - JSON LED status

Subscribed topics:
  /saltybot/battery/level (std_msgs/String) - Battery percentage
  /saltybot/balance/status (std_msgs/String) - Balance state
  /saltybot/social/state (std_msgs/String) - Social state
  /saltybot/emotion/state (std_msgs/String) - Emotion state
  /saltybot/health/status (std_msgs/String) - Health status
"""

import json
import math
import time
from dataclasses import dataclass, asdict
from typing import Dict, List, Optional, Tuple
from pathlib import Path
from enum import Enum

import yaml
import rclpy
from rclpy.node import Node
from rclpy.timer import Timer
from std_msgs.msg import String


class LEDState(Enum):
    """LED display states ordered by priority."""
    ERROR = 100
    LOW_BATTERY = 50
    CELEBRATE = 40
    FOLLOW = 30
    IDLE = 10


@dataclass
class LEDColor:
    """RGB color representation."""
    r: int
    g: int
    b: int

    def __post_init__(self):
        """Clamp values to 0-255."""
        self.r = max(0, min(255, self.r))
        self.g = max(0, min(255, self.g))
        self.b = max(0, min(255, self.b))

    def to_tuple(self) -> Tuple[int, int, int]:
        """Return as RGB tuple."""
        return (self.r, self.g, self.b)

    def scale(self, factor: float) -> "LEDColor":
        """Scale color by factor (0.0-1.0)."""
        return LEDColor(
            int(self.r * factor),
            int(self.g * factor),
            int(self.b * factor),
        )


class LEDPattern:
    """Base class for LED animation patterns."""

    def __init__(self, color: Optional[LEDColor], num_leds: int, speed: float = 1.0):
        """Initialize pattern.

        Args:
            color: RGB color or None for rainbow
            num_leds: Total number of LEDs
            speed: Animation speed multiplier
        """
        self.color = color
        self.num_leds = num_leds
        self.speed = speed
        self.time_offset = 0.0

    def update(self, dt: float, brightness: float) -> List[Tuple[int, int, int]]:
        """Generate LED frame.

        Args:
            dt: Delta time since last frame
            brightness: Brightness multiplier (0.0-1.0)

        Returns:
            List of RGB tuples for each LED
        """
        self.time_offset += dt * self.speed
        return self._generate_frame(brightness)

    def _generate_frame(self, brightness: float) -> List[Tuple[int, int, int]]:
        """Generate frame. Subclasses should override."""
        return [(0, 0, 0)] * self.num_leds


class BreathingPattern(LEDPattern):
    """Smooth breathing animation."""

    def _generate_frame(self, brightness: float) -> List[Tuple[int, int, int]]:
        """Breathing pulse based on sine wave."""
        # Sine wave from 0.3 to 1.0 for breathing effect
        breath_factor = 0.65 + 0.35 * math.sin(2 * math.pi * self.time_offset)
        scaled_color = self.color.scale(breath_factor * brightness)
        return [scaled_color.to_tuple()] * self.num_leds


class PulsePattern(LEDPattern):
    """Quick pulse animation."""

    def _generate_frame(self, brightness: float) -> List[Tuple[int, int, int]]:
        """Pulse with steeper rise/fall."""
        # Triangular wave for snappier pulse
        cycle = (self.time_offset % 1.0)  # 0 to 1
        pulse_factor = 1.0 - abs(2.0 * (cycle - 0.5)) if cycle < 1.0 else 0
        pulse_factor = 0.2 + 0.8 * pulse_factor
        scaled_color = self.color.scale(pulse_factor * brightness)
        return [scaled_color.to_tuple()] * self.num_leds


class FlashPattern(LEDPattern):
    """Fast on/off flashing."""

    def _generate_frame(self, brightness: float) -> List[Tuple[int, int, int]]:
        """Flash pattern: on/off."""
        cycle = self.time_offset % 1.0
        flash_on = cycle < 0.5
        factor = brightness if flash_on else 0.0
        scaled_color = self.color.scale(factor)
        return [scaled_color.to_tuple()] * self.num_leds


class RainbowPattern(LEDPattern):
    """Animated rainbow spectrum."""

    def _generate_frame(self, brightness: float) -> List[Tuple[int, int, int]]:
        """Rainbow spectrum with animation."""
        colors = []
        for i in range(self.num_leds):
            # Hue cycles across strip with animation
            hue = ((i / self.num_leds) + self.time_offset) % 1.0
            rgb = self._hsv_to_rgb(hue, 1.0, brightness)
            colors.append(rgb)
        return colors

    @staticmethod
    def _hsv_to_rgb(h: float, s: float, v: float) -> Tuple[int, int, int]:
        """Convert HSV to RGB."""
        c = v * s
        x = c * (1 - abs((h * 6) % 2 - 1))
        m = v - c

        if h < 1/6:
            r, g, b = c, x, 0
        elif h < 2/6:
            r, g, b = x, c, 0
        elif h < 3/6:
            r, g, b = 0, c, x
        elif h < 4/6:
            r, g, b = 0, x, c
        elif h < 5/6:
            r, g, b = x, 0, c
        else:
            r, g, b = c, 0, x

        return (
            int((r + m) * 255),
            int((g + m) * 255),
            int((b + m) * 255),
        )


class ChasePattern(LEDPattern):
    """Moving light chase across strip."""

    def _generate_frame(self, brightness: float) -> List[Tuple[int, int, int]]:
        """Chasing light pattern."""
        colors = [(0, 0, 0)] * self.num_leds
        pos = int((self.time_offset * self.num_leds) % self.num_leds)

        # Bright center with trailing fade
        for offset in range(-3, 4):
            idx = (pos + offset) % self.num_leds
            if offset == 0:
                factor = brightness
            else:
                factor = brightness * max(0, 1.0 - abs(offset) / 3.0)
            scaled = self.color.scale(factor)
            colors[idx] = scaled.to_tuple()

        return colors


class LEDControllerNode(Node):
    """ROS2 node for LED control."""

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

        # Load config
        self.declare_parameter("config_file", "led_config.yaml")
        config_path = self.get_parameter("config_file").value
        self.config = self._load_config(config_path)

        # Hardware
        self.num_leds = self.config["hardware"]["num_leds"]
        self.gpio_pin = self.config["hardware"]["gpio_pin"]
        self.brightness_max = self.config["hardware"]["brightness_max"]
        self.brightness_min = self.config["hardware"]["brightness_min"]
        self.update_freq = self.config["hardware"]["update_frequency"]

        # State tracking
        self.battery_level = 100
        self.balance_state = "idle"
        self.social_state = "idle"
        self.emotion_state = "idle"
        self.health_state = "ok"
        self.last_update = time.time()
        self.last_state = None
        self.current_pattern: Optional[LEDPattern] = None
        self.transition_time = 0.0

        # Subscribers
        self.create_subscription(String, "/saltybot/battery/level", self._battery_cb, 10)
        self.create_subscription(String, "/saltybot/balance/status", self._balance_cb, 10)
        self.create_subscription(String, "/saltybot/social/state", self._social_cb, 10)
        self.create_subscription(String, "/saltybot/emotion/state", self._emotion_cb, 10)
        self.create_subscription(String, "/saltybot/health/status", self._health_cb, 10)

        # Publisher
        self.led_state_pub = self.create_publisher(String, "/saltybot/led_state", 10)

        # Timer
        period = 1.0 / self.update_freq
        self.timer = self.create_timer(period, self._update_leds)

        self.get_logger().info(
            f"LED Controller initialized: {self.num_leds} LEDs on GPIO{self.gpio_pin}"
        )

    def _load_config(self, config_path: str) -> Dict:
        """Load YAML configuration."""
        try:
            with open(config_path) as f:
                return yaml.safe_load(f)
        except FileNotFoundError:
            self.get_logger().warn(
                f"Config not found: {config_path}, using defaults"
            )
            return self._default_config()

    def _default_config(self) -> Dict:
        """Return default configuration."""
        return {
            "hardware": {
                "num_leds": 30,
                "gpio_pin": 18,
                "brightness_max": 255,
                "brightness_min": 10,
                "update_frequency": 30,
            },
            "animation": {
                "transition_duration": 0.5,
                "breathing_speed": 1.5,
                "pulse_speed": 2.0,
                "chase_speed": 3.0,
            },
            "patterns": {
                "idle": {"color": [0, 50, 200], "pattern": "breathing", "speed": 1.5},
                "follow": {"color": [0, 200, 50], "pattern": "pulse", "speed": 2.0},
                "error": {"color": [255, 0, 0], "pattern": "flash", "speed": 3.0},
                "celebrate": {"color": None, "pattern": "rainbow", "speed": 2.0},
                "low_battery": {"color": [255, 165, 0], "pattern": "pulse", "speed": 1.5},
                "search": {"color": [255, 255, 255], "pattern": "chase", "speed": 3.0},
            },
            "thresholds": {"low_battery_percent": 20},
        }

    def _battery_cb(self, msg: String):
        """Handle battery level."""
        try:
            data = json.loads(msg.data)
            self.battery_level = data.get("level", 100)
        except (json.JSONDecodeError, KeyError):
            pass

    def _balance_cb(self, msg: String):
        """Handle balance state."""
        try:
            data = json.loads(msg.data)
            self.balance_state = data.get("state", "idle")
        except json.JSONDecodeError:
            pass

    def _social_cb(self, msg: String):
        """Handle social state."""
        try:
            data = json.loads(msg.data)
            self.social_state = data.get("state", "idle")
        except json.JSONDecodeError:
            pass

    def _emotion_cb(self, msg: String):
        """Handle emotion state."""
        try:
            data = json.loads(msg.data)
            self.emotion_state = data.get("state", "idle")
        except json.JSONDecodeError:
            pass

    def _health_cb(self, msg: String):
        """Handle health status."""
        try:
            data = json.loads(msg.data)
            self.health_state = data.get("status", "ok")
        except json.JSONDecodeError:
            pass

    def _determine_state(self) -> str:
        """Determine highest priority LED state."""
        if self.health_state == "error":
            return "error"

        if self.battery_level < self.config["thresholds"]["low_battery_percent"]:
            return "low_battery"

        if self.emotion_state == "celebrate":
            return "celebrate"

        if self.balance_state == "follow":
            return "follow"

        return "idle"

    def _get_brightness(self) -> float:
        """Calculate brightness (0.0-1.0)."""
        max_brightness = self.brightness_max
        min_brightness = self.brightness_min

        if self.config["brightness"]["auto_dim"]:
            # Simplified: use night brightness for now
            brightness = self.config["brightness"]["night_brightness"]
        else:
            brightness = max_brightness

        return brightness / 255.0

    def _create_pattern(self, state: str) -> LEDPattern:
        """Create pattern for state."""
        pattern_config = self.config["patterns"].get(state, self.config["patterns"]["idle"])
        color_list = pattern_config.get("color")
        pattern_type = pattern_config.get("pattern", "breathing")
        speed = pattern_config.get("speed", 1.0)

        color = None
        if color_list:
            color = LEDColor(color_list[0], color_list[1], color_list[2])

        if pattern_type == "breathing":
            return BreathingPattern(color, self.num_leds, speed)
        elif pattern_type == "pulse":
            return PulsePattern(color, self.num_leds, speed)
        elif pattern_type == "flash":
            return FlashPattern(color, self.num_leds, speed)
        elif pattern_type == "rainbow":
            return RainbowPattern(color, self.num_leds, speed)
        elif pattern_type == "chase":
            return ChasePattern(color, self.num_leds, speed)
        else:
            return BreathingPattern(color, self.num_leds, speed)

    def _update_leds(self):
        """Update LED display."""
        now = time.time()
        dt = now - self.last_update
        self.last_update = now

        # Determine state
        state = self._determine_state()

        # Transition to new state
        if state != self.last_state:
            self.last_state = state
            self.current_pattern = self._create_pattern(state)
            self.transition_time = 0.0
            self.get_logger().info(f"LED state changed to: {state}")

        # Update pattern
        brightness = self._get_brightness()
        if self.current_pattern:
            leds = self.current_pattern.update(dt, brightness)
        else:
            leds = [(0, 0, 0)] * self.num_leds

        # Publish state
        state_msg = {
            "state": self.last_state or "idle",
            "battery": self.battery_level,
            "brightness": int(brightness * 255),
            "num_leds": self.num_leds,
            "gpio_pin": self.gpio_pin,
        }
        self.led_state_pub.publish(String(data=json.dumps(state_msg)))


def main(args=None):
    """Main entry point."""
    rclpy.init(args=args)
    node = LEDControllerNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()
