# -*- coding: UTF-8 -*-
# Improved Fairino RoboDK driver based on the RoboDK draft driver received from support.
# Changes:
# - Robust joint parsing for different frrpc return formats
# - Fixed MOVL dispatch (was incorrectly calling MoveJ)
# - Safer connect/disconnect/stop handling
# - Safer Connected() check
# - Fixed DriverManager.Stop queue usage
# - Better diagnostics on connection/joint parsing failures

DRIVER_VERSION = "Fairino RoboDK Driver v0.0.2-patched"

import sys
import atexit
import threading
from queue import Queue
import time

global LOCK_STDOUT
LOCK_STDOUT = threading.Lock()


def print_status(status):
    global LOCK_STDOUT
    with LOCK_STDOUT:
        print("SMS:" + status)
        sys.stdout.flush()


def show_message(message):
    global LOCK_STDOUT
    with LOCK_STDOUT:
        print("SMS2:" + message)
        sys.stdout.flush()


def print_message(message):
    global LOCK_STDOUT
    with LOCK_STDOUT:
        print("DBG:" + message)
        sys.stdout.flush()


def print_response(response):
    global LOCK_STDOUT
    with LOCK_STDOUT:
        print("RE:" + response)
        sys.stdout.flush()


def print_commands(commands):
    global LOCK_STDOUT
    with LOCK_STDOUT:
        print("CMDLIST:" + commands)
        sys.stdout.flush()


def print_joints(joints, is_moving=False):
    global LOCK_STDOUT
    if joints is None:
        joints = []
    if not isinstance(joints, (list, tuple)):
        joints = [joints]
    with LOCK_STDOUT:
        if is_moving:
            print("JNTS_MOVING: " + " ".join(format(float(x), ".3f") for x in joints))
        else:
            print("JNTS: " + " ".join(format(float(x), ".6f") for x in joints))
        sys.stdout.flush()


try:
    import frrpc
except Exception:
    frrpc = None
    print_status("frrpc is not installed, please install it.")


ROBOTCOM_UNKNOWN = -1000
ROBOTCOM_CONNECTION_PROBLEMS = -3
ROBOTCOM_DISCONNECTED = -2
ROBOTCOM_NOT_CONNECTED = -1
ROBOTCOM_READY = 0
ROBOTCOM_WORKING = 1
ROBOTCOM_WAITING = 2

global STATUS
STATUS = ROBOTCOM_DISCONNECTED


def UpdateStatus(set_status=None):
    global STATUS
    if set_status is not None:
        STATUS = set_status

    if STATUS == ROBOTCOM_CONNECTION_PROBLEMS:
        print_status("Connection problems")
    elif STATUS == ROBOTCOM_DISCONNECTED:
        print_status("Disconnected")
    elif STATUS == ROBOTCOM_NOT_CONNECTED:
        print_status("Not connected")
    elif STATUS == ROBOTCOM_READY:
        print_status("Ready")
    elif STATUS == ROBOTCOM_WORKING:
        print_status("Working...")
    elif STATUS == ROBOTCOM_WAITING:
        print_status("Waiting...")
    else:
        print_status("Unknown status")


RDK_COMMANDS = {
    'CONNECT',
    'DISCONNECT',
    'STOP',
    'QUIT',
    'PAUSE',
    'MOVJ',
    'MOVL',
    'MOVLSEARCH',
    'MOVC',
    'CJNT',
    'SPEED',
    'SETROUNDING',
    'SETDO',
    'SETAO',
    'WAITDI',
    'SETTOOL',
    'RUNPROG',
    'POPUP',
    'SUPPORTSEARCH',
    'SEARCHDI',
}

DEBUG_COMMANDS = ['t', 'c', 'd', 's', 'q']

DRIVER_RESPONSES = {
    'SMS': None,
    'SMS2': None,
    'DBG': None,
    'RE': None,
    'CMDLIST': None,
}


def str_to_float(str_value):
    try:
        if '.' in str_value:
            return float(str_value)
        else:
            return int(str_value)
    except ValueError:
        return str_value


