#!/usr/bin/env python3
"""ROS Action Client for Agent ROS Bridge.

Supports ROS Actions (ROS1 actionlib, ROS2 rclpy actions) for long-running
tasks like navigation, motion planning, and manipulation.

Features:
- Send action goals with timeouts
- Receive feedback during execution
- Get final results
- Cancel/preempt ongoing actions
- Support for both ROS1 and ROS2

Usage:
    from agent_ros_bridge.actions import ActionClient

    client = ActionClient("navigate_to_pose", "nav2_msgs/NavigateToPose")
    result = await client.send_goal({"pose": {...}}, feedback_cb=on_feedback)
"""

import asyncio
import contextlib
import logging
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum, auto
from typing import Any

logger = logging.getLogger("ros_actions")


class ActionStatus(Enum):
    """Action client status."""

    IDLE = auto()
    PENDING = auto()  # Goal sent, waiting for acceptance
    ACTIVE = auto()  # Goal accepted and executing
    PREEMPTING = auto()  # Cancel requested
    SUCCEEDED = auto()  # Goal completed successfully
    ABORTED = auto()  # Goal aborted by server
    REJECTED = auto()  # Goal rejected by server
    CANCELED = auto()  # Goal canceled successfully
    LOST = auto()  # Action server lost


@dataclass
class ActionGoal:
    """Action goal definition."""

    goal_id: str
    goal_data: dict[str, Any]
    action_type: str  # e.g., "nav2_msgs/action/NavigateToPose"
    timeout_sec: float = 30.0
    created_at: datetime = field(default_factory=datetime.utcnow)


@dataclass
class ActionFeedback:
    """Action feedback during execution."""

    goal_id: str
    feedback_data: dict[str, Any]
    timestamp: datetime = field(default_factory=datetime.utcnow)


@dataclass
class ActionResult:
    """Action final result."""

    goal_id: str
    success: bool
    status: ActionStatus
    result_data: dict[str, Any] = field(default_factory=dict)
    error_message: str | None = None
    execution_time_sec: float = 0.0


class BaseActionClient:
    """Base class for ROS action clients."""

    def __init__(self, action_name: str, action_type: str, ros_version: str = "ros2"):
        """Initialize action client with name and type.

        Args:
            action_name: Name of the ROS action.
            action_type: Type of the action (e.g., "nav2_msgs/action/NavigateToPose").
            ros_version: ROS version ("ros1" or "ros2").
        """
        self.action_name = action_name
        self.action_type = action_type
        self.ros_version = ros_version
        self.status = ActionStatus.IDLE
        self._feedback_callbacks: list[Callable[[ActionFeedback], None]] = []
        self._result_callbacks: list[Callable[[ActionResult], None]] = []
        self._current_goal: ActionGoal | None = None
        self._start_time: datetime | None = None
        self._client: Any | None = None
        self._connected = False

    async def connect(self) -> bool:
        """Connect to action server."""
        raise NotImplementedError

    async def disconnect(self):
        """Disconnect from action server."""
        raise NotImplementedError

    async def send_goal(self, goal_data: dict[str, Any], timeout_sec: float = 30.0) -> ActionResult:
        """Send goal to action server."""
        raise NotImplementedError

    async def cancel_goal(self) -> bool:
        """Cancel current goal."""
        raise NotImplementedError

    def register_feedback_callback(self, callback: Callable[[ActionFeedback], None]):
        """Register callback for feedback updates."""
        self._feedback_callbacks.append(callback)

    def register_result_callback(self, callback: Callable[[ActionResult], None]):
        """Register callback for results."""
        self._result_callbacks.append(callback)

    def _notify_feedback(self, feedback: ActionFeedback):
        """Notify all feedback callbacks."""
        for cb in self._feedback_callbacks:
            try:
                cb(feedback)
            except Exception as e:
                logger.error(f"Feedback callback error: {e}")

    def _notify_result(self, result: ActionResult):
        """Notify all result callbacks."""
        for cb in self._result_callbacks:
            try:
                cb(result)
            except Exception as e:
                logger.error(f"Result callback error: {e}")


