# Copyright 2015 - RoboDK Inc. - https://robodk.com/
# Licensed under the Apache License, Version 2.0 (the "License");
# ----------------------------------------------------
# Custom Coint / Yaskawa style post processor for RoboDK
# Features:
#  - P points numbered from P001
#  - ///NPOS replacement and ///PULSE -> Pxxx insertion
#  - MOVJ/MOVL/MOVC output with V/VJ, PL, ACC, DEC using RoboDK set* values
#  - ProgSendRobot uses SFTP (username/password passed by RoboDK) via paramiko
# ----------------------------------------------------

RDK_COMMANDS = {"ProgMoveNames": "1", "ProgMoveExtAxesPoses": "1"}

from robodk import *
from robodk import robomath
from robodk import robodialogs
from robodk.robodialogs import mbox
import os
import sys
import math

try:
    import paramiko
except Exception:
    try:
        import robolink
        robolink.import_install('paramiko')
        import paramiko
    except Exception:
        paramiko = None
        print("WARNING: paramiko not available. ProgSendRobot will fail unless paramiko is installed.")

# --------------------
# Helpers
def pose_2_str(pose):
    [x, y, z, r, p, w] = robomath.Pose_2_Adept(pose)
    return f"{x:.6f},{y:.6f},{z:.6f},{w:.6f},{p:.6f},{r:.6f}"

def joints_2_str(joints):
    out = []
    for i in range(len(joints)):
        out.append(f"{joints[i]:.6f}")
    return ",".join(out)

def str_2_type(s):
    """Converts a string containing a numerical or true/false (case insensitive) to a basic data type (bool/int/float).

    :param s: The string to process
    :type s: str
    :return: The basic data type (bool/int/float), or str if it fails
    :rtype: str or bool or int or float
    """
    if type(s) != str:
        return s
    if s.lower() in ['true', 'false']:
        return bool(s)
    try:
        f = float(s)
        if int(f) == f:
            return int(f)
        return f
    except:
        pass
    return s
    
def FilterName(namefilter, safechar='P', reserved_names=None, max_len=16, space_to_underscore=False, invalid_chars=r' .-[]/\;,><&*:%=+@!#^|?^'):
    if namefilter is None:
        namefilter = safechar
    if space_to_underscore and not '_' in invalid_chars:
        namefilter = namefilter.replace(' ', '_')
    for c in invalid_chars:
        namefilter = namefilter.replace(c, '')
    namefilter = ''.join((c for c in namefilter if 0 < ord(c) < 127))
    if len(namefilter) <= 0:
        namefilter = safechar
    if namefilter[0].isdigit():
        namefilter = safechar + namefilter
    if max_len > 0:
        namefilter = namefilter[:max_len]
    if reserved_names:
        reserved_names_lower = [n.lower() for n in reserved_names]
        cnt = 1
        while namefilter.lower() in reserved_names_lower:
            if len(namefilter) == max_len:
                namefilter = namefilter[:-1]
            if cnt > 1:
                if max_len > 0:
                    namefilter = namefilter[:min(max_len - len(str(cnt)), len(namefilter))]
                namefilter = namefilter + "%i" % cnt
            cnt += 1
        if type(reserved_names) is list:
            reserved_names.append(namefilter)
        else:
            reserved_names.add(namefilter)
    return namefilter


# -------------------------------------------------------

import os
import sys

def UploadFTP(program, robot_ip, remote_path, ftp_user, ftp_pass, pause_sec=2):
    """Upload program file(s) using SFTP (UR method)."""
    try:
        import paramiko
    except:
        import robolink
        robolink.import_install('paramiko')
        import paramiko

    print(f"POPUP: Connecting via SFTP to {robot_ip}:22 ...")
    sys.stdout.flush()

    try:
        transport = paramiko.Transport((robot_ip, 22))
        transport.connect(username=ftp_user, password=ftp_pass)
        sftp = paramiko.SFTPClient.from_transport(transport)
    except Exception as e:
        print(f"POPUP: <font color='red'>SFTP connection failed: {e}</font>")
        sys.stdout.flush()
        return False

    try:
        dirs = remote_path.split('/')
        cumulative = ""
        for d in dirs:
            if d.strip() == "":
                continue
            cumulative += "/" + d
            try:
                sftp.chdir(cumulative)
            except IOError:
                try:
                    sftp.mkdir(cumulative)
                    sftp.chdir(cumulative)
                except:
                    pass
    except Exception as e:
        print(f"POPUP: Remote path creation failed: {e}")
        sys.stdout.flush()

    if isinstance(program, list):
        files_to_send = program
    else:
        files_to_send = [program]

    sent_count = 0
    for filepath in files_to_send:
        if os.path.isfile(filepath):
            filename = os.path.basename(filepath)
            print(f"POPUP: Uploading {filename} ...")
            sys.stdout.flush()
            try:
                sftp.put(filepath, filename)
                sent_count += 1
            except Exception as e:
                print(f"POPUP: Upload failed for {filename}: {e}")
        else:
            print(f"POPUP: Skipping non-file: {filepath}")
            sys.stdout.flush()

    sftp.close()
    transport.close()

    print(f"POPUP: <font color='blue'>SFTP upload completed ({sent_count} files)</font>")
    sys.stdout.flush()
    return True


