"""Robot arm controller — real DDS motor commands and JSONL motion replay.

Production-grade replay engine ported from G1_Lootah/Controller/g1_replay_trigger_r2x.py.
Features: body-lock, ramp-in/out interpolation, watchdog, speed control, CRC.
Falls back gracefully to simulation when the Unitree SDK is unavailable.
"""

from __future__ import annotations

import json
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from Project.Sanad.config import (
    G1_NUM_MOTOR,
    KD_HIGH,
    KD_LOW,
    KD_WRIST,
    KP_HIGH,
    KP_LOW,
    KP_WRIST,
    MOTIONS_DIR,
    REPLAY_HZ,
    WEAK_MOTORS,
    WRIST_MOTORS,
)
from Project.Sanad.core.config_loader import section as _cfg_section
from Project.Sanad.core.event_bus import bus
from Project.Sanad.core.logger import get_logger

log = get_logger("arm_controller")

_AC = _cfg_section("motion", "arm_controller")
# G1 hardware constants — single source in config/core_config.json
from Project.Sanad.config import ENABLE_ARM_SDK_INDEX
RAMP_IN_STEPS = _AC.get("ramp_in_steps", 60)           # ~1.0s smooth move to start pose
RAMP_OUT_STEPS = _AC.get("ramp_out_steps", 180)         # ~3.0s smooth return to home
WATCHDOG_TIMEOUT = _AC.get("watchdog_timeout_sec", 0.25)       # hold last pose if state stale
WATCHDOG_DISABLE_AFTER = _AC.get("watchdog_disable_after_sec", 1.0)  # abort if state stale this long
ARM_INDICES = range(
    _AC.get("arm_indices_start", 15),
    _AC.get("arm_indices_stop", 29),
)

# -- SDK import (optional) --
try:
    from unitree_sdk2py.core.channel import (
        ChannelFactoryInitialize,
        ChannelPublisher,
        ChannelSubscriber,
    )
    from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_, LowState_
    from unitree_sdk2py.utils.crc import CRC

    # IDL factory — LowCmd_() with no args fails because the dataclass
    # has 5 required positional fields. The SDK ships a default factory
    # named `unitree_hg_msg_dds__LowCmd_` that constructs a fully-zeroed
    # instance with the right number of motor_cmd entries.
    try:
        from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
        _make_low_cmd = unitree_hg_msg_dds__LowCmd_
    except ImportError:
        _make_low_cmd = None

    _HAS_SDK = True
except ImportError:
    _HAS_SDK = False
    _make_low_cmd = None
    log.warning("Unitree SDK not available — running in simulation mode")

# Optional G1 high-level SDK action client (built-in actions)
try:
    from unitree_sdk2py.g1.loco.g1_loco_client import LocoClient
    _HAS_LOCO = True
except ImportError:
    _HAS_LOCO = False


@dataclass
class Action:
    name: str
    id: int
    file: str = ""          # JSONL filename (empty = SDK built-in)
    category: str = "sdk"   # "sdk" | "jsonl"


# -- SDK actions (fixed — built into Unitree firmware) --
SDK_ACTIONS: list[Action] = [
    Action("release_arm", 0, category="sdk"),
    Action("shake_hand", 1, category="sdk"),
    Action("high_five", 2, category="sdk"),
    Action("hug", 3, category="sdk"),
    Action("high_wave", 4, category="sdk"),
    Action("clap", 5, category="sdk"),
    Action("face_wave", 6, category="sdk"),
    Action("left_kiss", 7, category="sdk"),
    Action("heart", 8, category="sdk"),
    Action("right_heart", 9, category="sdk"),
    Action("hands_up", 10, category="sdk"),
    Action("x_ray", 11, category="sdk"),
    Action("right_hand_up", 12, category="sdk"),
    Action("reject", 13, category="sdk"),
    Action("right_kiss", 14, category="sdk"),
    Action("two_hand_kiss", 15, category="sdk"),
]

# Next auto-ID for JSONL actions starts after SDK range.
_JSONL_ID_START = _AC.get("jsonl_id_start", 100)


