# run_skeleton_task.py (Completed Executable Code)

import time
import traceback
from typing import Dict, Tuple

import numpy as np
from pyrep.objects.shape import Shape

from env import setup_environment, shutdown_environment
from skill_code import pick, place, move, rotate, pull      # only the predefined skills
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


def _safe_call(skill_fn, *args, **kwargs):
    """
    Helper that wraps every skill invocation in a try / except so that the entire
    exploratory script never crashes because a single skill fails.  All exceptions
    are printed, but execution continues.
    """
    try:
        return skill_fn(*args, **kwargs)
    except Exception as e:            # noqa: broad-except   – we really do not want to crash
        print(f"[Warning] {skill_fn.__name__} raised an exception:")
        traceback.print_exc()
        # We always return a 3-tuple so the caller can keep going even on failure.
        return None, 0.0, False


def _find_drawers(positions: Dict[str, Tuple[float, float, float]]):
    """
    Naïve drawer detector: everything whose name contains the substring 'drawer'
    is treated as a drawer.  This is obviously not perfect, but it is sufficient
    for automatic exploration whose sole purpose is to discover whether the
    drawer is locked (→ the missing ­predicate ‘lock-known’).
    """
    return [name for name in positions if "drawer" in name.lower()]


def run_skeleton_task():
    """
    Generic task runner that (1) performs a short exploratory phase meant to
    discover the missing predicate ‘lock-known’ and (2) shows how to structure
    typical RLBench calls with the predefined skills only.
    """
    print("===== Starting Skeleton Task =====")

    # === 1)  Environment Setup =================================================
    env, task = setup_environment()
    try:
        # Reset the task into a fresh episode
        descriptions, obs = task.reset()

        # Optional: start video recording
        init_video_writers(obs)

        # Monkey-patch ‘step’ & ‘get_observation’ so that every environment step
        # is automatically recorded to disk.
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # === 2)  Obtain world information ======================================
        object_positions = get_object_positions()        # Dict[str, (x, y, z)]
        print(f"[Info] Detected objects: {list(object_positions.keys())}")

        # Detect all candidate drawers                                             
        drawer_names = _find_drawers(object_positions)
        print(f"[Info] Candidate drawers: {drawer_names}")

        # === 3)  EXPLORATION PHASE =============================================
        #
        # Goal: determine whether the *lock-known* predicate is required.  We
        # attempt a naïve sequence of actions on every detected drawer:
        #
        #   a) move the gripper close to the drawer front
        #   b) rotate the gripper so that it can grasp the handle
        #   c) pick (grasp) the handle
        #   d) pull the drawer
        #
        # If the drawer does not open we infer that it is probably locked, i.e.,
        # we have just gathered the previously missing information
        #      (lock-known ?d)     ; new knowledge about the drawer.
        # -----------------------------------------------------------------------

        for drawer in drawer_names:
            target_pos = object_positions[drawer]
            print(f"\n[Exploration] Working on drawer '{drawer}' @ {target_pos}")

            # ---- 3a) Move close to the drawer ---------------------------------
            _safe_call(
                move,
                env,
                task,
                target_pos=target_pos,
                approach_distance=0.20,
                max_steps=150,
                threshold=0.01,
                timeout=5.0
            )

            # ---- 3b) Rotate the wrist so the fingers are horizontal -----------
            #        ‘ninety_deg’ in PDDL roughly equals (180°, 0°, 90°) in XYZ-
            #        Euler which is quaternion (-0.707, 0.000, 0.000, 0.707).
            ninety_deg_quat = np.array([-0.70710678, 0.00000000, 0.00000000, 0.70710678])
            _safe_call(rotate, env, task, target_quat=ninety_deg_quat)

            # ---- 3c) Pick / grasp the drawer handle ---------------------------
            _safe_call(
                pick,
                env,
                task,
                target_pos=target_pos,
                approach_distance=0.07,
                max_steps=120,
                threshold=0.005,
                approach_axis='z',
                timeout=4.0
            )

            # ---- 3d) Pull the drawer ------------------------------------------
            obs, reward, done = _safe_call(pull, env, task)
            # Pull returns (obs, reward, done) by convention; we do not need
            # ‘reward’ for exploration – we only care if the drawer opened.
            if done:
                print(f"[Exploration] Drawer '{drawer}' opened ⇒ it was UNlocked.")
            else:
                print(f"[Exploration] Drawer '{drawer}' did NOT open ⇒ it is probably LOCKED.")
                print("               ⇒ Missing predicate discovered: lock-known")

            # Whenever ‘done’ becomes True the RLBench task terminates and we
            # must break to avoid operating on an un-reset environment.
            if done:
                break

        # === 4)  (Optional)  Additional task-specific logic goes here ==========
        #
        # In a real project you would now use the knowledge obtained above
        # (i.e., which drawers are locked / unlocked) to formulate the complete
        # task plan.  Because the purpose of this template is to showcase the
        # exploration strategy only, we immediately terminate.

    finally:
        # === 5)  Clean-up ======================================================
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()