"""Serial runner — real hardware over USB/serial (ESP32).

Implements the BaseRunner interface for a single physical robot.
All physics come from the real world; the runner translates between
the ESP32 serial protocol and the qpos/qvel tensors that BaseRunner
and BaseEnv expect.

Serial protocol (ESP32 firmware):
    Commands sent TO the ESP32:
        G           — start streaming state lines
        H           — stop streaming
        M<int>      — set motor PWM speed (-255 … 255)

    State lines received FROM the ESP32 (firmware sends SI units):
        S,<ms>,<motor_rad>,<motor_vel>,<pend_rad>,<pend_vel>,<motor_speed>
        (7 comma-separated fields after the ``S`` prefix)

        motor_rad    — motor joint angle (radians)
        motor_vel    — motor joint velocity (rad/s)
        pend_rad     — pendulum angle (radians, 0 = hanging down)
        pend_vel     — pendulum angular velocity (rad/s)
        motor_speed  — applied PWM (-255..255, for action recording)

A daemon thread continuously reads the serial stream so the control
loop never blocks on I/O.

Usage:
    python train.py env=rotary_cartpole runner=serial training=ppo_real
"""

from __future__ import annotations

import dataclasses
import logging
import threading
import time
from typing import Any

import numpy as np
import torch

from src.core.runner import BaseRunner, BaseRunnerConfig

logger = logging.getLogger(__name__)


@dataclasses.dataclass
class SerialRunnerConfig(BaseRunnerConfig):
    """Configuration for serial communication with the ESP32."""

    num_envs: int = 1  # always 1 — single physical robot
    device: str = "cpu"

    port: str = "/dev/cu.usbserial-0001"
    baud: int = 115200
    dt: float = 0.04  # control loop period (seconds), 25 Hz
    no_data_timeout: float = 2.0  # seconds of silence → disconnect

    # Physical reset procedure
    reset_drive_speed: int = 80       # PWM for bang-bang drive-to-center
    reset_deadband_rad: float = 0.01  # "centered" threshold (~0.6°)
    reset_drive_timeout: float = 3.0  # seconds to reach center
    reset_settle_timeout: float = 30.0  # seconds to wait for pendulum