# --------------------
class RobotPost(object):
    PROG_EXT = 'JBR'
    DEFAULT_LINEAR_SPEED = 500.0
    DEFAULT_LINEAR_ACCEL = 100.0
    DEFAULT_JOINT_SPEED = 100.0
    DEFAULT_JOINT_ACCEL = 100.0
    DEFAULT_BLENDING = 0.0
    def __init__(self, robotpost='', robotname='', robot_axes=6, port=22, *args, **kwargs):
        self.ROBOT_POST = robotpost
        self.ROBOT_NAME = robotname
        self.ROBOT_AXES = robot_axes
        self.AXES_TYPE = kwargs.get('axes_type', ['R']*robot_axes)
        self.PROG = []            # current program lines (list)
        self.PROG_LIST = []       # list of programs
        self.PROG_NAMES = []
        self.PROG_FILES = []
        self.LOG = ''
        self.ONETAB = '  '
        self.TAB = ''

        # runtime state
        self.POINT_COUNTER = 1
        self.POINTS = []
        self._frame_name = None
        self._tool_name = None
        self._target_name = None

        # Speed/accel/blend state (defaults)
        self._speed = self.DEFAULT_LINEAR_SPEED
        self._acceleration = self.DEFAULT_LINEAR_ACCEL
        self._speed_joints = self.DEFAULT_JOINT_SPEED
        self._acceleration_joints = self.DEFAULT_JOINT_ACCEL
        self._blend = int(round(self.DEFAULT_BLENDING))

    # --------------------
    # Low-level helpers
    def addline(self, newline):
        self.PROG.append(self.TAB + newline)

    def addlog(self, newline):
        self.LOG = self.LOG + newline + '\n'

    # --------------------
    # Program lifecycle
    def ProgStart(self, progname, *args, **kwargs):
        progname = FilterName(progname)
        self.PROG_NAME = progname
        self.PROG_NAMES.append(progname)

        # reset points and program buffer
        self.POINT_COUNTER = 1
        self.POINTS = []
        self.PROG = []
        self.TAB = ''

        # header (kept as requested, TOOL stays 9)
        self.addline('//DIR')
        self.addline('//JOB')
        self.addline(f'//NAME {progname}')
        self.addline('//POS')
        self.addline('///NPOS {TOTAL_POINTS},0,0,0,0,0')
        self.addline('///TOOL 9')
        self.addline('///POSTYPE PULSE')
        self.addline('///PULSE')
        # instruction header starts; P lines will be inserted after ProgFinish processing
        self.addline('//INSTRUCTION')

    def ProgFinish(self, progname, *args, **kwargs):
        # replace NPOS and insert P points under ///PULSE, add NOP under //INSTRUCTION
        self.TAB = ''
        total_points = len(self.POINTS)

        if isinstance(self.PROG, list):
            new_lines = []
            for line in self.PROG:
                if "{TOTAL_POINTS}" in line:
                    line = line.replace("{TOTAL_POINTS}", str(total_points))
                new_lines.append(line)
                if line.strip().startswith("///PULSE"):
                    for p_line in self.POINTS:
                        new_lines.append(p_line)
                if line.strip().startswith("//INSTRUCTION"):
                    new_lines.append("NOP")
            self.PROG = new_lines

        # Program end
        self.addline("END")

        # store and reset program buffer
        self.PROG_LIST.append(self.PROG)
        self.PROG = []

    def ProgSave(self, folder, progname, ask_user=False, show_result=False, *args, **kwargs):
        # Single-file save for all programs
        if len(self.PROG_NAMES) == 0:
            return

        if len(self.PROG_NAMES) == 1:
            prog_save = self.PROG_NAME + '.' + self.PROG_EXT
            filesave = folder + '/' + prog_save
            if ask_user or not os.path.isdir(folder):
                filesave = robodialogs.getSaveFileName(folder, prog_save)
                if not filesave:
                    return
                folder = filesave.replace('\\', '/').rsplit('/', 1)[0]
        else:
            if ask_user or not os.path.isdir(folder):
                folder = robodialogs.getSaveFolder(folder)
                if not folder or not os.path.isdir(folder):
                    return

        filesave = folder + '/' + self.PROG_NAME + '.' + self.PROG_EXT
        with open(filesave, "w", encoding="utf-8") as fid:
            for prog in self.PROG_LIST:
                for line in prog:
                    fid.write(line)
                    fid.write("\n")

        print("SAVED:", filesave)
        self.PROG_FILES.append(filesave)

        # optionally show
        if show_result:
            try:
                import subprocess
                subprocess.Popen([show_result, filesave])
            except Exception:
                try:
                    os.startfile(filesave)
                except Exception:
                    pass
    def ProgSendRobot(self, robot_ip, remote_path, ftp_user, ftp_pass):
        """Send files via SFTP using ftp_user / ftp_pass (provided by RoboDK)"""
        if len(self.PROG_FILES) == 0:
            print("POPUP: No files to transfer")
            return
        # Use port 22
        UploadFTP(self.PROG_FILES, robot_ip, remote_path, ftp_user, ftp_pass, port=22)

    # --------------------
    # Movement primitives
    def MoveJ(self, pose, joints, conf_RLF=None, *args, **kwargs):
        # ensure counters
        if not hasattr(self, "POINT_COUNTER"):
            self.POINT_COUNTER = 1
            self.POINTS = []

        point_name = f"P{self.POINT_COUNTER:03d}"
        self.POINT_COUNTER += 1

        joint_str = joints_2_str(joints)
        # Format in POS table: matches example "0,0,0,0,0,0,0,<j1>,<j2>..."
        self.POINTS.append(f"{point_name} = 0,0,0,0,0,0,0,{joint_str}")

        # dynamic params
        VJ = getattr(self, "_speed_joints", self.DEFAULT_JOINT_SPEED)
        ACCJ = getattr(self, "_acceleration_joints", self.DEFAULT_JOINT_ACCEL)
        DECJ = getattr(self, "_acceleration_joints", self.DEFAULT_JOINT_ACCEL)
        PL = getattr(self, "_blend", int(round(self.DEFAULT_BLENDING)))

        # write movement line (use % for VJ as in your examples)
        self.addline(f"MOVJ {point_name} VJ = {VJ:.1f} % PL = {PL} ACC = {ACCJ:.0f} DEC = {DECJ:.0f}")

    def MoveL(self, pose, joints, conf_RLF=None, *args, **kwargs):
        if not hasattr(self, "POINT_COUNTER"):
            self.POINT_COUNTER = 1
            self.POINTS = []

        point_name = f"P{self.POINT_COUNTER:03d}"
        self.POINT_COUNTER += 1

        joint_str = joints_2_str(joints)
        self.POINTS.append(f"{point_name} = 0,0,0,0,0,0,0,{joint_str}")

        V = getattr(self, "_speed", self.DEFAULT_LINEAR_SPEED)
        ACC = getattr(self, "_acceleration", self.DEFAULT_LINEAR_ACCEL)
        DEC = getattr(self, "_acceleration", self.DEFAULT_LINEAR_ACCEL)
        PL = getattr(self, "_blend", int(round(self.DEFAULT_BLENDING)))

        self.addline(f"MOVL {point_name} V = {V:.1f} mm/s PL = {PL} ACC = {ACC:.0f} DEC = {DEC:.0f}")

    def MoveC(self, pose1, joints1, pose2, joints2, conf_RLF_1=None, conf_RLF_2=None, *args, **kwargs):
        # keep last target pose/joints
        self._target_pose = pose2
        self._target_joints = joints2

        # Validate cartesian poses
        if pose1 is None or pose2 is None:
            msg = "Circular movement using joint targets is not supported."
            self.addlog(msg)
            self.RunMessage(msg, True)
            return

        # produce two P points (waypoint and end)
        for joints in [joints1, joints2]:
            point_name = f"P{self.POINT_COUNTER:03d}"
            self.POINT_COUNTER += 1
            joint_str = joints_2_str(joints)
            self.POINTS.append(f"{point_name} = 0,0,0,0,0,0,0,{joint_str}")

        # dynamic params
        V = getattr(self, "_speed", self.DEFAULT_LINEAR_SPEED)
        ACC = getattr(self, "_acceleration", self.DEFAULT_LINEAR_ACCEL)
        DEC = getattr(self, "_acceleration", self.DEFAULT_LINEAR_ACCEL)
        PL = getattr(self, "_blend", int(round(self.DEFAULT_BLENDING)))

        # Write a MOVC line that references last two created points
        p_last2 = f"P{self.POINT_COUNTER-2:03d}"
        p_last1 = f"P{self.POINT_COUNTER-1:03d}"
        self.addline(f"MOVC {p_last2} {p_last1} V = {V:.1f} mm/s PL = {PL} ACC = {ACC:.0f} DEC = {DEC:.0f}")

    # --------------------
    # Frames/tools/IO/controls
    def setFrame(self, pose, frame_id=None, frame_name=None, *args, **kwargs):
        self._frame_pose = pose
        self._frame_id = frame_id
        self._frame_name = FilterName(frame_name) if frame_name else None
        if self._frame_name:
            self.addline('// Frame: %s' % self._frame_name)
        self.addline('BASE_FRAME ' + pose_2_str(pose))

    def setTool(self, pose, tool_id=None, tool_name=None, *args, **kwargs):
        self._tool_pose = pose
        self._tool_id = tool_id
        self._tool_name = FilterName(tool_name) if tool_name else None
        if self._tool_name:
            self.addline('// Tool: %s' % self._tool_name)
        self.addline('TOOL_FRAME ' + pose_2_str(pose))

    def Pause(self, time_ms, *args, **kwargs):
        if time_ms < 0:
            self.addline('PAUSE')
        else:
            self.addline('WAIT %.3f' % (time_ms * 0.001))

    # --------------------
    # Speed / Acceleration / Blending (RoboDK calls these)
    def setSpeed(self, speed_mms, *args, **kwargs):
        self._speed = speed_mms

    def setAcceleration(self, accel_mmss, *args, **kwargs):
        self._acceleration = accel_mmss

    def setSpeedJoints(self, speed_degs, *args, **kwargs):
        # keep same numeric value (user expects deg/s — examples used percent; keep value as provided)
        # if desired, convert to percent mapping outside
        self._speed_joints = speed_degs

    def setAccelerationJoints(self, accel_degss, *args, **kwargs):
        self._acceleration_joints = accel_degss

    def setZoneData(self, zone_mm, *args, **kwargs):
        # store PL as integer rounding
        try:
            self._blend = int(round(zone_mm))
        except Exception:
            self._blend = 0

    # --------------------
    # I/O and code helpers
    def setDO(self, io_var, io_value, *args, **kwargs):
        if type(io_var) != str:
            io_var = f'DO[{io_var}]'
        io_value = io_value if type(io_value) is str else str_2_type(io_value)
        if type(io_value) != str:
            io_value = 'TRUE' if io_value else 'FALSE'
        self.addline(f'{io_var} = {io_value}')

    def setAO(self, io_var, io_value, *args, **kwargs):
        if type(io_var) != str:
            io_var = f'AO[{io_var}]'
        io_value = io_value if type(io_value) is str else ('%.3f' % float(io_value))
        self.addline(f'{io_var} = {io_value}')

    def waitDI(self, io_var, io_value, timeout_ms=-1, *args, **kwargs):
        if type(io_var) != str:
            io_var = f'DI[{io_var}]'
        io_value = io_value if type(io_value) is str else ('TRUE' if io_value else 'FALSE')
        if timeout_ms < 0:
            self.addline(f'WAIT FOR {io_var}=={io_value}')
        else:
            self.addline(f'WAIT FOR {io_var}=={io_value} TIMEOUT={timeout_ms}')

    def RunCode(self, code, is_function_call=False, *args, **kwargs):
        if is_function_call:
            code = FilterName(code, max_len=-1)
            if not code.endswith('()'):
                code = code + '()'
            self.addline(code)
        else:
            self.addline(code)

    def RunMessage(self, message, iscomment=False, *args, **kwargs):
        if iscomment:
            self.addline('// ' + message)
        else:
            self.addline(f'TP "{message}"')

