import numpy as np
from scipy.spatial.transform import Rotation as R
from pyrep.objects.shape import Shape            # kept – useful for debugging
from pyrep.objects.proximity_sensor import ProximitySensor

from env import setup_environment, shutdown_environment
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions

# === Pre-implemented Skills (DO NOT redefine) ===
from skill_code import move, rotate, pull, pick, place


def run_skeleton_task() -> None:
    """
    Execute the oracle plan:

      1) Rotate the wrist to 90 deg → allow side-way approach
      2) Move to the drawer’s side way-point
      3) Move to the drawer’s anchor (handle)
      4) Grasp the handle
      5) Pull the drawer open
      6) Pick up rubbish from the table
      7) Throw the rubbish into the bin
    """
    print("===== Starting Skeleton Task =====")

    # ------------------------------------------------------------------
    #  Environment set-up
    # ------------------------------------------------------------------
    env, task = setup_environment()

    try:
        # ‑- Reset & (optionally) start recording
        _, obs = task.reset()
        init_video_writers(obs)

        # Wrap RLBench step / observation so the video module can capture frames
        task.step            = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # ------------------------------------------------------------------
        #  Retrieve all required object poses from object_positions.py
        # ------------------------------------------------------------------
        # Returned format:  { name (str) : np.ndarray([x, y, z]) }
        positions = get_object_positions()

        # Drawer way-points (open the bottom drawer)
        side_pos_bottom   = positions.get('bottom_side_pos')
        anchor_pos_bottom = positions.get('bottom_anchor_pos')

        if side_pos_bottom is None or anchor_pos_bottom is None:
            raise RuntimeError(
                "Drawer waypoint positions 'bottom_side_pos' or "
                "'bottom_anchor_pos' missing in object_positions.py"
            )

        # Rubbish on the table – choose explicit key if present, else fall back
        rubbish_name = 'rubbish' if 'rubbish' in positions else 'item1'
        rubbish_pos  = positions.get(rubbish_name)
        if rubbish_pos is None:
            raise RuntimeError(
                f"Could not obtain a position for the rubbish object "
                f"'{rubbish_name}'."
            )

        # Bin position
        bin_pos = positions.get('bin')
        if bin_pos is None:
            raise RuntimeError(
                "Position for object 'bin' not supplied by "
                "object_positions.py."
            )

        # ==============================================================
        #                       EXECUTE ORACLE PLAN
        # ==============================================================

        done   = False
        reward = 0.0

        # ---------------- Step-1  ROTATE gripper to 90° ----------------
        print("\n[PLAN] Step-1  rotate  gripper to 90 deg about z-axis")
        quat_90_z = R.from_euler('z', np.pi / 2).as_quat()   # xyzw
        obs, reward, done = rotate(env, task, target_quat=quat_90_z)
        if done:
            print("[PLAN] Environment signalled completion after step-1.")
            return

        # ---------------- Step-2  MOVE to drawer side-waypoint ---------
        print("\n[PLAN] Step-2  move  to 'bottom_side_pos'")
        obs, reward, done = move(env, task, target_pos=side_pos_bottom)
        if done:
            print("[PLAN] Environment signalled completion after step-2.")
            return

        # ---------------- Step-3  MOVE to drawer anchor (handle) -------
        print("\n[PLAN] Step-3  move  to 'bottom_anchor_pos' (handle)")
        obs, reward, done = move(env, task, target_pos=anchor_pos_bottom)
        if done:
            print("[PLAN] Environment signalled completion after step-3.")
            return

        # ---------------- Step-4  PICK (grasp handle) ------------------
        # Using generic 'pick' skill – no explicit pick-drawer skill is in
        # the library, so we emulate the grasp of the drawer handle.
        print("\n[PLAN] Step-4  grasp drawer handle")
        obs, reward, done = pick(
            env, task,
            target_pos=anchor_pos_bottom,
            approach_distance=0.05,    # short approach, near handle
            approach_axis='-y'         # move forward towards the drawer
        )
        if done:
            print("[PLAN] Environment signalled completion after step-4.")
            return

        # ---------------- Step-5  PULL drawer open (+x) ---------------
        print("\n[PLAN] Step-5  pull  drawer by 0.30 m along +x")
        obs, reward, done = pull(
            env, task,
            pull_distance=0.30,
            pull_axis='x'
        )
        if done:
            print("[PLAN] Environment signalled completion after step-5.")
            return

        # Optionally: release the drawer handle (open gripper in place)
        print("\n[PLAN]   release drawer handle (open gripper)")
        current_pos = np.array(task.get_observation().gripper_pose[:3])
        obs, reward, done = place(
            env, task,
            target_pos=current_pos,
            approach_distance=0.00,
            approach_axis='z'
        )
        if done:
            print("[PLAN] Environment signalled completion after releasing handle.")
            return

        # ---------------- Step-6  PICK rubbish from table -------------
        print(f"\n[PLAN] Step-6  pick  '{rubbish_name}' from table")

        # First hover 10 cm above the rubbish to ensure collision-free descent
        rubbish_above = np.array(rubbish_pos) + np.array([0.0, 0.0, 0.10])
        obs, reward, done = move(env, task, target_pos=rubbish_above)
        if done:
            print("[PLAN] Environment signalled completion while approaching rubbish.")
            return

        obs, reward, done = pick(
            env, task,
            target_pos=rubbish_pos,
            approach_distance=0.15,
            approach_axis='-z'
        )
        if done:
            print("[PLAN] Environment signalled completion after picking rubbish.")
            return

        # ---------------- Step-7  PLACE rubbish into bin --------------
        print("\n[PLAN] Step-7  place rubbish into bin")
        bin_above = np.array(bin_pos) + np.array([0.0, 0.0, 0.10])
        obs, reward, done = move(env, task, target_pos=bin_above)
        if done:
            print("[PLAN] Environment signalled completion while approaching bin.")
            return

        obs, reward, done = place(
            env, task,
            target_pos=bin_pos,
            approach_distance=0.15,
            approach_axis='-z'
        )

        # ------------- PLAN COMPLETE ----------------------------------
        if done:
            print(f"\n===== Task completed successfully! Final reward = {reward} =====")
        else:
            print("\n===== Plan executed – environment did not set done=True. =====")

    finally:
        # Make sure the simulator is shut down under all circumstances
        shutdown_environment(env)
        print("===== End of Skeleton Task =====")


if __name__ == "__main__":
    run_skeleton_task()