class SerialRunner(BaseRunner[SerialRunnerConfig]):
    """BaseRunner implementation that talks to real hardware over serial.

    Maps the ESP32 serial protocol to qpos/qvel tensors so the existing
    RotaryCartPoleEnv (or any compatible env) works unchanged.

    qpos layout: [motor_angle_rad, pendulum_angle_rad]
    qvel layout: [motor_vel_rad_s, pendulum_vel_rad_s]
    """

    # ------------------------------------------------------------------
    # BaseRunner interface
    # ------------------------------------------------------------------

    @property
    def num_envs(self) -> int:
        return 1

    @property
    def device(self) -> torch.device:
        return torch.device("cpu")

    def _sim_initialize(self, config: SerialRunnerConfig) -> None:
        # Joint dimensions for the rotary cartpole (motor + pendulum).
        self._nq = 2
        self._nv = 2

        # Import serial here so it's not a hard dependency for sim-only users.
        import serial as _serial

        self._serial_mod = _serial

        # Explicitly disable hardware flow control and exclusive mode to
        # avoid termios.error (errno 22) on macOS with CH340/CP2102 adapters.
        self.ser: _serial.Serial = _serial.Serial(
            port=config.port,
            baudrate=config.baud,
            timeout=0.05,
            xonxoff=False,
            rtscts=False,
            dsrdtr=False,
            exclusive=False,
        )
        time.sleep(2)  # Wait for ESP32 boot.
        self.ser.reset_input_buffer()

        # Internal state tracking.
        self._rebooted: bool = False
        self._serial_disconnected: bool = False
        self._last_esp_ms: int = 0
        self._last_data_time: float = time.monotonic()
        self._streaming: bool = False

        # Latest parsed state (updated by the reader thread).
        # Firmware sends SI units — values are used directly as qpos/qvel.
        self._latest_state: dict[str, Any] = {
            "timestamp_ms": 0,
            "motor_rad": 0.0,
            "motor_vel": 0.0,
            "pend_rad": 0.0,
            "pend_vel": 0.0,
            "motor_speed": 0,
        }
        self._state_lock = threading.Lock()
        self._state_event = threading.Event()

        # Start background serial reader.
        self._reader_running = True
        self._reader_thread = threading.Thread(
            target=self._serial_reader, daemon=True
        )
        self._reader_thread.start()

        # Start streaming.
        self._send("G")
        self._streaming = True
        self._last_data_time = time.monotonic()

        # Derive max PWM from actuator ctrl_range so the serial
        # command range matches what MuJoCo training sees.
        ctrl_hi = self.env.robot.actuators[0].ctrl_range[1]
        self._max_pwm: int = round(ctrl_hi * 255)

        # Track wall-clock time of last step for PPO-gap detection.
        self._last_step_time: float = time.monotonic()

    def _sim_step(
        self, actions: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        now = time.monotonic()

        # Detect PPO update gap: if more than 0.5s since last step,
        # the optimizer was running and no motor commands were sent.
        # Trigger a full reset so the robot starts from a clean state.
        gap = now - self._last_step_time
        if gap > 0.5:
            logger.info(
                "PPO update gap detected (%.1f s) — resetting before resuming.",
                gap,
            )
            self._send("M0")
            all_ids = torch.arange(self.num_envs, device=self.device)
            self._sim_reset(all_ids)
            self.step_counts.zero_()

        # Map normalised action [-1, 1] → PWM, scaled by ctrl_range.
        action_val = float(actions[0, 0].clamp(-1.0, 1.0))
        motor_speed = int(action_val * self._max_pwm)
        self._send(f"M{motor_speed}")

        # Stream-driven: block until the firmware sends the next state
        # line (~20 ms at 50 Hz).
        state = self._read_state_blocking(timeout=0.1)

        # Firmware sends SI units — use directly.
        qpos, qvel = self._state_to_tensors(state)

        # Cache for _sync_viz().
        self._last_sync_state = state
        self._last_step_time = time.monotonic()

        return qpos, qvel

    def _sim_reset(
        self, env_ids: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # If ESP32 rebooted or disconnected, we can't recover.
        if self._rebooted or self._serial_disconnected:
            raise RuntimeError(
                "ESP32 rebooted or disconnected during training! "
                "Encoder center is lost. "
                "Please re-center the motor manually and restart."
            )

        # Stop motor and restart streaming.
        self._send("M0")
        self._send("H")
        self._streaming = False
        time.sleep(0.05)
        self._state_event.clear()
        self._send("G")
        self._streaming = True
        self._last_data_time = time.monotonic()
        time.sleep(0.05)

        # Physically return the motor to the centre position.
        self._drive_to_center()

        # Wait until the env considers the robot settled.
        self._wait_for_settle()

        # Refresh data timer so health checks don't false-positive.
        self._last_data_time = time.monotonic()

        # Read settled state and return as qpos/qvel.
        state = self._read_state_blocking()
        qpos, qvel = self._state_to_tensors(state)

        self._last_sync_state = state
        return qpos, qvel

    def _sim_close(self) -> None:
        self._reader_running = False
        self._streaming = False
        self._send("H")
        self._send("M0")
        time.sleep(0.1)
        self._reader_thread.join(timeout=1.0)
        self.ser.close()
        super()._sim_close()

    # ------------------------------------------------------------------
    # MuJoCo digital-twin rendering
    # ------------------------------------------------------------------

    def _ensure_viz_model(self) -> None:
        """Lazily load the MuJoCo model for visualisation (digital twin)."""
        if hasattr(self, "_viz_model"):
            return

        import mujoco
        from src.runners.mujoco import load_mujoco_model

        self._viz_model = load_mujoco_model(self.env.robot)
        self._viz_data = mujoco.MjData(self._viz_model)
        self._offscreen_renderer = None

    def _sync_viz(self) -> None:
        """Copy current serial sensor state into the MuJoCo viz model."""
        import mujoco

        self._ensure_viz_model()

        last_state = getattr(self, "_last_sync_state", None)
        if last_state is None:
            last_state = self._read_state()

        # Firmware sends radians — use directly.
        self._viz_data.qpos[0] = last_state["motor_rad"]
        self._viz_data.qpos[1] = last_state["pend_rad"]
        self._viz_data.qvel[0] = last_state["motor_vel"]
        self._viz_data.qvel[1] = last_state["pend_vel"]

        mujoco.mj_forward(self._viz_model, self._viz_data)

    def _render_frame(self, env_idx: int = 0) -> np.ndarray:
        """Offscreen render of the digital-twin MuJoCo model."""
        import mujoco

        self._sync_viz()

        if self._offscreen_renderer is None:
            self._offscreen_renderer = mujoco.Renderer(
                self._viz_model, width=640, height=480,
            )
        self._offscreen_renderer.update_scene(self._viz_data)
        return self._offscreen_renderer.render().copy()

    # ------------------------------------------------------------------
    # Override step() for runner-level safety
    # ------------------------------------------------------------------

    def step(
        self, actions: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]:
        # Check for ESP32 reboot / disconnect BEFORE stepping.
        if self._rebooted or self._serial_disconnected:
            self._send("M0")
            qpos, qvel = self._make_current_state()
            state = self.env.build_state(qpos, qvel)
            obs = self.env.compute_observations(state)
            reward = torch.tensor([[-100.0]])
            terminated = torch.tensor([[True]])
            truncated = torch.tensor([[False]])
            return obs, reward, terminated, truncated, {"reboot_detected": True}

        # Normal step via BaseRunner (calls _sim_step → env logic).
        obs, rewards, terminated, truncated, info = super().step(actions)

        # Check connection health after stepping.
        if not self._check_connection_health():
            self._send("M0")
            terminated = torch.tensor([[True]])
            rewards = torch.tensor([[-100.0]])
            info["reboot_detected"] = True

        # Always stop motor on episode end.
        if terminated.any() or truncated.any():
            self._send("M0")

        return obs, rewards, terminated, truncated, info

    # ------------------------------------------------------------------
    # Serial helpers
    # ------------------------------------------------------------------

    def _send(self, cmd: str) -> None:
        """Send a command to the ESP32."""
        try:
            self.ser.write(f"{cmd}\n".encode())
        except (OSError, self._serial_mod.SerialException):
            self._serial_disconnected = True

    def _serial_reader(self) -> None:
        """Background thread: continuously read and parse serial lines.

        Protocol: ``S,<ms>,<motor_rad>,<motor_vel>,<pend_rad>,<pend_vel>,<motor_speed>``
        (7 comma-separated fields).  Firmware sends SI units directly.
        """
        while self._reader_running:
            try:
                if self.ser.in_waiting:
                    line = (
                        self.ser.readline()
                        .decode("utf-8", errors="ignore")
                        .strip()
                    )

                    # Detect ESP32 reboot: it prints READY on startup.
                    if line.startswith("READY"):
                        self._rebooted = True
                        logger.critical("ESP32 reboot detected: %s", line)
                        continue

                    if line.startswith("S,"):
                        parts = line.split(",")
                        if len(parts) >= 7:
                            try:
                                esp_ms = int(parts[1])
                            except ValueError:
                                logger.debug(
                                    "Malformed state line (header): %s",
                                    line,
                                )
                                continue

                            # Detect reboot: timestamp jumped backwards.
                            if (
                                self._last_esp_ms > 5000
                                and esp_ms < self._last_esp_ms - 3000
                            ):
                                self._rebooted = True
                                logger.critical(
                                    "ESP32 reboot detected: timestamp"
                                    " %d -> %d",
                                    self._last_esp_ms,
                                    esp_ms,
                                )

                            self._last_esp_ms = esp_ms
                            self._last_data_time = time.monotonic()

                            try:
                                parsed: dict[str, Any] = {
                                    "timestamp_ms": esp_ms,
                                    "motor_rad": float(parts[2]),
                                    "motor_vel": float(parts[3]),
                                    "pend_rad": float(parts[4]),
                                    "pend_vel": float(parts[5]),
                                    "motor_speed": int(parts[6]),
                                }
                            except ValueError:
                                logger.debug(
                                    "Malformed state line (fields): %s",
                                    line,
                                )
                                continue
                            with self._state_lock:
                                self._latest_state = parsed
                            self._state_event.set()
                else:
                    time.sleep(0.001)  # Avoid busy-spinning.
            except (OSError, self._serial_mod.SerialException) as exc:
                self._serial_disconnected = True
                logger.critical("Serial connection lost: %s", exc)
                break

    def _check_connection_health(self) -> bool:
        """Return True if the ESP32 connection appears healthy."""
        if self._serial_disconnected:
            logger.critical("ESP32 serial connection lost.")
            return False
        if (
            self._streaming
            and (time.monotonic() - self._last_data_time)
            > self.config.no_data_timeout
        ):
            logger.critical(
                "No data from ESP32 for %.1f s — possible crash/disconnect.",
                time.monotonic() - self._last_data_time,
            )
            self._rebooted = True
            return False
        return True

    def _read_state(self) -> dict[str, Any]:
        """Return the most recent state from the reader thread (non-blocking)."""
        with self._state_lock:
            return dict(self._latest_state)

    def _read_state_blocking(self, timeout: float = 0.05) -> dict[str, Any]:
        """Wait for a fresh sample, then return it."""
        self._state_event.clear()
        self._state_event.wait(timeout=timeout)
        with self._state_lock:
            return dict(self._latest_state)

    def _state_to_tensors(
        self, state: dict[str, Any],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Convert a parsed state dict to (qpos, qvel) tensors."""
        qpos = torch.tensor(
            [[state["motor_rad"], state["pend_rad"]]], dtype=torch.float32
        )
        qvel = torch.tensor(
            [[state["motor_vel"], state["pend_vel"]]], dtype=torch.float32
        )
        return qpos, qvel

    def _make_current_state(self) -> tuple[torch.Tensor, torch.Tensor]:
        """Build qpos/qvel from current sensor data (utility)."""
        return self._state_to_tensors(self._read_state_blocking())

    # ------------------------------------------------------------------
    # Physical reset helpers
    # ------------------------------------------------------------------

    def _drive_to_center(self) -> None:
        """Drive the motor back toward center using bang-bang control."""
        cfg = self.config
        start = time.time()
        while time.time() - start < cfg.reset_drive_timeout:
            state = self._read_state_blocking()
            motor_rad = state["motor_rad"]
            if abs(motor_rad) < cfg.reset_deadband_rad:
                break
            speed = cfg.reset_drive_speed if motor_rad < 0 else -cfg.reset_drive_speed
            self._send(f"M{speed}")
            time.sleep(0.05)
        self._send("M0")
        time.sleep(0.2)

    def _wait_for_settle(self) -> None:
        """Block until the env considers the robot ready for a new episode."""
        cfg = self.config
        stable_since: float | None = None
        start = time.monotonic()

        while time.monotonic() - start < cfg.reset_settle_timeout:
            state = self._read_state_blocking()
            qpos, qvel = self._state_to_tensors(state)

            if self.env.is_reset_ready(qpos, qvel):
                if stable_since is None:
                    stable_since = time.monotonic()
                elif time.monotonic() - stable_since >= 0.5:
                    logger.info(
                        "Robot settled after %.2f s",
                        time.monotonic() - start,
                    )
                    return
            else:
                stable_since = None
            time.sleep(0.02)

        logger.warning(
            "Robot did not settle within %.1f s — proceeding anyway.",
            cfg.reset_settle_timeout,
        )