class ROS2ActionClient(BaseActionClient):
    """ROS2 action client using rclpy."""

    def __init__(self, action_name: str, action_type: str):
        """Initialize ROS2 action client.

        Args:
            action_name: Name of the ROS2 action.
            action_type: Type of the action.
        """
        super().__init__(action_name, action_type, "ros2")
        self._goal_handle: Any | None = None

    async def connect(self) -> bool:
        """Connect to ROS2 action server."""
        try:
            import rclpy
            from rclpy.action import ActionClient
            from rclpy.node import Node

            # Get or create node
            with contextlib.suppress(Exception):
                self._node = rclpy.create_node(
                    f"action_client_{self.action_name.replace('/', '_')}"
                )

            # Create action client
            self._client = ActionClient(self._node, self.action_type, self.action_name)

            # Wait for server
            if not self._client.wait_for_server(timeout_sec=5.0):
                logger.error(f"Action server '{self.action_name}' not available")
                return False

            self._connected = True
            logger.info(f"✅ Connected to ROS2 action: {self.action_name}")
            return True

        except ImportError:
            logger.error("rclpy not available")
            return False
        except Exception as e:
            logger.error(f"Failed to connect to action server: {e}")
            return False

    async def disconnect(self):
        """Disconnect from action server."""
        if self._client:
            self._client.destroy()
        self._connected = False

    async def send_goal(self, goal_data: dict[str, Any], timeout_sec: float = 30.0) -> ActionResult:
        """Send goal to ROS2 action server."""
        if not self._connected:
            return ActionResult(
                goal_id="",
                success=False,
                status=ActionStatus.LOST,
                error_message="Not connected to action server",
            )

        import uuid

        goal_id = str(uuid.uuid4())[:8]

        goal = ActionGoal(
            goal_id=goal_id,
            goal_data=goal_data,
            action_type=self.action_type,
            timeout_sec=timeout_sec,
        )

        self._current_goal = goal
        self._start_time = datetime.utcnow()
        self.status = ActionStatus.PENDING

        try:
            # Create goal message
            goal_msg = self._create_goal_message(goal_data)

            # Send goal
            if self._client is None:
                return ActionResult(
                    goal_id=goal_id,
                    success=False,
                    status=ActionStatus.LOST,
                    error_message="Client not connected",
                )
            self._goal_handle = await self._client.send_goal_async(
                goal_msg, feedback_callback=self._on_feedback
            )

            if self._goal_handle is None or not self._goal_handle.accepted:
                self.status = ActionStatus.REJECTED
                return ActionResult(
                    goal_id=goal_id,
                    success=False,
                    status=ActionStatus.REJECTED,
                    error_message="Goal rejected by server",
                )

            self.status = ActionStatus.ACTIVE
            logger.info(f"▶️  Action goal accepted: {goal_id}")

            # Wait for result with timeout
            if self._goal_handle is None:
                return ActionResult(
                    goal_id=goal_id,
                    success=False,
                    status=ActionStatus.LOST,
                    error_message="Goal handle is None",
                )
            result_future = self._goal_handle.get_result_async()

            try:
                result_response = await asyncio.wait_for(
                    asyncio.wrap_future(result_future), timeout=timeout_sec
                )

                execution_time = (datetime.utcnow() - self._start_time).total_seconds()

                # Check result status
                if result_response.status == 4:  # SUCCEEDED
                    self.status = ActionStatus.SUCCEEDED
                    result = ActionResult(
                        goal_id=goal_id,
                        success=True,
                        status=ActionStatus.SUCCEEDED,
                        result_data=self._parse_result(result_response.result),
                        execution_time_sec=execution_time,
                    )
                elif result_response.status == 5:  # ABORTED
                    self.status = ActionStatus.ABORTED
                    result = ActionResult(
                        goal_id=goal_id,
                        success=False,
                        status=ActionStatus.ABORTED,
                        error_message="Goal aborted by server",
                        execution_time_sec=execution_time,
                    )
                elif result_response.status == 2:  # CANCELED
                    self.status = ActionStatus.CANCELED
                    result = ActionResult(
                        goal_id=goal_id,
                        success=False,
                        status=ActionStatus.CANCELED,
                        execution_time_sec=execution_time,
                    )
                else:
                    self.status = ActionStatus.ABORTED
                    result = ActionResult(
                        goal_id=goal_id,
                        success=False,
                        status=ActionStatus.ABORTED,
                        error_message=f"Unknown status: {result_response.status}",
                        execution_time_sec=execution_time,
                    )

                self._notify_result(result)
                return result

            except TimeoutError:
                self.status = ActionStatus.ABORTED
                await self.cancel_goal()
                return ActionResult(
                    goal_id=goal_id,
                    success=False,
                    status=ActionStatus.ABORTED,
                    error_message=f"Timeout after {timeout_sec}s",
                    execution_time_sec=timeout_sec,
                )

        except Exception as e:
            logger.error(f"Action execution failed: {e}")
            self.status = ActionStatus.ABORTED
            return ActionResult(
                goal_id=goal_id, success=False, status=ActionStatus.ABORTED, error_message=str(e)
            )

    async def cancel_goal(self) -> bool:
        """Cancel current goal."""
        if self._goal_handle and self.status == ActionStatus.ACTIVE:
            cancel_future = self._goal_handle.cancel_goal_async()
            await asyncio.wrap_future(cancel_future)
            self.status = ActionStatus.CANCELED
            return True
        return False

    def _on_feedback(self, feedback_msg):
        """Handle feedback from action server."""
        feedback_data = self._parse_feedback(feedback_msg)

        feedback = ActionFeedback(
            goal_id=self._current_goal.goal_id if self._current_goal else "",
            feedback_data=feedback_data,
        )

        self._notify_feedback(feedback)

    def _create_goal_message(self, goal_data: dict[str, Any]):
        """Create goal message from dict (simplified)."""

        # This would need proper message type handling
        # For now, return the dict wrapped
        class GoalMsg:
            def __init__(self, data):
                for k, v in data.items():
                    setattr(self, k, v)

        return GoalMsg(goal_data)

    def _parse_feedback(self, feedback_msg) -> dict[str, Any]:
        """Parse feedback message to dict."""
        # Extract attributes from feedback message
        result = {}
        for attr in dir(feedback_msg):
            if not attr.startswith("_"):
                try:
                    val = getattr(feedback_msg, attr)
                    if not callable(val):
                        result[attr] = val
                except Exception:
                    pass
        return result

    def _parse_result(self, result_msg) -> dict[str, Any]:
        """Parse result message to dict."""
        return self._parse_feedback(result_msg)


