"""test_helpers.py — Shared helpers for social-bot integration tests.

No hardware dependencies. Provides:
- TopicRateChecker: measure actual publish rate on a topic
- NodeChecker: verify a named node is alive in the graph
- ServiceChecker: call a service and verify it responds
- GpuMemoryChecker: measure GPU allocation via pycuda / nvidia-smi
- TimeoutPoller: generic poll-until-true with timeout
"""

from __future__ import annotations

import subprocess
import threading
import time
from typing import Any, Callable, List, Optional

import rclpy
from rclpy.node import Node
from rclpy.qos import (
    QoSProfile, QoSReliabilityPolicy, QoSDurabilityPolicy, QoSHistoryPolicy
)


# ── Timeout poller ────────────────────────────────────────────────────────────

def poll_until(
    condition: Callable[[], bool],
    timeout_s: float = 30.0,
    interval_s: float = 0.5,
) -> bool:
    """Return True if condition() becomes True within timeout_s, else False."""
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        if condition():
            return True
        time.sleep(interval_s)
    return False


# ── Node checker ──────────────────────────────────────────────────────────────

class NodeChecker:
    """Check whether ROS2 nodes are alive in the graph."""

    def __init__(self, probe_node: Node) -> None:
        self._node = probe_node

    def is_alive(self, node_name: str, namespace: str = "/") -> bool:
        """Return True if node_name is present in the node graph."""
        names_and_ns = self._node.get_node_names_and_namespaces()
        for name, ns in names_and_ns:
            if name == node_name:
                return True
            if ns.rstrip("/") + "/" + name == node_name:
                return True
        return False

    def wait_for_nodes(
        self,
        node_names: List[str],
        timeout_s: float = 30.0,
    ) -> dict:
        """Block until all nodes appear or timeout. Returns {name: found}."""
        results = {n: False for n in node_names}
        deadline = time.time() + timeout_s
        while time.time() < deadline:
            for name in node_names:
                if not results[name]:
                    results[name] = self.is_alive(name)
            if all(results.values()):
                break
            time.sleep(0.5)
        return results


# ── Topic rate checker ─────────────────────────────────────────────────────────

class TopicRateChecker:
    """Measure the actual publish rate on a topic over a measurement window.

    Usage:
        checker = TopicRateChecker(node, "/social/faces/detections",
                                   FaceDetectionArray, window_s=3.0)
        checker.start()
        time.sleep(5.0)
        hz = checker.measured_hz()
        checker.stop()
    """

    def __init__(
        self,
        node: Node,
        topic: str,
        msg_type: Any,
        window_s: float = 3.0,
        qos: Optional[QoSProfile] = None,
    ) -> None:
        self._window_s = window_s
        self._timestamps: List[float] = []
        self._lock = threading.Lock()
        self._running = False

        if qos is None:
            qos = QoSProfile(
                reliability=QoSReliabilityPolicy.BEST_EFFORT,
                durability=QoSDurabilityPolicy.VOLATILE,
                history=QoSHistoryPolicy.KEEP_LAST,
                depth=10,
            )

        self._sub = node.create_subscription(
            msg_type, topic, self._cb, qos
        )

    def _cb(self, _msg: Any) -> None:
        now = time.time()
        with self._lock:
            self._timestamps.append(now)
            # Prune old samples outside window
            cutoff = now - self._window_s * 2
            self._timestamps = [t for t in self._timestamps if t > cutoff]

    def measured_hz(self) -> float:
        """Return Hz measured over the last window_s seconds."""
        now = time.time()
        cutoff = now - self._window_s
        with self._lock:
            recent = [t for t in self._timestamps if t > cutoff]
        if len(recent) < 2:
            return 0.0
        return (len(recent) - 1) / (recent[-1] - recent[0])

    def message_count(self) -> int:
        """Total messages received."""
        with self._lock:
            return len(self._timestamps)

    def received_any(self) -> bool:
        with self._lock:
            return len(self._timestamps) > 0


# ── Service checker ────────────────────────────────────────────────────────────

class ServiceChecker:
    """Verify a ROS2 service is available and call it."""

    def __init__(self, node: Node) -> None:
        self._node = node

    def is_available(self, service_name: str, srv_type: Any,
                     timeout_s: float = 5.0) -> bool:
        """Return True if the service becomes available within timeout_s."""
        client = self._node.create_client(srv_type, service_name)
        available = client.wait_for_service(timeout_sec=timeout_s)
        self._node.destroy_client(client)
        return available

    def call_service(
        self,
        service_name: str,
        srv_type: Any,
        request: Any,
        timeout_s: float = 10.0,
    ) -> Optional[Any]:
        """Call a service and return the response, or None on failure."""
        import rclpy.executors
        client = self._node.create_client(srv_type, service_name)
        if not client.wait_for_service(timeout_sec=timeout_s):
            self._node.destroy_client(client)
            return None
        future = client.call_async(request)
        # Spin until done
        deadline = time.time() + timeout_s
        while not future.done() and time.time() < deadline:
            rclpy.spin_once(self._node, timeout_sec=0.1)
        self._node.destroy_client(client)
        if future.done():
            return future.result()
        return None


# ── GPU memory checker ────────────────────────────────────────────────────────

class GpuMemoryChecker:
    """Query GPU memory usage. Gracefully no-ops if no GPU available."""

    @staticmethod
    def total_mb() -> Optional[float]:
        """Total GPU memory in MB."""
        return GpuMemoryChecker._query("memory.total")

    @staticmethod
    def free_mb() -> Optional[float]:
        """Free GPU memory in MB."""
        return GpuMemoryChecker._query("memory.free")

    @staticmethod
    def used_mb() -> Optional[float]:
        """Used GPU memory in MB."""
        return GpuMemoryChecker._query("memory.used")

    @staticmethod
    def _query(field: str) -> Optional[float]:
        try:
            out = subprocess.check_output(
                ["nvidia-smi", f"--query-gpu={field}",
                 "--format=csv,noheader,nounits"],
                timeout=5,
            ).decode().strip()
            return float(out.split("\n")[0].strip())
        except Exception:
            pass
        # Fallback: pycuda
        try:
            import pycuda.driver as drv
            drv.init()
            ctx = drv.Device(0).make_context()
            free_b, total_b = drv.mem_get_info()
            ctx.pop()
            if field == "memory.free":
                return free_b / 1024 / 1024
            elif field == "memory.total":
                return total_b / 1024 / 1024
            elif field == "memory.used":
                return (total_b - free_b) / 1024 / 1024
        except Exception:
            pass
        return None

    @staticmethod
    def gpu_available() -> bool:
        return GpuMemoryChecker.total_mb() is not None


# ── Process helpers ───────────────────────────────────────────────────────────

def count_zombie_processes(pattern: str = "ros") -> int:
    """Count zombie processes matching pattern via ps."""
    try:
        out = subprocess.check_output(
            ["ps", "aux"], timeout=5
        ).decode()
        count = 0
        for line in out.splitlines():
            if "<defunct>" in line and pattern in line:
                count += 1
        return count
    except Exception:
        return 0


def get_process_pids(name_pattern: str) -> List[int]:
    """Return PIDs of processes matching name_pattern."""
    try:
        out = subprocess.check_output(
            ["pgrep", "-f", name_pattern], timeout=5
        ).decode().strip()
        return [int(p) for p in out.splitlines() if p.strip()]
    except Exception:
        return []
