#!/usr/bin/env python3
# -- BEGIN LICENSE BLOCK ----------------------------------------------
# Copyright 2026 Universal Robots A/S
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#    * Redistributions of source code must retain the above copyright
#      notice, this list of conditions and the following disclaimer.
#
#    * Redistributions in binary form must reproduce the above copyright
#      notice, this list of conditions and the following disclaimer in the
#      documentation and/or other materials provided with the distribution.
#
#    * Neither the name of the {copyright_holder} nor the names of its
#      contributors may be used to endorse or promote products derived from
#      this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
# -- END LICENSE BLOCK ------------------------------------------------


"""Small helper script to start the tool communication interface."""

import subprocess
import logging
import argparse
import socket
import os
import pytest

RED = "\033[31m"
RESET = "\033[0m"


# Custom formatter to show both default values and description formatting in the help message
class Formatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter):
    pass


def get_args():
    # Arguments to configure socat
    arg = argparse.ArgumentParser(
        description=(
            "Starts socat to create a PTY symlink for the UR tool communication interface."
        ),
        epilog="""
IMPORTANT:

This script requires the ToolComm Forwarder URCap to be running on the robot.
Make sure it is installed and started before launching this script.

More information can be found in the following ToolComm Forwarder URCap repositories:
  - ToolComm Forwarder URCap (Polyscope X):
    https://github.com/UniversalRobots/Universal_Robots_ToolComm_Forwarder_URCapX
  - ToolComm Forwarder URCap (Polyscope 5)
    https://github.com/UniversalRobots/Universal_Robots_ToolComm_Forwarder_URCap

For background information on how tool communication works on UR robots, see:
https://docs.universal-robots.com/Universal_Robots_ROS2_Documentation/doc/ur_robot_driver/ur_robot_driver/doc/setup_tool_communication.html
        """,
        formatter_class=Formatter
    )

    arg.add_argument("robot_ip", help="IP address of the robot to connect to.")
    arg.add_argument("--tcp-port", type=int, default=54321, help="TCP Port. Likely, this should not be changed")
    arg.add_argument("--device-name", default="/tmp/ttyUR", help="PTY symlink device name.")
    return arg.parse_args()


def check_tcp(ip, port, timeout=5.0):
    try:
        with socket.create_connection((ip, port), timeout=timeout):
            return True
    except OSError:
        return False


def main(args):

    logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")

    # Get parameters from arguments
    robot_ip = args.robot_ip
    logging.info("Robot IP: " + robot_ip)
    tcp_port = args.tcp_port
    logging.info("TCP Port: " + str(tcp_port))
    local_device = args.device_name

    # Check IP and port reachability
    if not check_tcp(robot_ip, tcp_port):
        raise ConnectionError(
            f"Cannot reach {robot_ip}:{tcp_port}.\n"
            "Check that the IP address and port are correct.\n"
            "If so, ensure that the robot is powered on, reachable on the network, "
            f"and that the ToolCommForwarder URCap is running."
        )

    # Check if parent directory of device_name exists
    parent_dir = os.path.dirname(local_device)
    if parent_dir and not os.path.exists(parent_dir):
        raise FileNotFoundError(
            f"Parent directory '{parent_dir}' does not exist.\n"
            "Socat needs an existing directory to create the PTY symlink.\n"
            "Fix:\n"
            f"  - Create the parent directory, e.g. 'mkdir -p {parent_dir}'.\n"
            f"  - Use a different device name with an existing parent directory."
        )

    # Check if the device_name is a directory
    if os.path.isdir(local_device):
        raise FileExistsError(
            f"'{local_device}' exists and is a directory.\n"
            "Socat needs a file path to create a PTY symlink, but it cannot replace a directory.\n"
            "Fix:\n"
            "  - Remove the directory.\n"
            f"  - Use a different device name, e.g. '--device-name /tmp/ttyUR0'."
                                 )

    # Configure socat command
    socat_config = [
        "pty",
        f"link={local_device}",
        "raw",
        "ignoreeof",
        "waitslave",
    ]

    socat_command = [
        "socat",
        ",".join(socat_config),
        f"tcp:{robot_ip}:{tcp_port}",
    ]

    logging.info(f"Configuring PTY symlink at '{local_device}'")

    # Start socat
    try:
        logging.info("Starting socat with following command:\n" + " ".join(socat_command))
        subprocess.call(socat_command)
        logging.info("Socat terminated")

    # Error case when socat is not installed
    except FileNotFoundError:
        raise FileNotFoundError("Socat not found in PATH. Install it (e.g. apt-get install socat).")


def test_check_tcp():
    assert not check_tcp("127.0.0.1", 0.1)


def test_check_tcp_open_port():
    bind_ip = "127.0.0.1"
    bind_port = 54321
    args = argparse.Namespace(
        robot_ip=bind_ip,
        tcp_port=bind_port,
        device_name="/tmp/nonexistent_dir/ttyUR"
    )
    with pytest.raises(ConnectionError):
        main(args)


def test_parent_dir_doesnt_exist():
    bind_ip = "127.0.0.1"
    bind_port = 54321
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind((bind_ip, bind_port))
    server.listen(1)
    args = argparse.Namespace(
        robot_ip=bind_ip,
        tcp_port=bind_port,
        device_name="/tmp/nonexistent_dir/ttyUR"
    )
    with pytest.raises(FileNotFoundError) as exc_info:
        main(args)
    assert os.path.dirname(args.device_name) in str(exc_info.value)
    server.close()


def test_device_name_is_directory():
    bind_ip = "127.0.0.1"
    bind_port = 54321
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind((bind_ip, bind_port))
    server.listen(1)
    args = argparse.Namespace(
        robot_ip=bind_ip,
        tcp_port=bind_port,
        device_name="/tmp/ttyUR_dir"
    )
    os.makedirs(args.device_name, exist_ok=True)
    with pytest.raises(FileExistsError) as exc_info:
        main(args)
    assert args.device_name in str(exc_info.value)
    server.close()


if __name__ == "__main__":
    args = get_args()
    try:
        main(args)
    except Exception as e:
        logging.error(f"{RED}Unexpected error: {e} {RESET}")
        logging.info("Exiting tool communication script.")
