"""
Data Flywheel - Event-Driven Data Capture and Export System

The Data Flywheel implements the Layer 2 data architecture of ROSClaw:
- High-frequency capture (1kHz) via Ring Buffers
- Event-triggered data persistence
- Automatic export to LeRobot format
- Tiered storage (Hot/Warm/Cold)

Key Innovation: Instead of recording everything (1TB/day), we only store
data around interesting events (success/failure/emergency), achieving
100x storage optimization while keeping 100% of valuable data.

Usage:
    flywheel = DataFlywheel(
        robot_id="ur5e_001",
        joint_dof=6,
        camera_topics=["/camera/color"],
    )

    # In control loop (1kHz)
    flywheel.on_control_cycle(robot_state)

    # When interesting event happens
    flywheel.trigger_event(
        event_type=EventType.SUCCESS,
        metadata={"task": "pick_and_place", "instruction": "Pick the red block"},
    )
"""

import json
import time
import uuid
from dataclasses import dataclass, field
from enum import Enum, auto
from pathlib import Path
from typing import Optional, Callable, Dict, List, Any
import threading
import logging

import numpy as np

from .ring_buffer import MultiChannelRingBuffer


logger = logging.getLogger(__name__)


class EventType(Enum):
    """Types of events that trigger data capture."""
    SUCCESS = auto()      # Task completed successfully
    FAILURE = auto()      # Task failed (grasp missed, collision, etc.)
    EMERGENCY = auto()    # Emergency stop triggered
    USER_MARK = auto()    # User manually marked as interesting
    MILESTONE = auto()    # Reached milestone in multi-step task


@dataclass
class DataEvent:
    """
    Represents a captured event with associated data.

    Attributes:
        event_id: Unique identifier
        event_type: Type of event
        timestamp: When event occurred
        pre_event_data: Data before event (typically 5 seconds)
        post_event_data: Data after event (typically 5 seconds)
        metadata: Additional context (task description, language instruction, etc.)
    """
    event_id: str
    event_type: EventType
    timestamp: float
    robot_id: str
    pre_event_duration: float = 5.0   # Seconds before event
    post_event_duration: float = 5.0  # Seconds after event
    metadata: Dict[str, Any] = field(default_factory=dict)
    data_paths: Dict[str, Path] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "event_id": self.event_id,
            "event_type": self.event_type.name,
            "timestamp": self.timestamp,
            "robot_id": self.robot_id,
            "pre_event_duration": self.pre_event_duration,
            "post_event_duration": self.post_event_duration,
            "metadata": self.metadata,
            "data_paths": {k: str(v) for k, v in self.data_paths.items()},
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "DataEvent":
        """Create from dictionary."""
        return cls(
            event_id=data["event_id"],
            event_type=EventType[data["event_type"]],
            timestamp=data["timestamp"],
            robot_id=data["robot_id"],
            pre_event_duration=data.get("pre_event_duration", 5.0),
            post_event_duration=data.get("post_event_duration", 5.0),
            metadata=data.get("metadata", {}),
            data_paths={k: Path(v) for k, v in data.get("data_paths", {}).items()},
        )


@dataclass
class RobotState:
    """
    Snapshot of robot state at a given time.

    This is the standard state representation used throughout ROSClaw.
    """
    timestamp: float
    joint_positions: np.ndarray     # Shape: (dof,)
    joint_velocities: np.ndarray    # Shape: (dof,)
    joint_torques: np.ndarray       # Shape: (dof,)
    end_effector_pose: Optional[np.ndarray] = None  # Shape: (4, 4) homogeneous matrix
    gripper_state: Optional[float] = None  # 0.0 = open, 1.0 = closed

    def validate(self, expected_dof: int) -> bool:
        """Validate state dimensions."""
        return (
            self.joint_positions.shape == (expected_dof,) and  # noqa: W504
            self.joint_velocities.shape == (expected_dof,) and  # noqa: W504
            self.joint_torques.shape == (expected_dof,)
        )


