# run_skeleton_task.py  (Completed & ready-to-run)

import time
import numpy as np

from pyrep.objects.shape import Shape                 # <- kept exactly as in skeleton
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment
from skill_code import *                              # <- we only CALL the existing skills
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# --------------------------------------------------------
# Utility / helper functions – these do NOT redefine skills
# --------------------------------------------------------
def safe_call(fn, *args, **kwargs):
    """
    Generic wrapper that calls a skill and captures exceptions so
    that the overall task execution can continue gracefully.
    """
    try:
        print(f"[safe_call] Calling skill: {fn.__name__}")
        return fn(*args, **kwargs)
    except Exception as e:
        print(f"[safe_call]  -> Skill {fn.__name__} raised: {e}")
        # Signal “failure” with done=True so higher level code can react
        obs = kwargs.get("task").get_observation() if "task" in kwargs else None
        return obs, 0.0, True


def find_position_key(positions, substrings):
    """
    Returns the first key in `positions` whose name contains *any*
    of the provided substrings (case-insensitive).  If none are found
    returns None.
    """
    for key in positions:
        for sub in substrings:
            if sub.lower() in key.lower():
                return key
    return None


def run_skeleton_task():
    """Generic skeleton for running any task in your simulation."""
    print("\n===== Starting Skeleton Task =====")

    # =====================
    # 1) Environment setup
    # =====================
    env, task = setup_environment()
    try:
        # -- Reset --
        descriptions, obs = task.reset()

        # -- Video helpers (optional) --
        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)

        # ==============================
        # 2) Obtain useful environment info
        # ==============================
        positions = get_object_positions()     # <- user provided helper
        print("[Info] Available position keys:", list(positions.keys()))

        # Heuristically extract interesting objects
        drawer_key   = find_position_key(positions, ["drawer", "cabinet"])
        handle_key   = find_position_key(positions, ["handle", "knob"])
        plate_key    = find_position_key(positions, ["plate", "dish"])
        bin_key      = find_position_key(positions, ["bin", "disposal", "trash"])
        side_key     = find_position_key(positions, ["side"])
        anchor_key   = find_position_key(positions, ["anchor", "front"])

        # Fail early if mandatory objects are missing
        mandatory = [drawer_key, handle_key, plate_key, bin_key]
        if None in mandatory:
            print("[Warning] Could not identify every mandatory object in `get_object_positions()` – "
                  "exploration will continue but might not accomplish the full task.")

        # =========================================================
        # 3) Exploration PHASE – determine the ‘missing’ predicate
        # =========================================================
        #
        # From the feedback we expect that the domain might require us
        # to *learn* whether the drawer is locked (predicate lock-known ?d).
        # The only way to obtain that information in the PDDL is by
        # executing the “pull” skill once we are holding the drawer
        # handle.  Therefore the exploration phase tries exactly that:
        #  1) move next to the handle
        #  2) orient the gripper correctly
        #  3) pick the handle (anchor position)
        #  4) pull a little bit – regardless of success we now “know”
        #     the lock status according to the abstract domain
        #
        # Any success / failure is printed to the console.  The rest of
        # the actual task can then rely on the fact that the predicate
        # (lock-known ?drawer) exists in the belief state.
        # ---------------------------------------------------------

        # --- Move next to the drawer’s *side* position (if provided) ---
        if side_key is not None:
            side_pos = positions[side_key]
            obs, reward, done = safe_call(move,
                                          env, task,
                                          target_pos=side_pos,
                                          approach_distance=0.10,
                                          max_steps=120,
                                          threshold=0.01,
                                          approach_axis='z',
                                          timeout=10.0)
            if done:
                print("[Exploration] Task ended unexpectedly during side-approach.")
                return
        else:
            print("[Exploration] No explicit side position – skipping side-approach.")

        # --- Rotate the gripper 90° so that it is parallel to the drawer front ---
        ninety_deg_quat = np.array([0.0, 0.0, np.sin(np.pi/4), np.cos(np.pi/4)])   # simple Z-rotation
        obs, reward, done = safe_call(rotate, env, task,
                                      target_quat=ninety_deg_quat,
                                      max_steps=100,
                                      threshold=0.05,
                                      timeout=10.0)
        if done:
            print("[Exploration] Task ended during rotation.")
            return

        # --- Move to anchor position (center of handle) if known ---
        if anchor_key is not None:
            anchor_pos = positions[anchor_key]
            obs, reward, done = safe_call(move,
                                          env, task,
                                          target_pos=anchor_pos,
                                          approach_distance=0.07,
                                          max_steps=120,
                                          threshold=0.005,
                                          approach_axis='z',
                                          timeout=10.0)
            if done:
                print("[Exploration] Task ended while moving to anchor position.")
                return
        else:
            print("[Exploration] No explicit anchor position.  Using handle position instead.")
            anchor_pos = positions.get(handle_key, None)

        # --- Pick (grasp) the drawer handle ---
        if anchor_pos is not None:
            obs, reward, done = safe_call(pick,
                                          env, task,
                                          target_pos=anchor_pos,
                                          approach_distance=0.03,
                                          max_steps=80,
                                          threshold=0.002,
                                          approach_axis='z',
                                          timeout=8.0)
            if done:
                print("[Exploration] Task ended while picking the handle.")
                return
        else:
            print("[Exploration] No anchor / handle position available – cannot pick.")
            # Continue anyway; there is nothing else we can do.

        # --- SHORT PULL ACTION to discover the lock status ---
        print("[Exploration] Attempting a small pull to learn lock state...")
        obs, reward, done = safe_call(pull, env, task)
        # Even if the drawer is locked the domain considers the lock status ‘known’
        # after the pull action.

        # =======================================================
        # 4) MAIN TASK (example) – open drawer and fetch a plate
        # =======================================================
        if done:
            print("[Main] Early termination – cannot continue to main task.")
            return

        # -- If the drawer is still closed (or locked) print info and stop --
        if hasattr(task.get_observation(), "drawer_opened") and not task.get_observation().drawer_opened:
            print("[Main] Drawer appears to be locked; stopping as goal is unreachable.")
            return

        # At this point the drawer should be open (is-open ?d) so we can
        # finally grab the plate (or whatever object is inside).

        # --- Move end-effector into the drawer interior (reuse anchor_pos + small offset) ---
        if plate_key is not None:
            plate_pos = positions[plate_key]
            obs, reward, done = safe_call(move,
                                          env, task,
                                          target_pos=plate_pos,
                                          approach_distance=0.06,
                                          max_steps=120,
                                          threshold=0.005,
                                          approach_axis='z',
                                          timeout=10.0)
            if done:
                print("[Main] Task ended during move to plate.")
                return

            # --- Pick the plate ---
            obs, reward, done = safe_call(pick,
                                          env, task,
                                          target_pos=plate_pos,
                                          approach_distance=0.02,
                                          max_steps=80,
                                          threshold=0.002,
                                          approach_axis='z',
                                          timeout=8.0)
            if done:
                print("[Main] Task ended while picking the plate.")
                return
        else:
            print("[Main] No plate detected – skipping pick.")

        # --- Move to disposal / bin and place object ---
        if bin_key is not None and not done:
            bin_pos = positions[bin_key]

            obs, reward, done = safe_call(move,
                                          env, task,
                                          target_pos=bin_pos,
                                          approach_distance=0.10,
                                          max_steps=120,
                                          threshold=0.01,
                                          approach_axis='z',
                                          timeout=10.0)
            if done:
                print("[Main] Task ended during move to bin.")
                return

            obs, reward, done = safe_call(place,
                                          env, task,
                                          target_pos=bin_pos,
                                          approach_distance=0.08,
                                          max_steps=80,
                                          threshold=0.01,
                                          approach_axis='z',
                                          timeout=8.0)
            if done:
                print("[Main] Task ended during placing.")
                return
        else:
            print("[Main] No disposal location – skipping place.")

        # ==================================================
        # 5) Wrap-up
        # ==================================================
        print("[Result] Finished main task without early termination.")

    finally:
        # Always ensure a clean shutdown
        shutdown_environment(env)

    print("===== End of Skeleton Task =====")


if __name__ == "__main__":
    run_skeleton_task()