def _scan_jsonl_actions() -> list[Action]:
    """Auto-discover all .jsonl files in data/motions/ and create actions.

    Called at startup and whenever the dashboard requests a refresh.
    The action name is derived from the filename (without extension),
    with underscores replacing hyphens/spaces.
    """
    MOTIONS_DIR.mkdir(parents=True, exist_ok=True)
    actions = []
    for idx, path in enumerate(sorted(MOTIONS_DIR.glob("*.jsonl"))):
        name = path.stem.replace("-", "_").replace(" ", "_")
        actions.append(Action(
            name=name,
            id=_JSONL_ID_START + idx,
            file=path.name,
            category="jsonl",
        ))
    return actions


def rebuild_action_registry() -> tuple[list[Action], dict[int, Action], dict[str, Action]]:
    """Rebuild the full action list from SDK + disk scan. Called on startup and refresh."""
    jsonl_actions = _scan_jsonl_actions()
    all_actions = list(SDK_ACTIONS) + jsonl_actions
    by_id = {a.id: a for a in all_actions}
    by_name = {a.name: a for a in all_actions}
    log.info("Action registry: %d SDK + %d JSONL = %d total",
             len(SDK_ACTIONS), len(jsonl_actions), len(all_actions))
    return all_actions, by_id, by_name


# Initial build
ACTIONS, ACTION_BY_ID, ACTION_BY_NAME = rebuild_action_registry()


def _lerp_q(q_start: list[float], q_end: list[float], t: float) -> list[float]:
    """Linear interpolation between two joint-position vectors, t in [0,1]."""
    return [a + (b - a) * t for a, b in zip(q_start, q_end)]


def _load_frames(path: Path) -> list[dict[str, Any]]:
    """Read JSONL file, return list of frames with 't' and 'q' keys."""
    frames = []
    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            data = json.loads(line)
            if "q" in data:
                frames.append(data)
    return frames


def _load_home_q(home_file: str = "arm_home.jsonl") -> list[float] | None:
    path = MOTIONS_DIR / home_file
    if not path.exists():
        return None
    frames = _load_frames(path)
    return frames[0]["q"] if frames else None