class SimulatedActionClient(BaseActionClient):
    """Simulated action client for testing without ROS."""

    async def connect(self) -> bool:
        """Simulated connect."""
        self._connected = True
        logger.info(f"✅ Simulated action client: {self.action_name}")
        return True

    async def disconnect(self):
        """Simulated disconnect."""
        self._connected = False

    async def send_goal(
        self, goal_data: dict[str, Any], _timeout_sec: float = 30.0
    ) -> ActionResult:
        """Simulated goal execution with feedback."""
        import uuid

        goal_id = str(uuid.uuid4())[:8]

        self._current_goal = ActionGoal(
            goal_id=goal_id, goal_data=goal_data, action_type=self.action_type
        )

        self.status = ActionStatus.ACTIVE
        self._start_time = datetime.utcnow()

        logger.info(f"▶️  Simulated action: {self.action_name} goal {goal_id}")

        # Simulate execution with feedback
        steps = 10
        for i in range(steps):
            await asyncio.sleep(0.5)

            # Generate simulated feedback based on action type
            feedback_data = self._generate_simulated_feedback(i, steps, goal_data)

            feedback = ActionFeedback(goal_id=goal_id, feedback_data=feedback_data)
            self._notify_feedback(feedback)

            logger.debug(f"Simulated feedback: {feedback_data}")

        # Complete successfully
        execution_time = (datetime.utcnow() - self._start_time).total_seconds()
        self.status = ActionStatus.SUCCEEDED

        result = ActionResult(
            goal_id=goal_id,
            success=True,
            status=ActionStatus.SUCCEEDED,
            result_data={"message": "Simulated action completed"},
            execution_time_sec=execution_time,
        )

        self._notify_result(result)
        logger.info(f"✅ Simulated action completed: {goal_id}")

        return result

    async def cancel_goal(self) -> bool:
        """Simulated cancel."""
        self.status = ActionStatus.CANCELED
        return True

    def _generate_simulated_feedback(
        self, step: int, total: int, _goal_data: dict
    ) -> dict[str, Any]:
        """Generate simulated feedback based on action type."""
        progress = (step + 1) / total

        if "navigate" in self.action_name.lower():
            return {
                "distance_remaining": 10.0 * (1 - progress),
                "current_speed": 0.5,
                "estimated_time_remaining": 10.0 * (1 - progress),
                "progress": progress,
            }
        elif "manipulate" in self.action_name.lower() or "move" in self.action_name.lower():
            return {"current_joint": step, "progress": progress, "status": "moving"}
        else:
            return {"progress": progress, "step": step, "status": "active"}


def create_action_client(
    action_name: str, action_type: str, ros_version: str = "ros2"
) -> BaseActionClient:
    """Factory function to create appropriate action client."""
    # Try to use real ROS client
    if ros_version == "ros2":
        try:
            import rclpy

            return ROS2ActionClient(action_name, action_type)
        except ImportError:
            logger.warning("rclpy not available, using simulated action client")
            return SimulatedActionClient(action_name, action_type)
    else:
        # ROS1 not yet implemented
        logger.warning("ROS1 action client not implemented, using simulated")
        return SimulatedActionClient(action_name, action_type)
