from pathlib import Path
import sys
import time


ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

CLOSE_SCISSORS_ON_PICK = True

# Fill these before running: [x, y, z, roll, pitch, yaw]
PRE_PICK_POSE = []
PICK_POSE = []

control = None
control_core = None


def _load_project():
    global control, control_core
    if control is None or control_core is None:
        import control as control_module
        import control_core as control_core_module

        control = control_module
        control_core = control_core_module
    return control, control_core


def _checked_pose(pose, name):
    if (
        not isinstance(pose, (list, tuple))
        or len(pose) != 6
        or any(not isinstance(value, (int, float)) for value in pose)
    ):
        raise ValueError(f"{name} must be [x, y, z, roll, pitch, yaw]")
    return [float(value) for value in pose]


def _pulse_scissors(robot_name, io_index, off_index):
    project, _ = _load_project()
    project.control_tool_io(robot_name, off_index, False)
    project.control_tool_io(robot_name, io_index, True)
    time.sleep(float(project.GRIPPER_ACTION_DELAY))
    project.control_tool_io(robot_name, io_index, False)


def open_scissors(robot_name):
    _pulse_scissors(robot_name, control.GRIPPER_OPEN_IO, control.GRIPPER_CLOSE_IO)


def close_scissors(robot_name):
    _pulse_scissors(robot_name, control.GRIPPER_CLOSE_IO, control.GRIPPER_OPEN_IO)


def _move_pose_or_fail(arm, pose, reference_joints, description):
    joints = arm.move_pose(pose, reference_joints, description=description)
    if joints is None:
        raise RuntimeError(f"{description} failed")
    return joints


def run_teach_pick_once():
    arm = None
    core = None

    try:
        pre_pick_pose = _checked_pose(PRE_PICK_POSE, "PRE_PICK_POSE")
        pick_pose = _checked_pose(PICK_POSE, "PICK_POSE")
        project, core = _load_project()
        project._sync_config()
        core.reset_runtime_state()
        core.placement_manager.reset()

        arm = core.connect_robot_arm()
        open_scissors(arm.robot_name)

        if not arm.return_home():
            raise RuntimeError("return Home failed")

        reference_joints = arm.get_joint_positions() or project.HOME_JOINTS
        pre_joints = _move_pose_or_fail(arm, pre_pick_pose, reference_joints, "move to pre-pick pose")
        pick_joints = _move_pose_or_fail(arm, pick_pose, pre_joints, "move to pick pose")

        delay = float(project.PICK_POINT_ARRIVAL_DELAY)
        if delay > 0:
            print(f"Arrived at pick pose, waiting {delay:.1f}s")
            time.sleep(delay)

        if CLOSE_SCISSORS_ON_PICK:
            close_scissors(arm.robot_name)
        else:
            print("CLOSE_SCISSORS_ON_PICK=False, skip closing scissors")

        _move_pose_or_fail(arm, pre_pick_pose, pick_joints, "return to pre-pick pose")

        if not arm.return_home():
            raise RuntimeError("return Home after pick failed")

        place_pose, place_index = project.get_next_placement()
        reference_joints = arm.get_joint_positions() or project.HOME_JOINTS
        _move_pose_or_fail(arm, place_pose, reference_joints, f"move to place pose {place_index}")
        open_scissors(arm.robot_name)

        if not arm.return_home():
            raise RuntimeError("final return Home failed")

        print("Teach pick flow finished")
        return True
    except Exception as exc:
        print(f"Teach pick flow failed: {exc}")
        if arm is not None:
            try:
                arm.return_home()
            except Exception as home_exc:
                print(f"Return Home during cleanup failed: {home_exc}")
        return False
    finally:
        if core is not None:
            core.release_robot_arm()


def _self_test():
    assert _checked_pose([1, 2, 3, 4, 5, 6], "POSE") == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
    try:
        _checked_pose([], "POSE")
    except ValueError:
        pass
    else:
        raise AssertionError("empty pose must fail")
    assert isinstance(CLOSE_SCISSORS_ON_PICK, bool)


if __name__ == "__main__":
    if "--self-test" in sys.argv:
        _self_test()
    else:
        raise SystemExit(0 if run_teach_pick_once() else 1)