class ArmController:
    """Thread-safe arm controller with real DDS replay and simulation fallback."""

    def __init__(self):
        self._lock = threading.Lock()
        self._cancel = threading.Event()
        self._is_busy = False
        self._last_action_time = 0.0
        self.cooldown_sec = 1.0
        self._initialized = False

        # DDS handles (set in init())
        self._arm_pub = None
        self._state_sub = None
        self._low_cmd = None
        self._crc = None
        self._loco_client = None

        # Live state from LowState_ subscriber
        self._current_q: list[float] = [0.0] * G1_NUM_MOTOR
        self._last_state_time = 0.0
        self._state_lock = threading.Lock()
        self._first_state_event = threading.Event()

        # Cached motion file metadata
        self._motion_files_cache: dict[str, dict[str, Any]] = {}

    # -- initialization --

    def init(self, network_interface: str = "lo") -> bool:
        if self._initialized:
            return True
        if not _HAS_SDK:
            log.info("Simulation mode — DDS init skipped")
            return False
        try:
            ChannelFactoryInitialize(0, network_interface)
            self._arm_pub = ChannelPublisher("rt/arm_sdk", LowCmd_)
            self._arm_pub.Init()
            self._state_sub = ChannelSubscriber("rt/lowstate", LowState_)
            self._state_sub.Init(self._on_low_state, 10)
            # IDL types need the SDK's default factory — bare LowCmd_() fails
            # because the dataclass has 5 required positional fields.
            if _make_low_cmd is not None:
                self._low_cmd = _make_low_cmd()
            else:
                # Last-resort: try a few constructor signatures
                try:
                    self._low_cmd = LowCmd_()
                except TypeError:
                    # Build with explicit zeroed fields
                    from unitree_sdk2py.idl.unitree_hg.msg.dds_ import MotorCmd_
                    try:
                        from unitree_sdk2py.idl.default import (
                            unitree_hg_msg_dds__MotorCmd_ as _make_motor_cmd,
                        )
                    except ImportError:
                        _make_motor_cmd = lambda: MotorCmd_(
                            mode=0, q=0.0, dq=0.0, tau=0.0, kp=0.0, kd=0.0, reserve=0,
                        )
                    self._low_cmd = LowCmd_(
                        mode_pr=0,
                        mode_machine=0,
                        motor_cmd=[_make_motor_cmd() for _ in range(35)],
                        reserve=[0, 0, 0, 0],
                        crc=0,
                    )
            self._crc = CRC()

            # High-level SDK action client (for built-in actions)
            if _HAS_LOCO:
                try:
                    self._loco_client = LocoClient()
                    self._loco_client.SetTimeout(10.0)
                    self._loco_client.Init()
                    log.info("LocoClient initialized — built-in SDK actions available")
                except Exception as exc:
                    log.warning("LocoClient init failed: %s — built-in actions disabled", exc)
                    self._loco_client = None

            self._initialized = True
            log.info("DDS initialized on %s", network_interface)
        except Exception as exc:
            log.error("DDS init failed: %s", exc)
        return self._initialized

    def _on_low_state(self, msg):
        """Callback from DDS subscriber — updates current joint positions."""
        with self._state_lock:
            self._current_q = [float(msg.motor_state[i].q) for i in range(G1_NUM_MOTOR)]
            self._last_state_time = time.monotonic()
        if not self._first_state_event.is_set():
            self._first_state_event.set()

    def wait_for_state(self, timeout: float = 2.0) -> bool:
        """Block until first LowState_ callback fires (or timeout). Returns True if state received."""
        return self._first_state_event.wait(timeout=timeout)

    # -- internal API exposed for teaching/macro_player (encapsulation boundary) --

    def get_current_q(self) -> list[float]:
        """Public read of current joint positions."""
        return self._get_current_q()

    def send_frame(self, arm_target_q: list[float], body_lock_q: list[float]):
        """Public single-frame send. Use only inside a controlled playback loop."""
        self._send_frame(arm_target_q, body_lock_q)

    def disable(self):
        """Public disable — releases arm SDK control."""
        self._disable_sdk()

    def state_age(self) -> float:
        """Seconds since last LowState_ callback."""
        return self._state_age()

    def _get_current_q(self) -> list[float]:
        with self._state_lock:
            return list(self._current_q)

    def _state_age(self) -> float:
        with self._state_lock:
            return time.monotonic() - self._last_state_time if self._last_state_time else 999.0

    # -- frame sending (real DDS with CRC) --

    def _send_frame(self, arm_target_q: list[float], body_lock_q: list[float]):
        """Send one motor frame via DDS. Body stays locked, arms get target."""
        if not self._initialized or self._low_cmd is None:
            return

        # Enable ARM_SDK
        self._low_cmd.motor_cmd[ENABLE_ARM_SDK_INDEX].q = 1.0

        for i in range(G1_NUM_MOTOR):
            self._low_cmd.motor_cmd[i].mode = 1
            self._low_cmd.motor_cmd[i].dq = 0.0
            self._low_cmd.motor_cmd[i].tau = 0.0

            # Arms get replay data, body stays locked
            if i >= 15:
                self._low_cmd.motor_cmd[i].q = arm_target_q[i]
            else:
                self._low_cmd.motor_cmd[i].q = body_lock_q[i]

            # Per-motor gains
            if i in WEAK_MOTORS:
                kp, kd = KP_LOW, KD_LOW
            elif i in WRIST_MOTORS:
                kp, kd = KP_WRIST, KD_WRIST
            else:
                kp, kd = KP_HIGH, KD_HIGH
            self._low_cmd.motor_cmd[i].kp = kp
            self._low_cmd.motor_cmd[i].kd = kd

        self._low_cmd.crc = self._crc.Crc(self._low_cmd)
        self._arm_pub.Write(self._low_cmd)

    def _disable_sdk(self):
        """Send 5 disable frames to cleanly release arm control."""
        if not self._initialized or self._low_cmd is None:
            return
        for _ in range(5):
            self._low_cmd.motor_cmd[ENABLE_ARM_SDK_INDEX].q = 0.0
            self._low_cmd.crc = self._crc.Crc(self._low_cmd)
            self._arm_pub.Write(self._low_cmd)
            time.sleep(0.02)

    # -- replay engine --

    def replay_file(self, filepath: str, speed: float = 1.0):
        """Play a JSONL motion file with full production replay logic.

        Args:
            filepath: Path to .jsonl motion file.
            speed: Playback speed multiplier (1.0 = normal).
        """
        with self._lock:
            if self._is_busy:
                log.warning("replay_file: arm busy, skipping")
                return
            self._is_busy = True
            self._cancel.clear()

        try:
            self._replay_file_inner(filepath, speed)
        finally:
            with self._lock:
                self._is_busy = False
                self._last_action_time = time.monotonic()

    def _replay_file_inner(self, filepath: str, speed: float = 1.0):
        path = Path(filepath)
        if not path.is_absolute():
            path = MOTIONS_DIR / path
        if not path.exists():
            raise FileNotFoundError(f"Motion file not found: {path}")

        frames = _load_frames(path)
        if not frames:
            log.warning("Empty motion file: %s", path)
            return

        home_q = _load_home_q() or [0.0] * G1_NUM_MOTOR

        if not _HAS_SDK:
            duration = len(frames) / REPLAY_HZ / speed if speed else len(frames) / REPLAY_HZ
            log.info("[SIM] Replay %s (%.1fs, %d frames, speed=%.1f)", path.name, duration, len(frames), speed)
            self._sim_replay(frames, speed)
            return

        log.info("Replay %s (%d frames @ %.0f Hz, speed=%.1f)", path.name, len(frames), REPLAY_HZ, speed)

        # CRITICAL: wait for first LowState_ callback before reading _current_q
        # Otherwise we'd lock the body to all-zeros and crash the robot.
        if not self._first_state_event.is_set():
            log.warning("Waiting for first LowState message...")
            if not self._first_state_event.wait(timeout=2.0):
                log.error("No LowState received in 2s — refusing to replay (would lock body to zeros)")
                return

        body_lock_q = self._get_current_q()
        current_arm_q = list(body_lock_q)
        start_arm_q = frames[0]["q"]
        interval = 1.0 / REPLAY_HZ

        # Phase 1: Ramp-in — interpolate from current pose to first frame
        log.debug("Ramp-in (%d steps)", RAMP_IN_STEPS)
        for step in range(RAMP_IN_STEPS):
            if self._cancel.is_set():
                self._return_home(current_arm_q, body_lock_q, home_q)
                return
            t = (step + 1) / RAMP_IN_STEPS
            interp = _lerp_q(current_arm_q, start_arm_q, t)
            self._send_frame(interp, body_lock_q)
            time.sleep(interval)

        # Phase 2: Play frames with timing and watchdog
        log.debug("Playing %d frames", len(frames))
        play_elapsed = 0.0
        last_real = time.monotonic()
        t0_frame = frames[0].get("t", 0.0)

        for frame in frames:
            if self._cancel.is_set():
                break

            # Watchdog
            age = self._state_age()
            if age > WATCHDOG_DISABLE_AFTER:
                log.error("Watchdog abort — LowState stale %.2fs", age)
                self._disable_sdk()
                return
            if age > WATCHDOG_TIMEOUT:
                log.warning("Watchdog hold — LowState stale %.2fs", age)

            now_real = time.monotonic()
            play_elapsed += (now_real - last_real) * speed
            last_real = now_real

            frame_t = frame.get("t", 0.0) - t0_frame
            if frame_t > play_elapsed:
                time.sleep(min(frame_t - play_elapsed, interval))

            self._send_frame(frame["q"], body_lock_q)
            # NOTE: timing is driven by frame_t above — no extra sleep here

        # Phase 3: Return home
        last_arm_q = frames[-1]["q"] if frames else current_arm_q
        self._return_home(last_arm_q, body_lock_q, home_q)

    def _return_home(self, from_q: list[float], body_lock_q: list[float], home_q: list[float]):
        """Smooth interpolation back to home pose, then disable SDK."""
        log.debug("Returning home (%d steps)", RAMP_OUT_STEPS)
        interval = 1.0 / REPLAY_HZ
        for step in range(RAMP_OUT_STEPS):
            if self._cancel.is_set():
                break
            t = (step + 1) / RAMP_OUT_STEPS
            interp = _lerp_q(from_q, home_q, t)
            self._send_frame(interp, body_lock_q)
            time.sleep(interval)
        self._disable_sdk()
        log.info("Home reached, SDK disabled")

    def _sim_replay(self, frames: list[dict], speed: float):
        """Simulation replay — emit events, sleep for equivalent duration."""
        interval = 1.0 / REPLAY_HZ
        for i, frame in enumerate(frames):
            if self._cancel.is_set():
                log.info("[SIM] Replay cancelled at frame %d/%d", i, len(frames))
                return
            time.sleep(interval / max(speed, 0.1))
        log.info("[SIM] Replay complete")

    # -- public API --

    @property
    def is_busy(self) -> bool:
        return self._is_busy

    def cancel(self):
        """Cancel the currently running replay."""
        self._cancel.set()

    def refresh_actions(self):
        """Re-scan data/motions/ and rebuild the action registry."""
        global ACTIONS, ACTION_BY_ID, ACTION_BY_NAME
        ACTIONS, ACTION_BY_ID, ACTION_BY_NAME = rebuild_action_registry()

    def list_actions(self) -> list[dict[str, Any]]:
        return [
            {"id": a.id, "name": a.name, "file": a.file, "category": a.category}
            for a in ACTIONS
        ]

    def list_motion_files(self) -> list[dict[str, Any]]:
        """List all JSONL files in data/motions/ with metadata.

        Caches frame count by (path, mtime) to avoid re-parsing megabytes of
        JSONL on every dashboard refresh.
        """
        MOTIONS_DIR.mkdir(parents=True, exist_ok=True)
        result = []
        for p in sorted(MOTIONS_DIR.glob("*.jsonl")):
            stat = p.stat()
            cache_key = f"{p}:{stat.st_mtime_ns}"
            cached = self._motion_files_cache.get(cache_key)
            if cached is None:
                frames = _load_frames(p)
                duration = len(frames) / REPLAY_HZ if frames else 0
                cached = {
                    "name": p.name,
                    "path": str(p),
                    "frames": len(frames),
                    "duration_sec": round(duration, 2),
                    "size_kb": round(stat.st_size / 1024, 1),
                }
                # Drop stale entries for this path before adding new one
                stale = [k for k in self._motion_files_cache if k.startswith(f"{p}:")]
                for k in stale:
                    self._motion_files_cache.pop(k, None)
                self._motion_files_cache[cache_key] = cached
            result.append(cached)
        return result

    def trigger_by_id(self, action_id: int, speed: float = 1.0):
        action = ACTION_BY_ID.get(action_id)
        if action is None:
            raise KeyError(f"Unknown action id: {action_id}")
        self._execute(action, speed)

    def trigger_by_name(self, name: str, speed: float = 1.0):
        action = ACTION_BY_NAME.get(name)
        if action is None:
            raise KeyError(f"Unknown action: {name}")
        self._execute(action, speed)

    def _execute(self, action: Action, speed: float = 1.0):
        with self._lock:
            if self._is_busy:
                log.warning("Arm busy, skipping %s", action.name)
                return
            self._is_busy = True
            self._cancel.clear()

        try:
            bus.emit_sync("motion.action_started", action=action.name)
            if action.file:
                self._replay_file_inner(action.file, speed=speed)
            else:
                self._run_sdk_action(action)
        except Exception as exc:
            log.error("Action %s failed: %s", action.name, exc)
        finally:
            with self._lock:
                self._is_busy = False
                self._last_action_time = time.monotonic()
            bus.emit_sync("motion.action_done", action=action.name)

    def _run_sdk_action(self, action: Action):
        if not _HAS_SDK:
            log.info("[SIM] SDK action: %s (id=%d)", action.name, action.id)
            time.sleep(2.0)
            return
        if self._loco_client is None:
            log.warning(
                "SDK action %s requested but LocoClient not available — skipping",
                action.name,
            )
            return
        log.info("SDK action: %s (id=%d)", action.name, action.id)
        try:
            # G1 LocoClient exposes ExecuteAction(id) for built-in actions
            self._loco_client.ExecuteAction(action.id)
            # Built-in actions are blocking on the robot side, but the SDK call returns
            # immediately. Give it a sensible default duration so we don't hammer it.
            time.sleep(2.5)
        except Exception as exc:
            log.error("SDK action %s failed: %s", action.name, exc)

    def status(self) -> dict[str, Any]:
        return {
            "initialized": self._initialized,
            "sdk_available": _HAS_SDK,
            "busy": self._is_busy,
            "state_age_sec": round(self._state_age(), 3),
            "sdk_actions": len(SDK_ACTIONS),
            "jsonl_actions": len([a for a in ACTIONS if a.category == "jsonl"]),
            "total_actions": len(ACTIONS),
            "total_motion_files": len(list(MOTIONS_DIR.glob("*.jsonl"))) if MOTIONS_DIR.exists() else 0,
        }
