"""test_helpers.py - Utilities for SaltyBot integration tests."""
import time
from typing import Dict, List, Tuple
import rclpy
from rclpy.node import Node
from tf2_ros import TransformException
from tf2_ros.buffer import Buffer
from tf2_ros.transform_listener import TransformListener


class NodeChecker:
    """Wait for nodes to appear in the ROS graph."""
    def __init__(self, node: Node):
        self.node = node

    def wait_for_nodes(self, node_names: List[str], timeout_s: float = 30.0) -> Dict[str, bool]:
        deadline = time.time() + timeout_s
        results = {n: False for n in node_names}
        while time.time() < deadline:
            names_in_graph = [n for n, _ in self.node.get_node_names()]
            for name in node_names:
                if name in names_in_graph:
                    results[name] = True
            if all(results.values()):
                return results
            time.sleep(0.2)
        return results


class TFChecker:
    """Verify TF tree completeness for SaltyBot."""
    REQUIRED_FRAMES = {
        "odom": "base_link",
        "base_link": "base_footprint",
        "base_link": "camera_link",
        "base_link": "lidar_link",
        "base_link": "imu_link",
    }

    def __init__(self, node: Node):
        self.node = node
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, node)

    def wait_for_tf_tree(self, timeout_s: float = 30.0, spin_func=None) -> Tuple[bool, List[str]]:
        deadline = time.time() + timeout_s
        missing = []
        while time.time() < deadline:
            missing = []
            for parent, child in self.REQUIRED_FRAMES.items():
                try:
                    self.tf_buffer.lookup_transform(parent, child, rclpy.time.Time())
                except TransformException:
                    missing.append(f"{parent} -> {child}")
            if not missing:
                return (True, [])
            if spin_func:
                spin_func(0.2)
            else:
                time.sleep(0.2)
        return (False, missing)


class Nav2Checker:
    """Check if Nav2 stack is active."""
    def __init__(self, node: Node):
        self.node = node

    def wait_for_nav2_active(self, timeout_s: float = 30.0) -> bool:
        deadline = time.time() + timeout_s
        while time.time() < deadline:
            node_names = [n for n, _ in self.node.get_node_names()]
            if "nav2_behavior_tree_executor" in node_names:
                topics = {name for name, _ in self.node.get_topic_names_and_types()}
                if "/nav2_behavior_tree/status" in topics:
                    return True
            time.sleep(0.2)
        return False


class TopicMonitor:
    """Monitor topic message counts."""
    def __init__(self, node: Node):
        self.node = node
        self.message_counts: Dict[str, int] = {}
        self.subscriptions = {}

    def subscribe_to_topic(self, topic_name: str, msg_type) -> None:
        self.message_counts[topic_name] = 0
        def callback(msg):
            self.message_counts[topic_name] += 1
        self.subscriptions[topic_name] = self.node.create_subscription(
            msg_type, topic_name, callback, 10
        )

    def get_message_count(self, topic_name: str) -> int:
        return self.message_counts.get(topic_name, 0)

    def cleanup(self) -> None:
        for sub in self.subscriptions.values():
            self.node.destroy_subscription(sub)
        self.subscriptions.clear()
        self.message_counts.clear()
