import numpy as np
from scipy.spatial.transform import Rotation as R

from pyrep.objects.shape import Shape                   # kept from skeleton (may be unused)
from pyrep.objects.proximity_sensor import ProximitySensor  # kept from skeleton (may be unused)

from env import setup_environment, shutdown_environment

# KEEP the original wildcard-import line from the skeleton unchanged
from skill_code import *         # gives us rotate, move, pick, pull, place, …

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


# --------------------------------------------------------------------------- #
#                             Utility  helpers                                #
# --------------------------------------------------------------------------- #
def _safe_skill_call(skill_fn, *args, **kwargs):
    """Execute a skill while printing useful debug information."""
    try:
        print(f"[Plan] → Calling skill: {skill_fn.__name__}")
        return skill_fn(*args, **kwargs)
    except Exception as exc:
        print(f"[Error] Skill {skill_fn.__name__} failed: {exc}")
        raise


def _lookup_position(pos_dict, *candidate_keys):
    """
    Return the first object position that exists in pos_dict among candidate_keys.
    Raises RuntimeError if none are found.
    """
    for key in candidate_keys:
        if key in pos_dict:
            print(f"[Lookup] Using position for '{key}': {pos_dict[key]}")
            return pos_dict[key]
    raise RuntimeError(f"Missing object position – tried keys {candidate_keys}")


# --------------------------------------------------------------------------- #
#                        Main Task–specific routine                           #
# --------------------------------------------------------------------------- #
def run_tomato_drawer_task():
    """
    Executes the oracle plan required by the specification:
      1) rotate gripper from zero_deg to ninety_deg
      2) move to the side of the bottom drawer
      3) move to the anchor/handle of the bottom drawer
      4) pick the drawer handle (grasp)
      5) pull the drawer open
      6) pick the first tomato from the table
      7) place the first tomato on the plate
      8) pick the second tomato
      9) place the second tomato on the plate
    """
    print("===== Starting Tomato-Drawer Task =====")

    # ----------------------------------------------------------------------- #
    #  Environment initialisation                                             #
    # ----------------------------------------------------------------------- #
    env, task = setup_environment()
    try:
        # Reset the RLBench task
        _, obs = task.reset()

        # Optional video recording – keep as light overhead
        init_video_writers(obs)
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # ------------------------------------------------------------------- #
        #      Gather all relevant object positions from the scene            #
        # ------------------------------------------------------------------- #
        positions = get_object_positions()           # { name : np.ndarray([x,y,z]) }

        # Drawer related positions
        pos_bottom_side   = _lookup_position(
            positions, "bottom_side_pos", "bottom_side", "side-pos-bottom"
        )
        pos_bottom_anchor = _lookup_position(
            positions, "bottom_anchor_pos", "bottom_anchor", "anchor-pos-bottom"
        )

        # Tomatoes and plate
        pos_tomato1 = _lookup_position(positions, "item1", "tomato1")
        pos_tomato2 = _lookup_position(positions, "item2", "tomato2")
        pos_plate   = _lookup_position(positions, "plate", "plate_pos")

        # ------------------------------------------------------------------- #
        #       Execute the oracle plan, step-by-step                         #
        # ------------------------------------------------------------------- #

        # Step-1: rotate gripper 90° about Z
        quat_90deg_z = R.from_euler("xyz", [0.0, 0.0, np.pi / 2.0]).as_quat()
        obs, reward, done = _safe_skill_call(
            rotate, env, task, target_quat=quat_90deg_z
        )
        if done:
            print("[Plan] Finished immediately after rotate.")
            return

        # Step-2: move to drawer side position
        obs, reward, done = _safe_skill_call(
            move, env, task, target_pos=pos_bottom_side
        )
        if done:
            print("[Plan] Finished after move-to-side.")
            return

        # Step-3: move to drawer anchor (handle)
        obs, reward, done = _safe_skill_call(
            move, env, task, target_pos=pos_bottom_anchor
        )
        if done:
            print("[Plan] Finished after move-to-anchor.")
            return

        # Step-4: pick the drawer handle (gripper closes)
        obs, reward, done = _safe_skill_call(
            pick, env, task, target_pos=pos_bottom_anchor
        )
        if done:
            print("[Plan] Finished after pick-drawer.")
            return

        # Step-5: pull drawer open – pull 0.20 m along +X
        obs, reward, done = _safe_skill_call(
            pull, env, task, pull_distance=0.20, pull_axis="x"
        )
        if done:
            print("[Plan] Finished after pull.")
            return

        # Step-6: pick first tomato
        obs, reward, done = _safe_skill_call(
            pick, env, task, target_pos=pos_tomato1
        )
        if done:
            print("[Plan] Finished after picking tomato-1.")
            return

        # Step-7: place first tomato on plate
        obs, reward, done = _safe_skill_call(
            place, env, task, target_pos=pos_plate
        )
        if done:
            print("[Plan] Finished after placing tomato-1.")
            return

        # Step-8: pick second tomato
        obs, reward, done = _safe_skill_call(
            pick, env, task, target_pos=pos_tomato2
        )
        if done:
            print("[Plan] Finished after picking tomato-2.")
            return

        # Step-9: place second tomato on plate
        obs, reward, done = _safe_skill_call(
            place, env, task, target_pos=pos_plate
        )

        # ------------------------------------------------------------------- #
        #                     Final evaluation                                #
        # ------------------------------------------------------------------- #
        if done:
            print("[Plan] Task completed successfully! Reward:", reward)
        else:
            print("[Plan] All plan steps executed, but done flag is False. "
                  "Reward:", reward)

    finally:
        # Always ensure the environment shuts down cleanly
        shutdown_environment(env)

    print("===== End Tomato-Drawer Task =====")


# --------------------------------------------------------------------------- #
#                             Script entry point                              #
# --------------------------------------------------------------------------- #
if __name__ == "__main__":
    run_tomato_drawer_task()