"""Versioned H2 planner packet tail shared by predictor and controller."""

from __future__ import annotations

import math

import numpy as np


LEGACY_FIELDS = 20
H2_PLANNER_PROTOCOL_VERSION = 3.0
H2_PLANNER_PACKET_LEN = 27
PROTOCOL_VERSION_INDEX = 20
ROBOT_SOURCE_MONOTONIC_INDEX = 21
BALL_SOURCE_MONOTONIC_INDEX = 22
ROBOT_BASE_ORIGIN_START_INDEX = 23
FRAME_CONVENTION_INDEX = 26
FRAME_CONVENTION_WORLD_X = 1.0


def source_age_s(now_s: float, source_stamp_s: float) -> float:
    """Return a fail-closed same-host monotonic source age."""
    now = float(now_s)
    stamp = float(source_stamp_s)
    if not math.isfinite(now) or not math.isfinite(stamp) or stamp <= 0.0 or stamp > now + 0.01:
        return math.inf
    return max(0.0, now - stamp)


def build_h2_planner_packet(
    *,
    robot_pos_world,
    robot_quat_xyzw,
    time_to_hit_s: float,
    hit_pos_robot,
    racket_velocity_robot,
    command_id: int,
    incoming_velocity_robot,
    landing_xy_robot,
    robot_source_monotonic_s: float,
    ball_source_monotonic_s: float,
    robot_base_origin_world,
    allow_missing_sources: bool = False,
) -> np.ndarray:
    """Build the exact 27-value H2 planner-v3 wire payload.

    Both the ROS predictor and the isolated MuJoCo bridge call this function,
    preventing a simulation-only packet layout from drifting away from the
    real deployment producer.
    """

    robot_pos = np.asarray(robot_pos_world, dtype=np.float64).reshape(3)
    robot_quat = np.asarray(robot_quat_xyzw, dtype=np.float64).reshape(4)
    hit_pos = np.asarray(hit_pos_robot, dtype=np.float64).reshape(3)
    racket_velocity = np.asarray(racket_velocity_robot, dtype=np.float64).reshape(3)
    incoming_velocity = np.asarray(incoming_velocity_robot, dtype=np.float64).reshape(3)
    landing_xy = np.asarray(landing_xy_robot, dtype=np.float64).reshape(2)
    origin = np.asarray(robot_base_origin_world, dtype=np.float64).reshape(3)
    scalar_values = np.asarray(
        [time_to_hit_s, robot_source_monotonic_s, ball_source_monotonic_s],
        dtype=np.float64,
    )
    all_values = np.concatenate(
        [robot_pos, robot_quat, scalar_values, hit_pos, racket_velocity,
         incoming_velocity, landing_xy, origin]
    )
    if not np.isfinite(all_values).all():
        raise ValueError("H2 planner packet inputs must all be finite")
    if isinstance(command_id, bool) or int(command_id) != command_id or int(command_id) < 0:
        raise ValueError("H2 planner command_id must be a non-negative integer")
    quat_norm = float(np.linalg.norm(robot_quat))
    if not 0.5 <= quat_norm <= 1.5:
        raise ValueError(f"robot quaternion norm is invalid ({quat_norm:.3f})")
    if not allow_missing_sources and (robot_source_monotonic_s <= 0.0 or ball_source_monotonic_s <= 0.0):
        raise ValueError("H2 planner source timestamps must be positive monotonic seconds")
    if allow_missing_sources and (
        robot_source_monotonic_s < 0.0 or ball_source_monotonic_s < 0.0
        or (ball_source_monotonic_s == 0.0 and command_id != 0)
    ):
        raise ValueError("simulation missing-source packets use zero stamps and no strike before a ball")

    packet = np.empty(H2_PLANNER_PACKET_LEN, dtype=np.float64)
    packet[0:3] = robot_pos
    packet[3:7] = robot_quat
    packet[7] = float(time_to_hit_s)
    packet[8:11] = hit_pos
    packet[11:14] = racket_velocity
    packet[14] = float(command_id)
    packet[15:18] = incoming_velocity
    packet[18:20] = landing_xy
    packet[PROTOCOL_VERSION_INDEX] = H2_PLANNER_PROTOCOL_VERSION
    packet[ROBOT_SOURCE_MONOTONIC_INDEX] = float(robot_source_monotonic_s)
    packet[BALL_SOURCE_MONOTONIC_INDEX] = float(ball_source_monotonic_s)
    packet[ROBOT_BASE_ORIGIN_START_INDEX : ROBOT_BASE_ORIGIN_START_INDEX + 3] = origin
    packet[FRAME_CONVENTION_INDEX] = FRAME_CONVENTION_WORLD_X
    return packet


__all__ = [
    "BALL_SOURCE_MONOTONIC_INDEX",
    "FRAME_CONVENTION_INDEX",
    "FRAME_CONVENTION_WORLD_X",
    "H2_PLANNER_PACKET_LEN",
    "H2_PLANNER_PROTOCOL_VERSION",
    "LEGACY_FIELDS",
    "PROTOCOL_VERSION_INDEX",
    "ROBOT_SOURCE_MONOTONIC_INDEX",
    "ROBOT_BASE_ORIGIN_START_INDEX",
    "build_h2_planner_packet",
    "source_age_s",
]