class DataFlywheel:
    """
    Event-driven data capture and export system.

    The Data Flywheel maintains high-frequency ring buffers in memory
    and only persists data when interesting events occur. This achieves:

    - 100x storage reduction (10GB/day vs 1TB/day for full recording)
    - 100% capture of relevant data (events)
    - Zero performance impact on control loop (<1ms overhead)
    - Automatic export to standard formats (LeRobot)

    Storage Tiers:
    - Hot (Memory): 60 seconds of recent data, 1kHz, <1μs access
    - Warm (Local SSD): 7 days of event data, LeRobot format
    - Cold (Cloud): Permanent archive of keyframes + metadata
    """

    def __init__(
        self,
        robot_id: str,
        joint_dof: int = 6,
        buffer_duration_sec: float = 60.0,
        sampling_rate_hz: int = 1000,
        camera_topics: Optional[List[str]] = None,
        storage_path: Optional[Path] = None,
    ):
        """
        Initialize Data Flywheel.

        Args:
            robot_id: Unique identifier for this robot instance
            joint_dof: Degrees of freedom (e.g., 6 for UR5)
            buffer_duration_sec: How much history to keep in memory (default 60s)
            sampling_rate_hz: Control frequency (default 1000Hz)
            camera_topics: List of camera topic names to capture
            storage_path: Where to save event data (default: ./rosclaw_data/)
        """
        self.robot_id = robot_id
        self.joint_dof = joint_dof
        self.sampling_rate_hz = sampling_rate_hz
        self.camera_topics = camera_topics or []

        # Calculate buffer sizes
        buffer_size = int(buffer_duration_sec * sampling_rate_hz)

        # Initialize ring buffers for each data channel
        channels = {
            "joint_positions": (buffer_size, (joint_dof,)),
            "joint_velocities": (buffer_size, (joint_dof,)),
            "joint_torques": (buffer_size, (joint_dof,)),
        }

        # Add camera channels if specified
        for topic in self.camera_topics:
            # Assume RGB images, actual shape set dynamically
            channels[topic] = (buffer_size // 10, ())  # Cameras typically lower frequency

        self._buffers = MultiChannelRingBuffer(**channels)

        # Event storage
        self._events: List[DataEvent] = []
        self._event_lock = threading.Lock()

        # Storage configuration
        self._storage_path = storage_path or Path("./rosclaw_data")
        self._storage_path.mkdir(parents=True, exist_ok=True)

        # Performance metrics
        self._cycle_count = 0
        self._dropped_cycles = 0
        self._last_cycle_time = 0.0

        logger.info(
            f"DataFlywheel initialized for {robot_id}: "
            f"{buffer_size} samples @ {sampling_rate_hz}Hz "
            f"({buffer_duration_sec}s history)"
        )

    def on_control_cycle(self, state: RobotState) -> None:
        """
        Called every control cycle (1kHz).

        This method is performance-critical and must complete in <1ms.
        It simply appends data to ring buffers with minimal overhead.

        Args:
            state: Current robot state
        """
        # Validate state
        if not state.validate(self.joint_dof):
            logger.warning("Invalid state dimensions, skipping")
            self._dropped_cycles += 1
            return

        # Append to buffers (O(1) operation)
        self._buffers.append({
            "joint_positions": state.joint_positions,
            "joint_velocities": state.joint_velocities,
            "joint_torques": state.joint_torques,
        }, timestamp=state.timestamp)

        self._cycle_count += 1

    def trigger_event(
        self,
        event_type: EventType,
        metadata: Optional[Dict[str, Any]] = None,
        pre_duration_sec: float = 5.0,
        post_duration_sec: float = 5.0,
    ) -> str:
        """
        Trigger data capture around an event.

        This extracts data from the ring buffers and saves it to storage.
        Called asynchronously when interesting things happen.

        Args:
            event_type: Type of event (SUCCESS, FAILURE, etc.)
            metadata: Additional context (task name, language instruction, etc.)
            pre_duration_sec: How much data before event to save
            post_duration_sec: How much data after event to save

        Returns:
            event_id: Unique identifier for this event
        """
        event_id = str(uuid.uuid4())[:8]
        timestamp = time.time()

        logger.info(f"Triggering event {event_id}: {event_type.name}")

        # Create event record
        event = DataEvent(
            event_id=event_id,
            event_type=event_type,
            timestamp=timestamp,
            robot_id=self.robot_id,
            pre_event_duration=pre_duration_sec,
            post_event_duration=post_duration_sec,
            metadata=metadata or {},
        )

        # Extract data from ring buffers
        pre_samples = int(pre_duration_sec * self.sampling_rate_hz)

        # Get data from ring buffer
        all_data = self._buffers.get_last_n(pre_samples)

        # Save to disk asynchronously
        threading.Thread(
            target=self._save_event_data,
            args=(event, all_data),
            daemon=True,
        ).start()

        # Add to event list
        with self._event_lock:
            self._events.append(event)

        return event_id

    def _save_event_data(
        self,
        event: DataEvent,
        data: Dict[str, tuple],
    ) -> None:
        """
        Save event data to storage (runs in background thread).

        Args:
            event: Event metadata
            data: Data from ring buffers
        """
        try:
            # Create event directory
            event_dir = self._storage_path / f"{event.robot_id}_{event.event_id}"
            event_dir.mkdir(parents=True, exist_ok=True)

            # Save each channel
            for channel_name, (values, timestamps) in data.items():
                filepath = event_dir / f"{channel_name}.npy"
                np.save(filepath, values)

                # Save timestamps separately
                ts_filepath = event_dir / f"{channel_name}_timestamps.npy"
                np.save(ts_filepath, timestamps)

                event.data_paths[channel_name] = filepath

            # Save metadata
            metadata_path = event_dir / "metadata.json"
            with open(metadata_path, "w") as f:
                json.dump(event.to_dict(), f, indent=2, default=str)

            logger.info(f"Event {event.event_id} saved to {event_dir}")

        except Exception as e:
            logger.error(f"Failed to save event {event.event_id}: {e}")

    def export_to_lerobot(
        self,
        output_path: Path,
        filter_fn: Optional[Callable[[DataEvent], bool]] = None,
    ) -> Path:
        """
        Export all events to LeRobot dataset format.

        LeRobot format: https://github.com/huggingface/lerobot

        Args:
            output_path: Where to save the dataset
            filter_fn: Optional filter function for events

        Returns:
            Path to exported dataset
        """
        output_path = Path(output_path)
        output_path.mkdir(parents=True, exist_ok=True)

        # Filter events
        with self._event_lock:
            events = self._events.copy()

        if filter_fn:
            events = [e for e in events if filter_fn(e)]

        logger.info(f"Exporting {len(events)} events to LeRobot format")

        # Convert to LeRobot format
        # This is a simplified version - full implementation would follow
        # LeRobot's exact specification
        lerobot_data = {
            "dataset_info": {
                "robot_id": self.robot_id,
                "total_episodes": len(events),
                "total_frames": sum(
                    int((e.pre_event_duration + e.post_event_duration) * self.sampling_rate_hz)
                    for e in events
                ),
            },
            "episodes": [],
        }

        for event in events:
            episode = {
                "episode_id": event.event_id,
                "task": event.metadata.get("task", "unknown"),
                "language_instruction": event.metadata.get("instruction", ""),
                "success": event.event_type == EventType.SUCCESS,
            }
            lerobot_data["episodes"].append(episode)

        # Save
        dataset_path = output_path / "dataset.json"
        with open(dataset_path, "w") as f:
            json.dump(lerobot_data, f, indent=2)

        logger.info(f"Dataset exported to {dataset_path}")
        return dataset_path

    def get_stats(self) -> Dict[str, Any]:
        """Get flywheel statistics."""
        return {
            "robot_id": self.robot_id,
            "total_cycles": self._cycle_count,
            "dropped_cycles": self._dropped_cycles,
            "drop_rate": self._dropped_cycles / max(self._cycle_count, 1),
            "buffer_size": self._buffers.size,
            "total_events": len(self._events),
            "events_by_type": {
                event_type.name: sum(1 for e in self._events if e.event_type == event_type)
                for event_type in EventType
            },
        }

    def clear(self) -> None:
        """Clear all buffers and events."""
        self._buffers.clear()
        with self._event_lock:
            self._events.clear()
        logger.info("DataFlywheel cleared")