def command_parser(cmd_line):
    if not cmd_line.strip():
        return '', []

    args = [arg.strip() for arg in cmd_line.split(' ')]
    if not args:
        return '', []

    cmd = args[0]
    args = [str_to_float(arg) for arg in args[1:]]
    return cmd, args


def command_runner(cmd, args, robot: 'RobotDriver'):
    if cmd == 'c' and args:
        robot.RunCode(' '.join(str(a) for a in args))
        return

    if cmd not in RDK_COMMANDS:
        return

    if cmd not in ['CONNECT', 'DISCONNECT', 'QUIT'] and not robot.Connected():
        UpdateStatus(ROBOTCOM_NOT_CONNECTED)
        return

    if cmd == 'CONNECT':
        UpdateStatus(ROBOTCOM_WAITING)

        success = False
        if not args:
            success = robot.ConnectDefault()
        else:
            ip_or_com = args[0]
            port = args[1] if len(args) > 1 else -1
            ndof = args[2] if len(args) > 2 else 6
            success = robot.Connect(ip_or_com, port, ndof)

        if not success:
            UpdateStatus(ROBOTCOM_CONNECTION_PROBLEMS)
        else:
            joints = robot.Joints()
            print_joints(joints)
            UpdateStatus(ROBOTCOM_READY)

    elif cmd == 'DISCONNECT':
        robot.Disconnect()
        UpdateStatus(ROBOTCOM_DISCONNECTED)

    elif cmd in ['QUIT']:
        UpdateStatus(ROBOTCOM_WORKING)
        robot.Stop()
        robot.Disconnect()
        UpdateStatus(ROBOTCOM_DISCONNECTED)

    elif cmd in ['STOP']:
        UpdateStatus(ROBOTCOM_WORKING)
        robot.Stop()
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ["SUPPORTSEARCH"]:
        print_response("1" if robot.hasLinearSearch() else "0")
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ["SEARCHDI"]:
        print_message("SEARCHDI cmd not implemented yet")
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['MOVLSEARCH']:
        UpdateStatus(ROBOTCOM_WORKING)
        print_message("MOVLSEARCH cmd not implemented yet")
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['MOVJ', 'MOVL']:
        UpdateStatus(ROBOTCOM_WORKING)

        ndof = len(args) - 6
        joints = args[:ndof]
        xyzwpr = args[-6:]

        if cmd == 'MOVJ':
            robot.MoveJ(xyzwpr, joints)
        elif cmd == 'MOVL':
            robot.MoveL(xyzwpr, joints)

        joints = robot.Joints()
        print_joints(joints)

        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['MOVC']:
        UpdateStatus(ROBOTCOM_WORKING)
        print_message("MOVC cmd not implemented yet")
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['SPEED']:
        lin_vel, joint_vel, lin_accel, joint_accel = args[:4]
        if lin_vel >= 0:
            robot.setSpeed(lin_vel)
        if joint_vel >= 0:
            print_message("SPEEDJOINTS cmd not implemented yet")
        if lin_accel >= 0:
            robot.setAcceleration(lin_accel)
        if joint_accel >= 0:
            print_message("ACCJOINTS cmd not implemented yet")
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['CJNT']:
        joints = robot.Joints()
        print_joints(joints)
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['PAUSE']:
        UpdateStatus(ROBOTCOM_WORKING)
        time_ms = args[0]
        robot.Pause(time_ms)
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['SETROUNDING']:
        print_message("SETROUNDING cmd not implemented yet")
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['SETDO']:
        do_name, do_value, var_name, var_value = args[:4]
        robot.setDO(var_name, var_value)
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['WAITDI']:
        di_name, di_value, var_name, var_value = args[:4]
        robot.waitDI(var_name, var_value)
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['SETTOOL']:
        print_message("SETTOOL cmd not implemented yet")
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['RUNPROG']:
        UpdateStatus(ROBOTCOM_WORKING)
        print_message("RUNPROG cmd not implemented yet")
        UpdateStatus(ROBOTCOM_READY)

    elif cmd in ['POPUP']:
        UpdateStatus(ROBOTCOM_WORKING)
        print_message("POPUP cmd not implemented yet")
        UpdateStatus(ROBOTCOM_READY)


