from pyrep.objects.shape import Shape
from pyrep.objects.proximity_sensor import ProximitySensor
from env import setup_environment, shutdown_environment
from skill_code import move, pick, place, rotate, pull
from video import init_video_writers, recording_step, recording_get_observation
from object_positions import get_object_positions
import numpy as np

def run_skeleton_task():
    print("===== Starting Oracle Plan Execution =====")
    # === Environment Setup ===
    env, task = setup_environment()
    try:
        # Reset the task to its initial state and initialize video recording
        descriptions, obs = task.reset()
        init_video_writers(obs)
        # Wrap step and observation for recording
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # === Retrieve Object Positions and Orientations ===
        positions = get_object_positions()
        # We expect positions to contain:
        #   'zero_deg', 'ninety_deg'          : orientation quaternions (4‑element arrays)
        #   'bottom_side_pos', 'bottom_anchor_pos' : 3D positions for drawer handle
        #   'tomato1', 'tomato2', 'plate'     : 3D positions of the tomatoes and plate

        # === Oracle Plan ===

        # Step 1: Rotate gripper from zero_deg to ninety_deg
        print("[Task] Step 1: rotate(gripper, zero_deg, ninety_deg)")
        try:
            zero_quat   = np.array(positions['zero_deg'])
            ninety_quat = np.array(positions['ninety_deg'])
        except KeyError:
            raise RuntimeError("Missing orientation entries 'zero_deg' or 'ninety_deg' in positions.")
        obs, reward, done = rotate(env, task, target_quat=ninety_quat)
        if done: return

        # Step 2: move-to-side → use generic move to side position
        print("[Task] Step 2: move-to-side(gripper, bottom, nowhere-pos, side-pos-bottom)")
        try:
            side_pos = np.array(positions['bottom_side_pos'])
        except KeyError:
            raise RuntimeError("Missing 'bottom_side_pos' in positions.")
        obs, reward, done = move(env, task, target_pos=side_pos)
        if done: return

        # Step 3: move-to-anchor → use generic move to anchor position
        print("[Task] Step 3: move-to-anchor(gripper, bottom, side-pos-bottom, anchor-pos-bottom)")
        try:
            anchor_pos = np.array(positions['bottom_anchor_pos'])
        except KeyError:
            raise RuntimeError("Missing 'bottom_anchor_pos' in positions.")
        obs, reward, done = move(env, task, target_pos=anchor_pos)
        if done: return

        # Step 4: pick-drawer → pick drawer handle at anchor position
        print("[Task] Step 4: pick-drawer(gripper, bottom, anchor-pos-bottom)")
        obs, reward, done = pick(env, task, target_pos=anchor_pos)
        if done: return

        # Step 5: pull the drawer open
        print("[Task] Step 5: pull(gripper, bottom)")
        # We choose a nominal pull distance; adjust if needed
        pull_distance = 0.1
        obs, reward, done = pull(env, task, pull_distance=pull_distance, pull_axis='x')
        if done: return

        # Step 6: pick tomato1 from table
        print("[Task] Step 6: pick(tomato1, table)")
        try:
            tomato1_pos = np.array(positions['tomato1'])
        except KeyError:
            raise RuntimeError("Missing 'tomato1' in positions.")
        obs, reward, done = pick(env, task, target_pos=tomato1_pos)
        if done: return

        # Step 7: place tomato1 onto the plate
        print("[Task] Step 7: place(tomato1, plate)")
        try:
            plate_pos = np.array(positions['plate'])
        except KeyError:
            raise RuntimeError("Missing 'plate' in positions.")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done: return

        # Step 8: pick tomato2 from table
        print("[Task] Step 8: pick(tomato2, table)")
        try:
            tomato2_pos = np.array(positions['tomato2'])
        except KeyError:
            raise RuntimeError("Missing 'tomato2' in positions.")
        obs, reward, done = pick(env, task, target_pos=tomato2_pos)
        if done: return

        # Step 9: place tomato2 onto the plate
        print("[Task] Step 9: place(tomato2, plate)")
        obs, reward, done = place(env, task, target_pos=plate_pos)
        if done:
            print("[Task] Completed all steps with done=True. Final reward:", reward)
        else:
            print("[Task] Completed all steps. done=False. Final reward:", reward)

    except Exception as e:
        print("[Task] Exception encountered:", e)
    finally:
        # Ensure the environment is properly shutdown
        shutdown_environment(env)
    print("===== End of Oracle Plan Execution =====")


if __name__ == "__main__":
    run_skeleton_task()