# run_skeleton_task.py (Completed)

import sys
import time
import math
import inspect
from typing import Dict, Tuple, Any, List

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

from env import setup_environment, shutdown_environment

# Import every predefined skill exactly as-is
from skill_code import pick, place, move, rotate, pull      # noqa: F401

from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# ------------------------------------------------------------
# Helper utilities (NOT new skills ─ just small wrappers)
# ------------------------------------------------------------
def is_position_reachable(target_pos: Tuple[float, float, float],
                          reach_radius: float = 1.25) -> bool:
    """
    Very conservative workspace check.  We assume the robot base is
    at (0, 0, 0).  Most RLBench arms have ≈1 m spherical workspace.
    We give a little margin.  This prevents infinite retries on
    positions that are clearly impossible.
    """
    distance = np.linalg.norm(target_pos[:3])
    return distance <= reach_radius


def safe_skill_call(skill_fn, *args, **kwargs):
    """
    Executes an existing skill while handling *any* exception so the
    main script can keep running and collect additional information.
    """
    try:
        print(f"[safe] Calling skill: {skill_fn.__name__}")
        return skill_fn(*args, **kwargs)
    except Exception as e:      # noqa: BLE001
        print(f"[safe]   Skill {skill_fn.__name__} raised: {e}")
        # Return sensible default triple so upstream logic never breaks
        return None, 0.0, False


def exploration_phase(env, task, positions: Dict[str, Tuple[float, float, float]]):
    """
    Very light-weight ‘exploration’ routine.  We simply *move* to each
    location we know of so that conditional effects like
    (identified ?obj) or (temperature-known ?obj) can potentially fire.
    """
    print("\n===== [Exploration] BEGIN =====")
    # We do not know the robot name in the PDDL, but the low-level
    # ‘move’ skill does not need it.  We therefore iterate over every
    # location and try a best-effort move.
    for obj_name, obj_pos in positions.items():
        print(f"[Exploration] Moving close to {obj_name} @ {obj_pos}")
        if not is_position_reachable(obj_pos):
            print(f"[Exploration]   {obj_name} is OUTSIDE workspace ‑ skipping.")
            continue
        # Adapt to any possible ‘move’ signature using Python introspection
        move_sig = inspect.signature(move)
        try_args: List[Any] = []
        try_kwargs: Dict[str, Any] = {}
        # Typical signature: (env, task, target_pos, …)
        if len(move_sig.parameters) >= 3:
            try_args = [env, task, obj_pos]
        else:                                 # Fallback for odd signatures
            try_kwargs = dict(env=env, task=task, target_pos=obj_pos)
        safe_skill_call(move, *try_args, **try_kwargs)
    print("===== [Exploration] END =====\n")


# ------------------------------------------------------------
# Main task routine
# ------------------------------------------------------------
def run_skeleton_task():
    """Generic skeleton for running any task in your simulation."""
    print("===== Starting Skeleton Task =====")
    env, task = setup_environment()

    try:
        # 1) Reset task and obtain the very first observation
        descriptions, obs = task.reset()
        init_video_writers(obs)

        # Wrap ‘step’ and ‘get_observation’ to enable automatic video capture
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # 2) Fetch all known object positions
        positions: Dict[str, Tuple[float, float, float]] = get_object_positions()
        print(f"[Init] Known objects: {list(positions.keys())}")

        # 3) Exploration phase to discover missing predicate information
        exploration_phase(env, task, positions)

        # 4) Example logical plan (very generic):
        #
        # Goal hypothesis: “All small objects must be placed in ‘bin’.”
        # So we will search for an object called 'bin' or
        # 'trash_bin' and treat every other object as candidate to dispose.
        #
        # NOTE: this is purely demonstrative; adapt to your real task.

        disposal_key_candidates = [k for k in positions if 'bin' in k.lower()
                                                     or 'trash' in k.lower()]
        if disposal_key_candidates:
            disposal_key = disposal_key_candidates[0]
            disposal_pos = positions[disposal_key]
            print(f"[Task] Disposal location assumed to be: {disposal_key} @ {disposal_pos}")
        else:
            disposal_key = None
            disposal_pos = None
            print("[Task] No obvious disposal container found – will just pick each object.")

        # main manipulation loop
        for obj_name, obj_pos in positions.items():
            if obj_name == disposal_key:     # never try to pick the bin itself
                continue

            print(f"\n[Task] Handling {obj_name}")
            if not is_position_reachable(obj_pos):
                print(f"[Task]   {obj_name} unreachable – skipping.")
                continue

            # --- PICK -------------------------------------------------------
            pick_sig = inspect.signature(pick)
            pick_args: List[Any] = []
            pick_kwargs: Dict[str, Any] = {}
            if len(pick_sig.parameters) >= 3:
                pick_args = [env, task, obj_pos]
            else:
                pick_kwargs = dict(env=env, task=task, target_pos=obj_pos)

            obs, reward, done = safe_skill_call(pick, *pick_args, **pick_kwargs)

            if done:
                print("[Task] Episode finished by environment during pick.")
                break

            # Optional weight/durability identified here via conditional
            # effects (from exploration domain).

            # --- PLACE ------------------------------------------------------
            if disposal_pos is not None:
                place_sig = inspect.signature(place)
                place_args: List[Any] = []
                place_kwargs: Dict[str, Any] = {}
                if len(place_sig.parameters) >= 3:
                    place_args = [env, task, disposal_pos]
                else:
                    place_kwargs = dict(env=env, task=task, target_pos=disposal_pos)
                obs, reward, done = safe_skill_call(place, *place_args, **place_kwargs)
                if done:
                    print("[Task] Episode finished by environment during place.")
                    break
            else:
                # If no disposal location, simply drop back (place at original)
                safe_skill_call(place, env, task, obj_pos)

        # 5) A very small example of using ‘rotate’ and ‘pull’
        # to open a drawer if such affordances exist:
        if hasattr(obs, 'gripper_pose'):   # only attempt if an observation is present
            print("\n[Drawer] Attempting rotation of gripper by 90° about Z.")
            ninety_deg_quat = np.array([0, 0, math.sin(math.pi/4), math.cos(math.pi/4)])
            safe_skill_call(rotate, env, task, ninety_deg_quat)

            # A blind ‘pull’ attempt; in a realistic scenario you would
            # need to *move* first to a handle pose
            safe_skill_call(pull, env, task)

    finally:
        # 6) Ensure a clean shutdown irrespective of success / failure
        shutdown_environment(env)
        print("===== End of Skeleton Task =====")


if __name__ == "__main__":
    run_skeleton_task()