class DriverManager:

    def __init__(self, robot):
        self.commands_queue = Queue()
        self.thread_driver = None
        self.thread_console = None
        self.robot = robot

    def ConsoleRunner(self):
        for line in sys.stdin:
            cmd, args = command_parser(line)
            if not cmd:
                continue

            if cmd in ['QUIT']:
                with self.commands_queue.mutex:
                    self.commands_queue.queue.clear()

            self.commands_queue.put((cmd, args))

            if cmd in ['QUIT']:
                return

    def DriverRunner(self):
        while True:
            if self.commands_queue.empty():
                time.sleep(0.005)
                continue

            cmd, args = self.commands_queue.get()
            if cmd == 't':
                self.TestDriver()
                continue

            if cmd == 'c' and not args:
                cmd = 'CONNECT'
            elif cmd == 'd' and not args:
                cmd = 'DISCONNECT'
            elif cmd == 's' and not args:
                cmd = 'STOP'
            elif cmd == 'q' and not args:
                cmd = 'QUIT'

            command_runner(cmd, args, self.robot)

            if cmd in ['QUIT']:
                return

    def Run(self):
        if self.thread_driver is not None:
            return

        print_commands(self.robot.CustomCommandList())

        self.thread_console = threading.Thread(target=self.ConsoleRunner, name="ConsoleRunner", args=())
        self.thread_console.daemon = True
        self.thread_console.start()

        self.DriverRunner()

    def Stop(self):
        self.commands_queue.put(('QUIT', []))
        if self.thread_console is not None and self.thread_console.is_alive():
            self.thread_console.join(5)
            self.thread_console = None

        if self.thread_driver is not None and self.thread_driver.is_alive():
            self.thread_driver.join(5)
            self.thread_driver = None

    def TestDriver(self):
        commands = [
            "CONNECT",
            "SPEED 10 -1 -1 -1",
            "SETTOOL -0.025 -41.046 50.920 60.000 -0.000 90.000",
            "CJNT",
            "MOVJ 0 -5.362010 46.323420 20.746290 74.878840 -50.101680 61.958500",
            "MOVJ 500 -5.362010 46.323420 20.746290 74.878840 -50.101680 61.958500",
            "SPEED 100 -1 -1 -1",
            "MOVL 0 -5.362010 50.323420 20.746290 74.878840 -50.101680 61.958500",
            "MOVL 500 -5.362010 50.323420 20.746290 74.878840 -50.101680 61.958500",
            "PAUSE 1000",
            "DISCONNECT",
        ]

        for line in commands:
            cmd, args = command_parser(line)
            if not cmd:
                continue
            self.commands_queue.put((cmd, args))


