import numpy as np
from scipy.spatial.transform import Rotation as R
from pyrep.objects.shape import Shape          # kept for completeness; may be unused
from pyrep.objects.proximity_sensor import ProximitySensor  # kept for completeness

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 rotate, move, pick, pull, place
# ==============================================================================


def run_skeleton_task() -> None:
    """
    Execute the oracle plan that
      (1) opens an accessible drawer, and
      (2) transfers the two tomatoes onto the plate.

    The nine steps follow the provided specification exactly, using only the
    predefined skills: rotate, move, pick, pull, and place.
    """
    print("\n==========  Combined-Domain Task  ==========")

    # --------------------------------------------------------------------------
    # 1)  Environment initialisation
    # --------------------------------------------------------------------------
    env, task = setup_environment()
    try:
        descriptions, obs = task.reset()

        # ------------- (Optional) Enable video recording ----------------------
        try:
            init_video_writers(obs)
            original_step = task.step
            task.step = recording_step(original_step)
            original_get_obs = task.get_observation
            task.get_observation = recording_get_observation(original_get_obs)
        except Exception as e:
            # Video is nice-to-have – never fail the whole run
            print("[Warning] Video initialisation failed:", e)
        # ----------------------------------------------------------------------

        # ----------------------------------------------------------------------
        # 2)  Query all required 3-D positions from the helper module
        # ----------------------------------------------------------------------
        positions = get_object_positions()      # raises if a key is missing

        # Drawer related positions
        pos_bottom_side   = np.asarray(positions['bottom_side_pos'])
        pos_bottom_anchor = np.asarray(positions['bottom_anchor_pos'])

        # Tomatoes and plate
        pos_tomato1 = np.asarray(positions['tomato1'])
        pos_tomato2 = np.asarray(positions['tomato2'])
        pos_plate   = np.asarray(positions['plate'])

        # ----------------------------------------------------------------------
        # 3)  Oracle plan – steps 1‒9 (specification matched)
        # ----------------------------------------------------------------------
        done   = False   # RLBench “done” flag
        reward = 0.0     # placeholder

        # STEP-1: rotate(gripper, zero_deg → ninety_deg)
        if not done:
            print("\n[Plan-1] rotate gripper to 90° about z-axis")
            target_quat = R.from_euler('z', 90.0, degrees=True).as_quat()
            obs, reward, done = rotate(env, task, target_quat)

        # STEP-2: move-to-side (→ side-pos-bottom)
        if not done:
            print("\n[Plan-2] move to bottom_side_pos")
            obs, reward, done = move(env, task, pos_bottom_side)

        # STEP-3: move-to-anchor (side-pos-bottom → anchor-pos-bottom)
        if not done:
            print("\n[Plan-3] move to bottom_anchor_pos")
            obs, reward, done = move(env, task, pos_bottom_anchor)

        # STEP-4: pick-drawer – grab the handle at anchor-pos-bottom
        if not done:
            print("\n[Plan-4] pick drawer handle (bottom drawer)")
            obs, reward, done = pick(
                env, task,
                target_pos=pos_bottom_anchor,
                approach_distance=0.10,
                approach_axis='z'
            )

        # STEP-5: pull(gripper, bottom) – pull drawer straight out (+x)
        if not done:
            print("\n[Plan-5] pull drawer along +x (0.20 m)")
            obs, reward, done = pull(
                env, task,
                pull_distance=0.20,
                pull_axis='x'
            )

        # STEP-6: pick tomato1
        if not done:
            print("\n[Plan-6] pick tomato1")
            obs, reward, done = pick(
                env, task,
                target_pos=pos_tomato1,
                approach_distance=0.15,
                approach_axis='z'
            )

        # STEP-7: place tomato1 onto plate
        if not done:
            print("\n[Plan-7] place tomato1 on plate")
            obs, reward, done = place(
                env, task,
                target_pos=pos_plate,
                approach_distance=0.15,
                approach_axis='z'
            )

        # STEP-8: pick tomato2
        if not done:
            print("\n[Plan-8] pick tomato2")
            obs, reward, done = pick(
                env, task,
                target_pos=pos_tomato2,
                approach_distance=0.15,
                approach_axis='z'
            )

        # STEP-9: place tomato2 onto plate
        if not done:
            print("\n[Plan-9] place tomato2 on plate")
            obs, reward, done = place(
                env, task,
                target_pos=pos_plate,
                approach_distance=0.15,
                approach_axis='z'
            )

        # ----------------------------------------------------------------------
        # 4)  Summary
        # ----------------------------------------------------------------------
        if done:
            print("\n[Result] Task signalled completion early.  Reward =", reward)
        else:
            print("\n[Result] Oracle plan executed.  Verify goal state in simulation.")

    except KeyError as ke:
        print("[Error] Missing key in position dictionary:", ke)
    except Exception as e:
        print("[Error] Unexpected exception:", e)
        raise
    finally:
        # Always ensure clean teardown of the environment
        shutdown_environment(env)
        print("==========  Task finished – Environment closed  ==========")


if __name__ == "__main__":
    run_skeleton_task()