import numpy as np
from scipy.spatial.transform import Rotation as R
from pyrep.objects.shape import Shape

from env import setup_environment, shutdown_environment

# bring the low-level skills into scope
from skill_code import rotate, move, pick, pull, place

# utilities for optional video recording
from video import init_video_writers, recording_step, recording_get_observation

# helper that returns a dict  {name: (x, y, z)}
from object_positions import get_object_positions


def _fetch_position(name, cache):
    """
    Convert an object / waypoint name into a 3-D numpy position.

    1) Try the cached dictionary provided by get_object_positions().
    2) If not cached, fall back to querying the simulator directly.
    """
    if name in cache:
        return np.asarray(cache[name], dtype=float)

    try:
        return np.asarray(Shape(name).get_position(), dtype=float)
    except Exception as exc:
        raise RuntimeError(f"[Task] Could not obtain position for `{name}`") from exc


def run_skeleton_task() -> None:
    print("========== [Task] START ==========")

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

        # ---- optional video capture --------------------------------
        init_video_writers(obs)
        task.step            = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # ---- static positions -------------------------------------
        pos_dict   = get_object_positions()
        pos        = lambda n: _fetch_position(n, pos_dict)

        ninety_deg_quat = R.from_euler('z', 90, degrees=True).as_quat()

        # drawer related way-points
        target_side   = pos('bottom_side_pos')
        target_anchor = pos('bottom_anchor_pos')

        # tomatoes and goal plate
        tomato1_pos   = pos('tomato1')
        tomato2_pos   = pos('tomato2')
        plate_pos     = pos('plate')

        # ---------------------------------------------------------------
        #  Oracle plan execution
        #  NOTE: steps 1-7 are frozen – do not edit
        # ---------------------------------------------------------------
        done   = False
        reward = 0.0

        # [Frozen Code Start]
        obs, reward, done = rotate(env, task, ninety_deg_quat)
        obs, reward, done = move(env, task, target_side)
        obs, reward, done = move(env, task, target_anchor)
        obs, reward, done = pick(env, task, target_anchor,
                                 approach_distance=0.05,  # small because handle is nearby
                                 approach_axis='-y')
        obs, reward, done = pull(env, task,
                                 pull_distance=0.15,       # heuristic pull length
                                 pull_axis='x')
        obs, reward, done = pick(env, task, tomato1_pos,
                                 approach_distance=0.15,
                                 approach_axis='z')
        obs, reward, done = place(env, task, plate_pos,
                                  approach_distance=0.15,
                                  approach_axis='z')
        # [Frozen Code End]

        if done:
            print("[Task] Task ended during the frozen section!")
            return

        # ---------------------------------------------------------------
        #  Step-8 : pick up tomato2 from the table
        # ---------------------------------------------------------------
        print("\n[Plan] Step-8 : pick tomato2")
        obs, reward, done = pick(
            env, task,
            tomato2_pos,
            approach_distance=0.15,
            approach_axis='z'
        )
        if done:
            print("[Task] Task ended after picking tomato2!")
            return

        # ---------------------------------------------------------------
        #  Step-9 : place tomato2 on the plate
        # ---------------------------------------------------------------
        print("\n[Plan] Step-9 : place tomato2 on plate")
        obs, reward, done = place(
            env, task,
            plate_pos,
            approach_distance=0.15,
            approach_axis='z'
        )

        # ---------------------------------------------------------------
        #  Final status report
        # ---------------------------------------------------------------
        if done:
            print(f"[Task] Completed – reward: {reward}")
        else:
            print("[Task] Plan executed, but task not marked done – verify success criteria.")

    finally:
        # always shut down the simulator
        shutdown_environment(env)
        print("========== [Task] END ==========")


if __name__ == "__main__":
    run_skeleton_task()