class RobotDriver(object):
    def __init__(self):
        self.robot = None
        self.vel = 20.0
        self.acc = 0.0
        self._manager = DriverManager(self)

    def DriverRun(self):
        self._manager.Run()
        return self

    def DriverKill(self):
        try:
            self.Disconnect()
        except Exception:
            pass
        self._manager.Stop()

    def __del__(self):
        try:
            self.DriverKill()
        except Exception:
            pass

    def CustomCommandList(self):
        commands = {}
        return "|".join([f'c {cmd}|{info}' for cmd, info in commands.items()])

    def Connect(self, ip_or_com: str, port, dof, *args, **kwargs):
        if frrpc is None:
            print_message("frrpc import failed")
            return False
        try:
            self.robot = frrpc.RPC(ip_or_com)
            # Probe the controller once so connect failures surface here.
            _ = self.Joints()
            return True
        except Exception as e:
            print_message("Connect failed: " + repr(e))
            self.robot = None
            return False

    def ConnectDefault(self, *args, **kwargs):
        if frrpc is None:
            print_message("frrpc import failed")
            return False
        try:
            self.robot = frrpc.RPC()
            _ = self.Joints()
            return True
        except Exception as e:
            print_message("ConnectDefault failed: " + repr(e))
            self.robot = None
            return False

    def Connected(self, force_check=False, *args, **kwargs):
        if self.robot is None:
            return False
        try:
            state = getattr(self.robot, "sock_cli_state_state", None)
            if state is not None:
                return bool(state)
            # Fallback: any successful lightweight read means we're connected.
            self.Joints()
            return True
        except Exception:
            return False

    def Disconnect(self, *args, **kwargs):
        if self.robot is None:
            return
        try:
            self.robot.CloseRPC()
        except Exception:
            pass
        finally:
            self.robot = None

    def Stop(self, *args, **kwargs):
        if self.robot is None:
            return
        try:
            self.robot.StopMotion()
        except Exception as e:
            print_message("StopMotion failed: " + repr(e))

    def _normalize_joints(self, data):
        # Common cases:
        # 1) [j1, j2, ...]
        # 2) (err, [j1, j2, ...])
        # 3) (err, j1, j2, j3, j4, j5, j6)
        # 4) nested lists/tuples
        if isinstance(data, (list, tuple)):
            if len(data) > 0 and all(isinstance(x, (int, float)) for x in data):
                # If first element looks like an error code and there are 6+ more values,
                # drop the first item. Otherwise treat the whole list as joints.
                if len(data) >= 7 and isinstance(data[0], (int, float)):
                    return list(data[1:])
                return list(data)

            for item in data:
                if isinstance(item, (list, tuple)) and all(isinstance(x, (int, float)) for x in item):
                    return list(item)

        if isinstance(data, (int, float)):
            raise TypeError("Joint query returned scalar instead of joint list: %r" % (data,))

        raise TypeError("Unexpected joint format: %r" % (data,))

    def Joints(self, *args, **kwargs):
        if self.robot is None:
            return []
        data = self.robot.GetActualJointPosDegree(0)
        try:
            return self._normalize_joints(data)
        except Exception as e:
            print_message("GetActualJointPosDegree raw: " + repr(data))
            raise

    def MoveJ(self, pose, joints, conf_RLF=None, *args, **kwargs):
        # frrpc does not accept keyword arguments here on your setup.
        # Use positional arguments instead.
        self.robot.MoveJ(joints, 0, 0, self.vel, self.acc)

    def MoveL(self, pose, joints, conf_RLF=None, *args, **kwargs):
        # Keep the same convention as the RoboDK draft driver, but use positional args.
        self.robot.MoveL(pose, 0, 0, self.vel, self.acc)

    def MoveC(self, pose1, joints1, pose2, joints2, conf_RLF_1=None, conf_RLF_2=None, *args, **kwargs):
        raise NotImplementedError

    def SearchL(self, joints, xyzwpr, *args, **kwargs):
        raise NotImplementedError

    def hasLinearSearch(self):
        return 0

    def setTool(self, pose, tool_id=None, tool_name=None, *args, **kwargs):
        pass

    def Pause(self, time_ms, *args, **kwargs):
        self.robot.WaitMs(time_ms)

    def setSpeed(self, speed_mms, *args, **kwargs):
        self.vel = min(float(speed_mms), 100.0)
        self.robot.SetSpeed(self.vel)

    def setAcceleration(self, accel_mmss, *args, **kwargs):
        self.acc = min(float(accel_mmss), 100.0)

    def setSpeedJoints(self, speed_degs, *args, **kwargs):
        raise NotImplementedError

    def setAccelerationJoints(self, accel_degss, *args, **kwargs):
        raise NotImplementedError

    def setZoneData(self, zone_mm, *args, **kwargs):
        raise NotImplementedError

    def setDO(self, io_var, io_value, *args, **kwargs):
        self.robot.SetDO(io_var, io_value)

    def setAO(self, io_var, io_value):
        self.robot.SetAO(io_var, io_value)

    def waitDI(self, io_var, io_value, timeout_ms=-1, *args, **kwargs):
        self.robot.WaitDI(io_var, io_value, timeout_ms, 2)

    def RunCode(self, code, is_function_call=False, *args, **kwargs):
        raise NotImplementedError

    def RunMessage(self, message, iscomment=False, *args, **kwargs):
        raise NotImplementedError


def RunMain():
    print_message(DRIVER_VERSION)
    UpdateStatus(ROBOTCOM_DISCONNECTED)
    driver = RobotDriver().DriverRun()
    driver.DriverKill()
    print_message("Process terminated cleanly")


if __name__ == "__main__":
    RunMain()