import numpy as np
from scipy.spatial.transform import Rotation as R
from pyrep.objects.shape import Shape                                  # noqa: F401 (side-effects only)
from pyrep.objects.proximity_sensor import ProximitySensor             # noqa: F401

from env import setup_environment, shutdown_environment

# Low-level primitives delivered in “skill_code.py”
from skill_code import rotate, move, pull, pick, place

# Optional helpers for automatic video capture / debugging
from video import init_video_writers, recording_step, recording_get_observation

# Helper that returns a dict {name → np.ndarray(position)}
from object_positions import get_object_positions


# --------------------------------------------------------------------------- #
#  Small utility – robust dict access                                         #
# --------------------------------------------------------------------------- #
def _safe_get(mapping, key, *, assume=None):
    """
    Robustly fetch a key from a dictionary.
    If the key is missing
        •  return `assume` if provided
        •  otherwise raise a clear RuntimeError.
    """
    if key not in mapping:
        if assume is not None:
            return assume
        raise RuntimeError(f"[run_task] Cannot find key '{key}' in object_positions.")
    return mapping[key]


# --------------------------------------------------------------------------- #
#  Main entry-point                                                           #
# --------------------------------------------------------------------------- #
def run_skeleton_task():
    """
    Oracle policy that fulfils the goal:
        1) Rotate and align the gripper with the bottom drawer handle.
        2) Open the (already unlocked) drawer.
        3) Pick two tomatoes and place them on the plate.
    Only the predefined skills (rotate, move, pull, pick, place) are used,
    exactly matching the 9-step specification.
    """
    print("===== Starting Skeleton Task =====")

    env, task = setup_environment()
    try:
        # ------------------------------------------------------------------ #
        #  Reset task & (optionally) enable video recording                   #
        # ------------------------------------------------------------------ #
        descriptions, obs = task.reset()

        init_video_writers(obs)                                 # comment out if not needed
        task.step            = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # ------------------------------------------------------------------ #
        #  Obtain all relevant positions                                     #
        # ------------------------------------------------------------------ #
        positions = get_object_positions()

        # Drawer related (bottom drawer assumed to be accessible)
        bottom_side_pos   = np.asarray(_safe_get(positions, 'bottom_side_pos'))
        bottom_anchor_pos = np.asarray(_safe_get(positions, 'bottom_anchor_pos'))

        # Items & goal receptacle
        item1_pos  = np.asarray(_safe_get(positions, 'item1'))
        item2_pos  = np.asarray(_safe_get(positions, 'item2'))
        plate_pos  = np.asarray(_safe_get(positions, 'plate'))

        # ------------------------------------------------------------------ #
        #  Oracle Plan (1‒9)                                                 #
        # ------------------------------------------------------------------ #
        done   = False
        reward = 0.0

        # STEP-1 ── rotate gripper from 0° → 90° (vertical fingers)
        print("\n[Step-1] rotate gripper to 90° about Z")
        quat_z_90 = R.from_euler('z', 90, degrees=True).as_quat()   # xyzw
        obs, reward, done = rotate(env, task, quat_z_90)
        if done:
            print("[run_task] Terminated during rotate.")
            return

        # STEP-2 ── move to the drawer’s side approach point
        print("\n[Step-2] move to bottom_side_pos")
        obs, reward, done = move(env, task, bottom_side_pos)
        if done:
            print("[run_task] Terminated during move to side position.")
            return

        # STEP-3 ── move straight into the anchor (handle) position
        print("\n[Step-3] move to bottom_anchor_pos")
        obs, reward, done = move(env, task, bottom_anchor_pos)
        if done:
            print("[run_task] Terminated during move to anchor.")
            return

        # STEP-4 ── grip the drawer handle  (substitutes PDDL pick-drawer)
        print("\n[Step-4] close fingers on the drawer handle")
        obs, reward, done = pick(env, task,
                                 target_pos=bottom_anchor_pos,
                                 approach_distance=0.02,      # tiny vertical tuck
                                 approach_axis='z')
        if done:
            print("[run_task] Terminated while gripping handle.")
            return

        # STEP-5 ── pull the drawer open 15 cm along the +X axis
        print("\n[Step-5] pull drawer outward (+X, 0.15 m)")
        obs, reward, done = pull(env, task,
                                 pull_distance=0.15,
                                 pull_axis='x')
        if done:
            print("[run_task] Terminated during pull.")
            return

        # Because the gripper just opened the drawer, it is still closed on the
        # handle.  The upcoming pick will reopen the fingers automatically,
        # releasing the handle implicitly – hence no extra “place” step needed.

        # STEP-6 ── pick tomato-1
        print("\n[Step-6] pick tomato-1")
        obs, reward, done = pick(env, task, target_pos=item1_pos)
        if done:
            print("[run_task] Terminated picking tomato-1.")
            return

        # STEP-7 ── place tomato-1 on plate
        print("\n[Step-7] place tomato-1 on plate")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("[run_task] Terminated after placing tomato-1.")
            return

        # STEP-8 ── pick tomato-2
        print("\n[Step-8] pick tomato-2")
        obs, reward, done = pick(env, task, target_pos=item2_pos)
        if done:
            print("[run_task] Terminated picking tomato-2.")
            return

        # STEP-9 ── place tomato-2 on plate
        print("\n[Step-9] place tomato-2 on plate")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("[run_task] Environment reports done=True – likely success!")
        else:
            print("[run_task] Plan executed; env did not signal done. Verify criteria.")

    except Exception as exc:
        print(f"[run_task] Exception encountered: {exc}")
        raise

    finally:
        shutdown_environment(env)

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


# --------------------------------------------------------------------------- #
#  Stand-alone execution guard                                                #
# --------------------------------------------------------------------------- #
if __name__ == "__main__":
    run_skeleton_task()