"""
Simulated X2, for building and testing the dashboard without a robot present.

It is deliberately more than a random-number generator: the battery drains at a
rate that depends on what the robot is doing, odometry integrates the velocity
you command, joints ease toward their targets, and mode transitions enforce the
same preconditions the real robot does. That means UI logic exercised here
behaves the same way once a real bridge is attached.
"""

from __future__ import annotations

import asyncio
import math
import struct
import time
import zlib

from . import x2_spec
from .bridge_base import Bridge, CommandResult


def _png(width: int, height: int, rgb_rows: list[bytes]) -> bytes:
    """Minimal PNG encoder - stdlib only, no Pillow dependency."""
    raw = b"".join(b"\x00" + row for row in rgb_rows)

    def chunk(tag: bytes, data: bytes) -> bytes:
        body = tag + data
        return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body) & 0xFFFFFFFF)

    return (
        b"\x89PNG\r\n\x1a\n"
        + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
        + chunk(b"IDAT", zlib.compress(raw, 6))
        + chunk(b"IEND", b"")
    )


class MockBridge(Bridge):
    name = "mock"
    simulated = True

    def __init__(self, hub, config: dict):
        super().__init__(hub, config)
        self._task: asyncio.Task | None = None
        self._started = 0.0
        self._joint_targets: dict[str, dict[str, float]] = {}
        self._preset_until = 0.0
        self._preset_label = ""
        self._last_command = 0.0
        self._frame_seq = 0
        # Every feed starts off, matching the real agent.
        self._streams: dict[str, bool] = {
            **{camera["key"]: False for camera in x2_spec.CAMERAS},
            x2_spec.LIDAR["key"]: False,
        }

    # -- lifecycle ----------------------------------------------------------

    async def start(self) -> None:
        self._started = time.time()

        self.state.connection.online = True
        self.state.connection.transport = "mock"
        self.state.connection.host = self.config.get("robot_host") or "simulated"
        self.state.connection.ros_domain_id = self.config.get("ros_domain_id", 0)
        self.state.connection.since = self._started
        self.state.connection.simulated = True
        self.state.connection.error = ""

        self.state.mode = "DAMPING_DEFAULT"
        self.state.mode_desc = "Damping"
        self.state.mode_status = "ready"

        self.state.battery_pct = 87.0
        self.state.battery_voltage = 50.4
        self.state.battery_current = -2.1
        self.state.battery_temp = 31.5
        self.state.battery_cycles = 142
        self.state.pmu_temp = 38.0
        self.state.fan_rpm = 2400.0
        self.state.fan_pct = 42.0
        self.state.hand_type = "OmniHand Dynamic Edition 2025"

        self.state.hand_type = "None"
        for rail in x2_spec.PMU_RAILS:
            self.state.rails[rail["key"]] = {
                "label": rail["label"],
                "voltage": rail["nominal"],
                "current": 0.8,
                "nominal": rail["nominal"],
                "ok": True,
            }

        for group in x2_spec.JOINT_GROUPS:
            self.state.joints[group["key"]] = [
                {"name": j["name"], "label": j["label"], "position": 0.0,
                 "velocity": 0.0, "effort": 0.0, "error": 0}
                for j in group["joints"]
            ]
            self._joint_targets[group["key"]] = {j["name"]: 0.0 for j in group["joints"]}

        # This unit reports hand type NONE, so the simulator matches it: no hand
        # joints are reported until hardware is attached.
        self.state.hand_state = {"left": [], "right": [], "left_type": 0, "right_type": 0}

        self.state.input_sources = list(x2_spec.BUILTIN_INPUT_SOURCES)
        self.state.touch_head = {"touched": False, "zones": [False] * x2_spec.TOUCH_ZONE_COUNT}

        self._task = asyncio.create_task(self._loop(), name="mock-bridge")
        await self.hub.emit("info", "bridge", "Simulation bridge started - no robot required")

    async def stop(self) -> None:
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
        self.state.connection.online = False
        self.state.connection.transport = "none"

    # -- simulation loop ----------------------------------------------------

    async def _loop(self) -> None:
        hz = max(1, int(self.config.get("telemetry_hz", 10)))
        dt = 1.0 / hz
        tick = 0

        while True:
            try:
                await asyncio.sleep(dt)
                tick += 1
                now = time.time()

                self._step_motion(dt, now)
                self._step_joints(dt)
                self._step_imu(now)
                self._step_power(dt, tick, hz)
                self._step_sensors(now)

                self.state.updated = now
                for topic in (x2_spec.TOPIC_IMU_CHEST, x2_spec.TOPIC_IMU_TORSO,
                              x2_spec.TOPIC_LOCOMOTION_VELOCITY):
                    self.state.mark_topic(topic)

                self._record_series(now)
                await self.hub.broadcast("state", self.state.as_dict())

            except asyncio.CancelledError:
                raise
            except Exception as exc:  # keep the loop alive; surface the fault
                await self.hub.emit("error", "bridge", f"Simulation step failed: {exc}")
                await asyncio.sleep(1.0)

    def _step_motion(self, dt: float, now: float) -> None:
        cmd = self.state.velocity_command
        moving_mode = self.state.mode in x2_spec.DRIVEABLE_MODES

        # Dead-man: a browser that stops sending joystick frames must not leave
        # the robot walking. Same rule the real bridge enforces.
        deadman = float(self.config.get("locomotion_deadman_s", x2_spec.LOCOMOTION_DEADMAN_S))
        if self._last_command and (now - self._last_command) > deadman:
            cmd = {"forward": 0.0, "lateral": 0.0, "angular": 0.0}
            self.state.velocity_command = dict(cmd)

        target = cmd if moving_mode else {"forward": 0.0, "lateral": 0.0, "angular": 0.0}

        # Apply the documented activation dead-band.
        eff = {}
        for axis, key in (("forward", "forward"), ("lateral", "lateral"), ("angular", "angular")):
            value = target[axis]
            if abs(value) < x2_spec.VELOCITY_THRESHOLDS[key] and abs(self.state.velocity[axis]) < 1e-3:
                value = 0.0
            eff[axis] = value

        # First-order lag toward the commanded velocity.
        alpha = min(1.0, dt * 4.0)
        for axis in ("forward", "lateral", "angular"):
            self.state.velocity[axis] = round(
                self.state.velocity[axis] + (eff[axis] - self.state.velocity[axis]) * alpha, 4
            )

        yaw = self.state.odom["yaw"] + self.state.velocity["angular"] * dt
        fwd = self.state.velocity["forward"] * dt
        lat = self.state.velocity["lateral"] * dt
        self.state.odom = {
            "x": round(self.state.odom["x"] + fwd * math.cos(yaw) - lat * math.sin(yaw), 4),
            "y": round(self.state.odom["y"] + fwd * math.sin(yaw) + lat * math.cos(yaw), 4),
            "yaw": round(math.atan2(math.sin(yaw), math.cos(yaw)), 4),
        }

        if self._preset_until and now > self._preset_until:
            self._preset_until = 0.0
            self._preset_label = ""

    def _step_joints(self, dt: float) -> None:
        speed = min(1.0, dt * 3.0)
        walking = abs(self.state.velocity["forward"]) + abs(self.state.velocity["angular"]) > 0.05
        phase = time.time() * 4.0

        for group_key, joints in self.state.joints.items():
            targets = self._joint_targets.get(group_key, {})
            for idx, joint in enumerate(joints):
                target = targets.get(joint["name"], 0.0)

                # Legs get a gait overlay while walking so the joint view is alive.
                if group_key == "leg" and walking:
                    side = 0.0 if idx < 6 else math.pi
                    target += 0.22 * math.sin(phase + side) * (1 if idx % 6 in (0, 3) else 0.4)
                if group_key == "arm" and self._preset_until:
                    target += 0.5 * math.sin(phase * 0.8) * (1 if idx % 7 < 3 else 0.3)

                previous = joint["position"]
                joint["position"] = round(previous + (target - previous) * speed, 4)
                joint["velocity"] = round((joint["position"] - previous) / dt, 4)
                joint["effort"] = round(abs(joint["velocity"]) * 3.2 + abs(joint["position"]) * 1.4, 3)

    def _step_imu(self, now: float) -> None:
        walking = abs(self.state.velocity["forward"]) > 0.05
        wobble = 0.035 if walking else 0.004
        base = now * 3.0

        for key, offset in (("chest", 0.0), ("torso", 1.1)):
            roll = wobble * math.sin(base * 2.1 + offset)
            pitch = wobble * math.cos(base * 1.7 + offset) - self.state.velocity["forward"] * 0.08
            self.state.imu[key] = {
                "roll": round(roll, 4),
                "pitch": round(pitch, 4),
                "yaw": round(self.state.odom["yaw"], 4),
                "accel_x": round(self.state.velocity["forward"] * 1.2 + wobble * 9 * math.sin(base * 5), 3),
                "accel_y": round(self.state.velocity["lateral"] * 1.2 + wobble * 7 * math.cos(base * 4), 3),
                "accel_z": round(9.81 + wobble * 5 * math.sin(base * 6 + offset), 3),
                "gyro_x": round(wobble * 4 * math.cos(base * 3), 4),
                "gyro_y": round(wobble * 4 * math.sin(base * 3.3), 4),
                "gyro_z": round(self.state.velocity["angular"], 4),
                "temp": round(34.0 + math.sin(base * 0.05) * 1.5, 2),
            }

    def _step_power(self, dt: float, tick: int, hz: int) -> None:
        # Draw scales with what the robot is actually doing.
        load = 0.5
        if self.state.mode in ("STAND_DEFAULT", "LOCOMOTION_DEFAULT"):
            load = 1.4
        load += abs(self.state.velocity["forward"]) * 2.0
        load += abs(self.state.velocity["angular"]) * 1.2
        if self._preset_until:
            load += 1.0

        drain_per_hour = 6.0 * load
        self.state.battery_pct = max(0.0, round(self.state.battery_pct - drain_per_hour * dt / 3600.0, 4))
        self.state.battery_current = round(-3.0 * load, 3)
        self.state.battery_voltage = round(44.0 + (self.state.battery_pct / 100.0) * 8.4, 3)
        self.state.battery_temp = round(30.0 + load * 3.5 + math.sin(time.time() * 0.02) * 0.6, 2)
        self.state.pmu_temp = round(35.0 + load * 4.0, 2)
        self.state.fan_pct = round(min(100.0, 30.0 + load * 28.0), 1)
        self.state.fan_rpm = round(1200 + self.state.fan_pct * 42, 0)
        self.state.charging = False

        for key, rail in self.state.rails.items():
            nominal = rail["nominal"]
            rail["voltage"] = round(nominal - load * 0.12 + math.sin(time.time() * 0.7) * 0.03, 3)
            rail["current"] = round(0.6 + load * (1.8 if key in ("bus_48v", "orin") else 0.5), 3)
            rail["ok"] = rail["voltage"] > nominal * 0.85

        if tick % max(1, hz * 5) == 0:
            self.state.mark_topic(x2_spec.TOPIC_PMU_STATE, x2_spec.PMU_RATE_HZ)

    def _step_sensors(self, now: float) -> None:
        # A periodic touch event so the sensors tab shows something happening.
        touched = (int(now) % 37) < 2
        zones = [False] * x2_spec.TOUCH_ZONE_COUNT
        if touched:
            zones[int(now) % x2_spec.TOUCH_ZONE_COUNT] = True
        self.state.touch_head = {"touched": touched, "zones": zones}
        self.state.mark_topic(x2_spec.TOPIC_TOUCH_HEAD, 100)

    def _record_series(self, now: float) -> None:
        self.hub.record("battery_pct", self.state.battery_pct, now)
        self.hub.record("battery_voltage", self.state.battery_voltage, now)
        self.hub.record("battery_current", self.state.battery_current, now)
        self.hub.record("battery_temp", self.state.battery_temp, now)
        self.hub.record("pmu_temp", self.state.pmu_temp, now)
        self.hub.record("fan_pct", self.state.fan_pct, now)
        self.hub.record("vel_forward", self.state.velocity["forward"], now)
        self.hub.record("vel_lateral", self.state.velocity["lateral"], now)
        self.hub.record("vel_angular", self.state.velocity["angular"], now)
        chest = self.state.imu.get("chest", {})
        self.hub.record("imu_roll", chest.get("roll"), now)
        self.hub.record("imu_pitch", chest.get("pitch"), now)
        self.hub.record("imu_yaw", chest.get("yaw"), now)

    # -- commands -----------------------------------------------------------

    async def set_mode(self, mode_id: str) -> CommandResult:
        match = next((m for m in x2_spec.MC_MODES if m["id"] == mode_id), None)
        if not match:
            return CommandResult.failure(f"Unknown mode '{mode_id}'")

        self.state.mode = mode_id
        self.state.mode_desc = match["label"]
        self.state.mode_status = "Running"
        if mode_id not in x2_spec.DRIVEABLE_MODES:
            self.state.velocity_command = {"forward": 0.0, "lateral": 0.0, "angular": 0.0}
        await self.hub.emit("info", "mode", f"Mode set to {match['label']}")
        return CommandResult.success(f"Mode set to {match['label']}", {"mode": mode_id})

    async def set_velocity(self, forward: float, lateral: float, angular: float) -> CommandResult:
        if self.state.mode not in x2_spec.DRIVEABLE_MODES:
            return CommandResult.failure("Switch to Stable stand or Locomotion before commanding velocity")
        if not self.state.source_registered:
            return CommandResult.failure("Register an MC input source first (Control tab)")

        self.state.velocity_command = {
            "forward": round(float(forward), 4),
            "lateral": round(float(lateral), 4),
            "angular": round(float(angular), 4),
        }
        self._last_command = time.time()
        return CommandResult.success("Velocity accepted", self.state.velocity_command)

    async def stop_motion(self) -> CommandResult:
        self.state.velocity_command = {"forward": 0.0, "lateral": 0.0, "angular": 0.0}
        self._last_command = time.time()
        return CommandResult.success("Motion stopped")

    async def play_preset(self, motion: int, area: int, interrupt: bool = True) -> CommandResult:
        if self.state.mode != "STAND_DEFAULT":
            return CommandResult.failure("Preset motions require Stable stand mode")
        preset = next((p for p in x2_spec.PRESET_MOTIONS
                       if p["motion"] == motion and p["area"] == area), None)
        label = preset["label"] if preset else f"motion {motion}"
        self._preset_until = time.time() + 3.0
        self._preset_label = label
        await self.hub.emit("info", "motion", f"Playing preset: {label}")
        return CommandResult.success(f"Playing {label}",
                                     {"task_id": int(time.time() * 1000) % 1_000_000})

    async def set_joints(self, group: str, mode: str, targets: dict[str, float],
                         stiffness=None, damping=None) -> CommandResult:
        if group not in self._joint_targets:
            return CommandResult.failure(f"Unknown joint group '{group}'")
        spec = x2_spec.JOINT_GROUP_BY_KEY[group]
        limits = {j["name"]: j for j in spec["joints"]}

        applied, rejected = {}, {}
        for name, value in targets.items():
            joint = limits.get(name)
            if joint is None:
                rejected[name] = "unknown joint"
                continue
            lo = math.radians(joint["min_deg"])
            hi = math.radians(joint["max_deg"])
            clamped = max(lo, min(hi, float(value)))
            if abs(clamped - float(value)) > 1e-6:
                rejected[name] = f"clamped to [{joint['min_deg']}, {joint['max_deg']}] deg"
            self._joint_targets[group][name] = clamped
            applied[name] = round(clamped, 4)

        return CommandResult.success(f"{len(applied)} joint target(s) applied",
                                     {"applied": applied, "adjusted": rejected})

    async def set_hand(self, side: str, positions: list[float]) -> CommandResult:
        if side not in ("left", "right"):
            return CommandResult.failure(f"Unknown side '{side}'")
        # Mirror the real unit: with hand type NONE there is nothing to drive.
        self.state.hand_state[side] = [
            {"name": name, "position": round(float(positions[i]), 4) if i < len(positions) else 0.0,
             "velocity": 0.0, "effort": 0.0, "fault": 0}
            for i, name in enumerate(x2_spec.DEXHAND_JOINTS)
        ]
        return CommandResult.success(f"{side} hand updated")

    async def speak(self, text: str, priority: int = 6, interrupt: bool = False) -> CommandResult:
        if not text.strip():
            return CommandResult.failure("Nothing to say")
        await self.hub.emit("info", "voice", f'TTS: "{text[:120]}"')
        return CommandResult.success("Queued for speech", {"text": text, "priority": priority})

    async def set_volume(self, volume: int) -> CommandResult:
        self.state.volume = max(0, min(100, int(volume)))
        return CommandResult.success(f"Volume {self.state.volume}")

    async def set_mute(self, muted: bool) -> CommandResult:
        self.state.muted = bool(muted)
        return CommandResult.success("Muted" if muted else "Unmuted")

    async def play_emoji(self, emotion_id: int, mode: int = 1, priority: int = 6) -> CommandResult:
        self.state.emoji_id = int(emotion_id)
        entry = next((e for e in x2_spec.EMOJIS if e["id"] == int(emotion_id)), None)
        return CommandResult.success(f"Showing {entry['label'] if entry else emotion_id}")

    async def set_led(self, mode: int, r: int, g: int, b: int, priority: int = 6,
                      keep: bool = True) -> CommandResult:
        self.state.led = {"mode": int(mode), "r": int(r), "g": int(g),
                          "b": int(b), "keep": bool(keep)}
        return CommandResult.success("LED updated", self.state.led)

    async def register_input_source(self, name: str, priority: int, timeout: int) -> CommandResult:
        existing = next((s for s in self.state.input_sources if s["name"] == name), None)
        if existing:
            existing.update({"priority": priority, "timeout": timeout})
        else:
            self.state.input_sources.append(
                {"name": name, "priority": priority, "timeout": timeout, "desc": "This dashboard"}
            )
        self.state.source_registered = True
        # The built-in sources are listed for reference but nothing is actually
        # driving them in simulation, so the newly registered source wins
        # arbitration. On a real robot the firmware decides this.
        self.state.input_source = name
        await self.hub.emit("info", "control", f"Input source '{name}' registered at priority {priority}")
        return CommandResult.success(f"Registered '{name}'", {"current": self.state.input_source})

    async def publish_raw(self, topic: str, message_type: str, payload: dict) -> CommandResult:
        self.state.mark_topic(topic)
        return CommandResult.success(f"[simulated] published to {topic}", payload)

    async def call_service(self, service: str, service_type: str, payload: dict) -> CommandResult:
        return CommandResult.success(f"[simulated] called {service}", {"request": payload})

    # -- synthetic camera ---------------------------------------------------

    async def camera_frame(self, camera_key: str, stream: str = "rgb",
                           flip: bool | None = None) -> bytes | None:
        """A moving synthetic scene, so the Vision tab is testable without hardware."""
        # Off means off in simulation too, or the Vision tab's switches would
        # look broken here and only work against real hardware.
        if not self._streams.get(camera_key):
            return None

        width, height = 192, 144
        self._frame_seq += 1
        t = time.time()
        yaw = self.state.odom["yaw"]
        depth = stream == "depth" or camera_key.startswith("depth")

        # Horizon shifts with pitch, scene pans with yaw - so driving the robot
        # visibly changes the view.
        pitch = self.state.imu.get("chest", {}).get("pitch", 0.0)
        horizon = int(height * 0.55 + pitch * 220)
        pan = (yaw * 90.0 + t * 6.0) % width

        rows = []
        for y in range(height):
            row = bytearray()
            for x in range(width):
                if depth:
                    # Distance rises toward the horizon; a nearby object sweeps past.
                    d = abs(y - horizon) / height
                    obj = math.exp(-(((x - pan) % width - width / 2) ** 2) / 400.0)
                    v = max(0.0, min(1.0, d * 1.6 - obj * 0.5))
                    # Single-hue blue ramp, light = near, dark = far.
                    r = int(205 * (1 - v) + 13 * v)
                    g = int(226 * (1 - v) + 54 * v)
                    b = int(251 * (1 - v) + 107 * v)
                else:
                    if y < horizon:
                        shade = y / max(1, horizon)
                        r, g, b = int(24 + 26 * shade), int(28 + 32 * shade), int(38 + 46 * shade)
                    else:
                        shade = (y - horizon) / max(1, height - horizon)
                        r, g, b = int(30 + 24 * shade), int(34 + 20 * shade), int(32 + 16 * shade)
                    # Vertical markers that slide with yaw, giving visible motion.
                    if (int(x + pan) // 24) % 2 == 0 and y > horizon:
                        r, g, b = min(255, r + 26), min(255, g + 22), min(255, b + 18)
                    if abs(y - horizon) <= 1:
                        r, g, b = 57, 135, 229
                row += bytes((r, g, b))
            rows.append(bytes(row))

        if flip:
            # 180 degrees = reverse the row order, and the pixels within each
            # row. Pixels are 3 bytes, so reverse in triples, not bytewise.
            def reverse_pixels(row: bytes) -> bytes:
                return b"".join(row[i - 3:i] for i in range(len(row), 0, -3))

            rows = [reverse_pixels(row) for row in reversed(rows)]

        return _png(width, height, rows)

    def frame_content_type(self, camera_key: str = "", stream: str = "rgb") -> str:
        return "image/png"

    # -- on-demand streams --------------------------------------------------

    async def set_stream(self, key: str, active: bool) -> CommandResult:
        self._streams[key] = bool(active)
        self.state.custom["streams"] = {
            k: {"key": k, "active": v} for k, v in self._streams.items()
        }
        return CommandResult.success(
            f"[simulated] {key} {'on' if active else 'off'}", {"key": key, "active": active})

    async def list_streams(self) -> CommandResult:
        return CommandResult.success("[simulated] streams", {
            "streams": {k: {"key": k, "active": v} for k, v in self._streams.items()},
        })

    async def lidar_points(self) -> CommandResult:
        """A synthetic room, so the LiDAR view is testable with no robot attached."""
        if not self._streams.get(x2_spec.LIDAR["key"]):
            return CommandResult.failure("LiDAR is off")

        t = time.time()
        points = []
        # Four walls and a floor patch, swept by a rotating scan line so the
        # view visibly updates.
        for i in range(900):
            angle = (i / 900.0) * math.tau
            radius = 3.0 + 0.8 * math.sin(angle * 4 + t * 0.3)
            for z in (-0.4, 0.1, 0.6, 1.1):
                points.append([round(radius * math.cos(angle), 3),
                               round(radius * math.sin(angle), 3),
                               z, round(40 + 60 * abs(math.sin(angle + t)), 1)])
        return CommandResult.success("[simulated] points", {
            "ts": t, "count": len(points), "points": points,
            "frame": x2_spec.LIDAR["frame"],
        })