# --------------------
# Test routine (can be used to validate output)
def test_post():
    from robodk.robomath import PosePP as p
    r = RobotPost("Coint_Post", "Coint eRob10", 6)
    r.ProgStart("Prog8")
    r.RunMessage("## Program generated by custom post", True)
    # set some speeds and rounding
    r.setSpeed(20.0)
    r.setAcceleration(10.0)
    r.setSpeedJoints(50.0)
    r.setAccelerationJoints(10.0)
    r.setZoneData(3.0)
    # create a few moves (simulate RoboDK calls)
    r.MoveJ(None, [94.38, -33.36, -93.24, 33.89, 135.03, -92.71])
    r.MoveC(None, [96.066, -35.6977, -89.551, 34.2264, 135.085, -90.3252], None, [98.1477, -33.5738, -93.6575, 38.2972, 135.084, -87.3766])
    r.MoveJ(None, [98.1477, -33.5738, -93.6575, 38.2972, 135.084, -87.3766])
    r.MoveC(None, [96.4803, -31.1929, -97.3623, 37.9484, 135.091, -89.7384], None, [94.38, -33.3601, -93.2396, 33.8887, 135.03, -92.7107])
    r.ProgFinish("Prog8")
    # Save to current folder for inspection
    out_folder = '.'
    r.ProgSave(out_folder, "Prog8", ask_user=False, show_result=False)
    print("Generated file(s):", r.PROG_FILES)
    for prog in r.PROG_LIST:
        for line in prog:
            print(line)

if __name__ == "__main__":
    test_post()