"""Hardware Abstraction Layer — bridges the MuJoCo env and the brain."""
from __future__ import annotations

import threading
import time
from collections import deque
from typing import Optional

import numpy as np

from env.humanoid_env import HumanoidTouchEnv, RobotObs
from hal.obs_buffer import RingBuffer


class HAL:
    """
    Runs the MuJoCo control loop and owns the obs ring-buffer.

    The brain pushes action chunks here via push_actions().
    HAL pops one action per step; zeros are used when the buffer drains.

    Usage (headless):
        hal.start()          # launches background control thread
        ...
        hal.stop()

    Usage (with MuJoCo passive viewer — viewer owns the main thread):
        with mujoco.viewer.launch_passive(env.model, env.data) as v:
            while v.is_running():
                with v.lock():
                    hal.step()
                v.sync()
    """

    def __init__(self, env: HumanoidTouchEnv, config: dict):
        self.env = env
        self._cfg = config

        buf_size = config.get("obs_buffer_size", 200)
        self.obs_buffer: RingBuffer[RobotObs] = RingBuffer(buf_size)

        self._chunk_len = config.get("action_chunk_len", 10)
        self._action_queue: deque[np.ndarray] = deque()
        self._action_lock = threading.Lock()

        self._step_count = 0
        self._vision_every = None          # set from robot config in start()
        self._zero_action = np.zeros(env.N_ACTUATORS, dtype=np.float32)
        # Zero-order hold: when the action queue drains we re-issue the last
        # command instead of zeros, so the robot holds its pose rather than
        # lurching toward the all-zero configuration (proven-control behaviour).
        self._last_action = self._zero_action.copy()

        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._target_dt = 1.0 / config.get("control_freq", 50)

        # Offscreen camera rendering is only needed when a consumer (a real
        # world-model / VLA) actually uses the frames. The RX1 IK stub does not,
        # and rendering on the control thread both slows the loop (which made the
        # action queue back up) and can crash the GL context. Off unless enabled.
        self._vision_enabled = bool(
            env.config.get("vision", {}).get("enabled", False))

    # ------------------------------------------------------------------ public

    def start(self):
        """Launch the HAL control loop in a daemon thread."""
        obs = self.env.reset()   # render_vision=False on reset for headless compat
        self.obs_buffer.push(obs)
        self._running = True
        self._thread = threading.Thread(target=self._loop, daemon=True, name="HAL")
        self._thread.start()

    def stop(self, timeout: float = 2.0):
        self._running = False
        if self._thread:
            self._thread.join(timeout=timeout)

    def step(self, vision_every: int = 5) -> RobotObs:
        """
        Advance the simulation by one control step and record the observation.
        Called externally when the viewer owns the main thread.
        Renders primary vision AND multi-view cameras together (same cadence).
        """
        action = self._pop_action()
        render = self._vision_enabled and (self._step_count % vision_every == 0)
        obs = self.env.step(action, render_vision=render, render_multi_view=render)
        self.obs_buffer.push(obs)
        self._step_count += 1
        return obs

    def push_actions(self, actions: np.ndarray):
        """
        Load a fresh chunk of actions from the policy (latest-wins).

        The queue is REPLACED, not appended to: the brain re-plans from the
        current state every think cycle, so any unconsumed actions from the
        previous chunk are stale. Appending them made the queue grow without
        bound (the control loop ran slightly slower than the brain pushed),
        so the robot executed commands many seconds old. Replacing keeps
        latency bounded to a single chunk and the robot always tracks the
        freshest target.
        """
        with self._action_lock:
            self._action_queue.clear()
            for a in actions:
                self._action_queue.append(a.astype(np.float32))

    def get_recent_obs(self, n: int = 8) -> list[RobotObs]:
        return self.obs_buffer.get_recent(n)

    def latest_obs(self) -> Optional[RobotObs]:
        return self.obs_buffer.latest()

    @property
    def step_count(self) -> int:
        return self._step_count

    # ------------------------------------------------------------------ internals

    def _pop_action(self) -> np.ndarray:
        with self._action_lock:
            if self._action_queue:
                self._last_action = self._action_queue.popleft()
                return self._last_action
        # Queue drained — hold the last full command (zero-order hold). Holding
        # the weld bits too is intentional: dropping them would release a
        # carried object the instant the 5 Hz brain lags behind the 50 Hz loop.
        return self._last_action.copy()

    def _loop(self):
        """Background control loop — runs at ~control_freq Hz."""
        obs = self.env.reset()
        self.obs_buffer.push(obs)
        vision_every = self.env.config.get("vision", {}).get("update_every_n_steps", 5)

        while self._running:
            t0 = time.monotonic()
            self.step(vision_every=vision_every)   # renders vision + multi_view together
            elapsed = time.monotonic() - t0
            sleep = self._target_dt - elapsed
            if sleep > 0:
                time.sleep(sleep)