#!/usr/bin/env python3
"""Add measured racket positions without replacing the selected ball estimator."""
from __future__ import annotations

import argparse
import importlib.util
from pathlib import Path
import sys


def load_adapter(root):
    path = Path(root) / "scripts/mocap_ego_zmq_adapter_physics.py"
    spec = importlib.util.spec_from_file_location("h2_original_mocap_adapter", path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def load_v2_adapter(root):
    # Keep the upstream's runtime package separate from the deployment runtime.
    root = Path(root).resolve()
    spec = importlib.util.spec_from_file_location(
        "_h2_eval_v2_adapter", root / "run_adapter.py",
        submodule_search_locations=[str(root)])
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def racket_adapter_class(base, adapter_class=None):
    class RacketAdapter(adapter_class or base.MocapEgoZmqAdapter):
        def __init__(self):
            super().__init__()
            with self._payload_cv:
                self._latest_payloads["racket"] = None
            self.declare_parameter("racket_pose_topic", "/test_pat1/pose")
            topic = self.get_parameter("racket_pose_topic").value
            self._racket_subscription = self.create_subscription(
                base.PoseStamped, topic, self._racket_callback, base.make_sensor_qos())
            self.get_logger().info(f"Racket center: {topic}; same coordinates/units as ball")

        def _racket_callback(self, message):
            sample = self._pose_to_reference_frame(base.pose_msg_to_sample(message))
            with self._state_lock:
                position = self._position_to_output_table_center(sample.pos)
            self._queue_payload(dict(type="racket", t=sample.t, unit=self._payload_unit,
                                     pos=list(self._vector_from_meters_to_payload_unit(position))))

        def _take_pending_snapshot(self):
            snapshot = super()._take_pending_snapshot()
            if snapshot is not None:
                with self._payload_cv:
                    snapshot["racket"] = self._latest_payloads.get("racket")
            return snapshot

    return RacketAdapter


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    roots = parser.add_mutually_exclusive_group(required=True)
    roots.add_argument("--adapter-root")
    roots.add_argument("--upstream-v2-root")
    args, ros_args = parser.parse_known_args()
    if args.upstream_v2_root:
        entry = load_v2_adapter(args.upstream_v2_root)
        entry.main(ros_args, adapter_class_factory=racket_adapter_class)
        return
    base = load_adapter(args.adapter_root)
    base.rclpy.init(args=ros_args)
    node = racket_adapter_class(base)()
    try:
        base.rclpy.spin(node)
    except (KeyboardInterrupt, base.ExternalShutdownException):
        pass
    finally:
        node.destroy_node()
        if base.rclpy.ok():
            base.rclpy.shutdown()


if __name__ == "__main__":
    main()