"""
Dataset recorder for Feather + LeRobot with start/stop hooks.
Designed for teleop integration: use the same ZED instance, pass frames from main loop.
Records actual robot joint positions from FeatherLeRobotAdapter.
"""
from __future__ import annotations

import time
from typing import TYPE_CHECKING, Optional

import numpy as np

from .feather_robot_adapter import ACTION_NAMES, STATE_NAMES

if TYPE_CHECKING:
    from .feather_robot_adapter import FeatherLeRobotAdapter


def _ensure_lerobot():
    import sys
    import os

    cwd = os.getcwd()
    idx = cwd.find("featheros")
    featheros_path = cwd[: idx + len("featheros")] if idx >= 0 else cwd
    sys.path.insert(0, featheros_path)
    sys.path.insert(0, os.path.join(featheros_path, "feathersdk/src"))


class FeatherDatasetRecorder:
    """
    Start/stop dataset recorder for teleop. Uses the same ZED instance as teleop;
    caller passes image frames from the main loop.

    Records actual robot joint positions from adapter.get_observation().
    Requires a FeatherLeRobotAdapter wrapping a real FeatherRobot.

    Usage:
        recorder = FeatherDatasetRecorder(adapter)
        recorder.start_recording(task="Pick object")
        # each frame: recorder.record_frame(left_rgb) when recorder.is_recording
        recorder.stop_recording()
        recorder.finalize()
    """

    def __init__(
        self,
        adapter: "FeatherLeRobotAdapter",
        *,
        repo_id: str = "feather",
        fps: int = 30,
        config_path: Optional[str] = None,
    ):
        _ensure_lerobot()
        self._adapter = adapter
        self._repo_id = repo_id
        self._fps = fps
        self._config_path = config_path
        self._dataset = None
        self._recording = False
        self._current_task: str = ""
        self._last_frame_time = 0.0
        self._min_frame_interval = 1.0 / fps
        self._robot_type = adapter.name

    @property
    def is_recording(self) -> bool:
        return self._recording and self._dataset is not None

    def start_recording(
        self,
        task: str = "Manipulation task",
        repo_id: Optional[str] = None,
        fps: Optional[int] = None,
    ) -> None:
        """Start a new episode. Creates dataset on first call."""
        from lerobot.datasets.lerobot_dataset import LeRobotDataset
        from lerobot.datasets.utils import hw_to_dataset_features
        from lerobot.utils.constants import ACTION, OBS_STR

        repo_id = repo_id or self._repo_id
        fps = fps or self._fps

        if self._dataset is None:
            obs_hw = {n: float for n in STATE_NAMES}
            obs_hw["top"] = (720, 1280, 3)
            action_hw = {n: float for n in ACTION_NAMES}
            features = {
                **hw_to_dataset_features(action_hw, ACTION, use_video=True),
                **hw_to_dataset_features(obs_hw, OBS_STR, use_video=True),
            }
            self._dataset = LeRobotDataset.create(
                repo_id=repo_id,
                fps=fps,
                features=features,
                robot_type=self._robot_type,
                use_videos=True,
                streaming_encoding=True,
                encoder_threads=2,
            )
            self._dataset.start_image_writer(num_threads=4)
            self._min_frame_interval = 1.0 / fps

        self._current_task = task
        self._recording = True
        self._last_frame_time = 0.0

    def record_frame(self, image: np.ndarray) -> None:
        """
        Add a frame if recording and rate-limited to fps.
        image: RGB (H, W, 3) uint8, e.g. (720, 1280, 3).
        State comes from adapter.get_observation() (actual robot joint positions).
        """
        if not self.is_recording:
            return

        now = time.perf_counter()
        if now - self._last_frame_time < self._min_frame_interval:
            return
        self._last_frame_time = now

        from lerobot.datasets.utils import build_dataset_frame
        from lerobot.utils.constants import ACTION, OBS_STR

        obs = self._adapter.get_observation()
        state_vec = obs["observation.state"]

        obs_flat = {STATE_NAMES[i]: state_vec[i] for i in range(len(STATE_NAMES))}
        obs_flat["top"] = np.asarray(image)
        if obs_flat["top"].shape[2] == 4:
            obs_flat["top"] = obs_flat["top"][:, :, :3]
        action_flat = {ACTION_NAMES[i]: state_vec[i] for i in range(len(ACTION_NAMES))}

        frame = {
            **build_dataset_frame(self._dataset.features, obs_flat, prefix=OBS_STR),
            **build_dataset_frame(self._dataset.features, action_flat, prefix=ACTION),
            "task": self._current_task,
        }
        self._dataset.add_frame(frame)

    def stop_recording(self) -> None:
        """Save current episode and stop recording."""
        if self._recording and self._dataset is not None:
            self._dataset.save_episode()
        self._recording = False

    def finalize(self) -> None:
        """Finalize dataset. Call when done recording session."""
        if self._dataset is not None:
            self._dataset.finalize()
            self._dataset = None
        self._recording = False
