# run_skeleton_task.py  (Completed Version – no new skills are defined)

import inspect
import time
import numpy as np

from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment
from skill_code import *                 # <- keep all existing skills
from video import (init_video_writers,
                   recording_step,
                   recording_get_observation)

from object_positions import get_object_positions


# ==============  Helper / Utility functions  ============== #
def safe_skill_call(skill_fn, *args, **kwargs):
    """
    Try to execute a skill.  If the signature requires more positional
    arguments than what was provided, we pad the call with `None` so that
    we never raise a TypeError due to missing parameters.  Any other
    exception coming from inside the skill is caught and printed, but the
    program keeps running (so that exploration can continue).
    """
    try:
        return skill_fn(*args, **kwargs)
    except TypeError as e:
        # Pad missing positional arguments with None
        sig            = inspect.signature(skill_fn)
        total_required = len(sig.parameters)
        padded_args    = list(args) + [None] * max(0, total_required - len(args))
        try:
            return skill_fn(*padded_args, **kwargs)
        except Exception as e_inner:
            print(f"[safe_skill_call]  '{skill_fn.__name__}' failed  ({e_inner})")
    except Exception as e:
        print(f"[safe_skill_call]  '{skill_fn.__name__}' threw an exception  ({e})")


def validate_action_shape(env):
    """
    A tiny safeguard suggested by the feedback: verify that `env.action_shape`
    looks like a 7-DoF pose (x,y,z,qx,qy,qz,qw) plus a gripper scalar.
    Typical RLBench tasks expose 8 elements (7 pose + 1 gripper) or 9
    (adding one for optional collision flag).  Anything wildly different
    is suspicious and we warn the user.
    """
    shape = getattr(env, "action_shape", None)
    if shape is None:
        print("[validate_action_shape]  env.action_shape missing!")
    elif shape not in (7, 8, 9) and not isinstance(shape, (tuple, list)):
        raise ValueError(f"[validate_action_shape]  Unexpected action shape: {shape}")
    else:
        print(f"[validate_action_shape]  action shape OK: {shape}")


def identify_domain_objects(positions_dict):
    """
    Very light-weight naming heuristics so that the generic skeleton can
    still do *something* useful without hard-coding exact scene names.
    """
    drawers   = [k for k in positions_dict if "drawer" in k.lower()]
    handles   = [k for k in positions_dict if "handle" in k.lower()]
    bins      = [k for k in positions_dict if "bin"    in k.lower()]
    rubbish   = [k for k in positions_dict
                 if k not in drawers + handles + bins]
    # If we could not find a bin fall back to any ‘basket’ keyword
    if not bins:
        bins = [k for k in positions_dict if "basket" in k.lower()]
    return drawers, handles, bins, rubbish


def drawer_is_locked(env, task, drawer_handle_name):
    """
    Tiny probing routine: try a gentle `pull` once.  We treat success as
    ‘unlocked’ and failure as ‘locked/missing predicate’.
    """
    print(f"[drawer_is_locked]  probing '{drawer_handle_name}'")
    outcome = safe_skill_call(pull, env, task, drawer_handle_name)
    # Very naïve heuristic: if the skill returns a tuple we assume it was
    # executed; if it returned None we treat it as failure.
    locked  = outcome is None
    msg     = "LOCKED" if locked else "UNLOCKED"
    print(f"[drawer_is_locked]  drawer '{drawer_handle_name}' appears {msg}")
    return locked


# ==============              MAIN              ============== #
def run_skeleton_task():
    """
    Generic task-runner that (i) sets up the environment, (ii) performs a
    light exploration pass to discover missing predicates such as
    ‘is-locked’, and (iii) – when possible – picks every loose object and
    drops it into the nearest bin.
    """
    print("==========  Starting Skeleton Task  ==========")

    env, task = setup_environment()
    try:
        # ----------  Environment reset & video helpers  ----------
        descriptions, obs = task.reset()
        init_video_writers(obs)

        original_step      = task.step
        task.step          = recording_step(original_step)
        original_get_obs   = task.get_observation
        task.get_observation = recording_get_observation(original_get_obs)

        # ----------  Feedback:  action shape validation  ----------
        validate_action_shape(env)

        # ----------  Retrieve object positions  ----------
        positions = get_object_positions()
        if not positions:
            print("[Task]  Warning: object_positions empty – nothing to do.")
            return

        drawers, handles, bins, rubbish = identify_domain_objects(positions)
        print(f"[Task]  Found {len(drawers)} drawers, "
              f"{len(handles)} handles, {len(bins)} bins, "
              f"{len(rubbish)} candidate loose objects.")

        # ==========================================================
        #  EXPLORATION  – discover the “is-locked” predicate
        # ==========================================================
        locked_drawers = {}
        for handle_name in handles:
            is_locked = drawer_is_locked(env, task, handle_name)
            locked_drawers[handle_name] = is_locked

        print("[Exploration]  Summary (handle -> locked?):", locked_drawers)

        # ==========================================================
        #  SIMPLE CLEAN-UP ROUTINE
        # ==========================================================
        if not bins:
            print("[Task]  No disposal bin found – skipping clean-up.")
        else:
            main_bin_name  = bins[0]
            main_bin_pos   = positions[main_bin_name]

            for obj_name in rubbish:
                print(f"[Task]  Picking '{obj_name}'")
                obj_pos = positions[obj_name]

                # ---  pick  ---
                safe_skill_call(pick,
                                env,
                                task,
                                target_pos        = obj_pos,
                                approach_distance = 0.15,
                                max_steps         = 120,
                                threshold         = 0.01,
                                approach_axis     = 'z',
                                timeout           = 10.0)

                # ---  move (generic)  ---
                safe_skill_call(move,
                                env,
                                task,
                                target_pos = main_bin_pos,
                                speed      = 0.25,
                                threshold  = 0.01,
                                timeout    = 10.0)

                # ---  place  ---
                safe_skill_call(place,
                                env,
                                task,
                                target_pos        = main_bin_pos,
                                retreat_distance  = 0.12,
                                max_steps         = 80,
                                threshold         = 0.01,
                                approach_axis     = 'z',
                                timeout           = 6.0)

        print("==========  Skeleton Task Completed  ==========")

    finally:
        shutdown_environment(env)


if __name__ == "__main__":
    run_skeleton_task()