# run_skeleton_task.py (Completed)

import numpy as np
import traceback
from env import setup_environment, shutdown_environment

# Import every predefined skill exactly once
from skill_code import move, pick, place, rotate, pull   # noqa: F401
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions


# ------------------------------------------------------------------
# Helper Utilities (pure-python, no new skills)
# ------------------------------------------------------------------
def _safe_call(fn, *args, **kwargs):
    """
    Helper that safely calls a skill and prints the traceback if the
    call fails.  Returns (obs, reward, done) on success, else (None, 0, True).
    """
    try:
        return fn(*args, **kwargs)
    except Exception as exc:
        print("====================================================")
        print(f"[ERROR] Exception while executing skill {fn.__name__}: {exc}")
        traceback.print_exc()
        print("====================================================")
        # Signal to the main loop that we should abort the episode
        return None, 0.0, True


def run_skeleton_task():
    """
    Generic entry-point for running any task in our simulated environment.
    The flow:
        1.  Set up environment / video recording.
        2.  Validate object list (feedback requirement).
        3.  Perform an exploration phase – attempt to visit every object
            in the scene and learn about it with predefined skills.
        4.  Shutdown the environment gracefully.
    """
    print("===== Starting Skeleton Task =====")

    # ------------------------------------------------------------------
    # 1)  Environment Setup
    # ------------------------------------------------------------------
    env, task = setup_environment()
    try:
        # Reset environment and obtain initial observation / description
        descriptions, obs = task.reset()

        # Optional: initialise video capture
        init_video_writers(obs)

        # Wrap task.step & task.get_observation so that every interaction
        # is automatically recorded to disk
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # ------------------------------------------------------------------
        # 2)  Retrieve & Validate Object List
        # ------------------------------------------------------------------
        positions = get_object_positions() or {}
        print(f"[Info] Detected objects from object_positions: {list(positions.keys())}")

        # Feedback specified that "rubbish" and "waypoint1" caused problems.
        missing_critical_objs = [name for name in ["rubbish", "waypoint1"] if name not in positions]
        if missing_critical_objs:
            print(f"[Warning] The following expected objects are NOT present in the scene: {missing_critical_objs}")

        # Identify a plausible disposal / drop site if it exists
        bin_candidates = [k for k in positions if "bin" in k.lower() or "trash" in k.lower()]
        disposal_site_name = bin_candidates[0] if bin_candidates else None
        if disposal_site_name:
            print(f"[Info] Using '{disposal_site_name}' as disposal site.")
            disposal_site_pos = np.array(positions[disposal_site_name])
        else:
            disposal_site_pos = None

        # ------------------------------------------------------------------
        # 3)  Exploration Phase (learn missing predicate / gather info)
        # ------------------------------------------------------------------
        # The exploration policy: iterate over every object, try to move to
        # it, then attempt a pick.  If pick succeeds and a disposal site
        # exists, place the object into the bin.
        #
        # NOTE:  We solely rely on the predefined skills – no new low-level
        #        control is introduced here.
        #
        done = False
        for obj_name, obj_pos in positions.items():
            if done:
                break

            # Skip the disposal site itself to avoid recursive pick/place
            if disposal_site_name and obj_name == disposal_site_name:
                continue

            print(f"\n----- [Exploration] Handling object '{obj_name}' -----")

            # Step 1: move close to the object (approach along +Z by default)
            obs, reward, done = _safe_call(
                move,
                env,
                task,
                target_pos=np.array(obj_pos),
                approach_distance=0.25,
                threshold=0.02,
                approach_axis="z",
                timeout=10.0,
            )
            if done:
                break

            # Step 2: attempt to pick the object
            obs, reward, done = _safe_call(
                pick,
                env,
                task,
                target_pos=np.array(obj_pos),
                approach_distance=0.15,
                threshold=0.01,
                approach_axis="z",
                timeout=10.0,
            )
            if done or obs is None:
                # Either task ended or pick failed – continue to next object
                continue

            # Step 3:  If we successfully picked and have a disposal site,
            #          place the object there
            if disposal_site_pos is not None:
                print(f"[Exploration] Placing '{obj_name}' at the disposal site.")
                obs, reward, done = _safe_call(
                    place,
                    env,
                    task,
                    target_pos=disposal_site_pos,
                    approach_distance=0.20,
                    threshold=0.01,
                    approach_axis="z",
                    timeout=10.0,
                )
                if done:
                    break
            else:
                print("[Exploration] No disposal site detected – keeping object held.")

            # Optional: rotate or pull actions could be added here for drawers
            if "drawer" in obj_name.lower():
                # Try to rotate gripper (arbitrary example – 90 degree turn)
                print(f"[Exploration] Attempting 'rotate' on drawer-like object '{obj_name}'.")
                _safe_call(
                    rotate,
                    env,
                    task,
                    gripper_name="gripper",  # assuming skill signature matches
                    from_angle="zero_deg",
                    to_angle="ninety_deg",
                )
                # Attempt pull (open drawer)
                _safe_call(
                    pull,
                    env,
                    task,
                    target_name=obj_name,
                )

        # ------------------------------------------------------------------
        # 4)  All objects processed or task finished
        # ------------------------------------------------------------------
        print("[Result] Exploration phase finished.")

    finally:
        # ------------------------------------------------------------------
        # Always shutdown environment (even if errors occurred)
        # ------------------------------------------------------------------
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()