#!/usr/bin/env python3
"""
ROS2 Bridge for FlowTorch Basin Control
Enables robot control via multi-state steering.

Example: Trit control for robot arm (retreat/hold/advance)
"""

import requests
import json
from typing import Optional

# ROS2 imports (conditional - only if ROS2 installed)
try:
    import rclpy
    from rclpy.node import Node
    from std_msgs.msg import String
    ROS2_AVAILABLE = True
except ImportError:
    ROS2_AVAILABLE = False
    print("ROS2 not available - install: pip install rclpy std_msgs")


class FlowTorchBridge(Node):
    """
    ROS2 node bridging FlowTorch basin control to robot commands.

    Subscribed topics:
        /flowtorch/command (String): State letter command ('a', 'b', 'c', etc.)

    Published topics:
        /flowtorch/receipt (String): Signed receipt JSON
        /flowtorch/basin (String): Basin measurement (phi, R, MI)
    """

    def __init__(
        self,
        flowtorch_url: str = "http://127.0.0.1:8089",
        cardinality: int = 3,  # Default: Trit (retreat/hold/advance)
    ):
        super().__init__('flowtorch_bridge')

        self.flowtorch_url = flowtorch_url
        self.cardinality = cardinality

        # Subscribers
        self.command_sub = self.create_subscription(
            String,
            '/flowtorch/command',
            self.command_callback,
            10
        )

        # Publishers
        self.receipt_pub = self.create_publisher(String, '/flowtorch/receipt', 10)
        self.basin_pub = self.create_publisher(String, '/flowtorch/basin', 10)

        self.get_logger().info(f"FlowTorch bridge ready (n={cardinality})")
        self.get_logger().info(f"FlowTorch URL: {flowtorch_url}")
        self.get_logger().info(f"Subscribed to: /flowtorch/command")
        self.get_logger().info(f"Publishing to: /flowtorch/receipt, /flowtorch/basin")

    def command_callback(self, msg: String):
        """
        Process state letter command from ROS2.

        Args:
            msg: String message with state letter ('a', 'b', 'c', etc.)
        """
        state_letter = msg.data.strip().lower()

        self.get_logger().info(f"Received command: {state_letter}")

        # Validate state letter
        valid_letters = {
            2: ['a', 'b'],
            3: ['a', 'b', 'c'],
            4: ['a', 'b', 'c', 'd'],
            8: ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'],
        }

        if state_letter not in valid_letters[self.cardinality]:
            self.get_logger().error(
                f"Invalid state '{state_letter}' for n={self.cardinality}. "
                f"Valid: {valid_letters[self.cardinality]}"
            )
            return

        # Call FlowTorch API
        try:
            response = requests.post(
                f"{self.flowtorch_url}/v1/chat/completions",
                json={
                    "model": "flowtorch/basin-control",
                    "messages": [
                        {"role": "user", "content": f"Execute state {state_letter}"}
                    ],
                    "max_tokens": 16,
                    "quant": {
                        "cardinality": self.cardinality,
                        "state_letter": state_letter,
                        "control_consent": True,  # Required for basin control
                    }
                },
                timeout=10.0
            )

            response.raise_for_status()
            data = response.json()

            # Extract basin control metadata
            if "receipt" in data:
                receipt = data["receipt"]
                basin = receipt["basin"]

                # Publish receipt
                receipt_msg = String()
                receipt_msg.data = json.dumps(receipt, indent=2)
                self.receipt_pub.publish(receipt_msg)

                # Publish basin measurement
                basin_msg = String()
                basin_msg.data = json.dumps({
                    "state": state_letter,
                    "phi_deg": basin["phi_deg"],
                    "R_spectral": basin["R_spectral"],
                    "MI_bits": basin["MI_bits"],
                    "energy_uJ": receipt["energy"]["uJ"],
                }, indent=2)
                self.basin_pub.publish(basin_msg)

                self.get_logger().info(
                    f"Basin lock: φ={basin['phi_deg']:.1f}°, "
                    f"R={basin['R_spectral']:.3f}, "
                    f"MI={basin['MI_bits']:.3f} bits"
                )
            else:
                self.get_logger().warning("No receipt in response")

        except requests.exceptions.RequestException as e:
            self.get_logger().error(f"FlowTorch request failed: {e}")
        except Exception as e:
            self.get_logger().error(f"Unexpected error: {e}")


def main(args=None):
    """Main entry point for ROS2 node."""
    if not ROS2_AVAILABLE:
        print("ERROR: ROS2 not available. Install: pip install rclpy std_msgs")
        return

    rclpy.init(args=args)

    # Create bridge node
    bridge = FlowTorchBridge(
        flowtorch_url="http://127.0.0.1:8089",
        cardinality=3,  # Trit: retreat/hold/advance
    )

    # Spin
    try:
        rclpy.spin(bridge)
    except KeyboardInterrupt:
        pass
    finally:
        bridge.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
