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

from env import setup_environment, shutdown_environment

# Import every predefined primitive exactly as provided by the platform
from skill_code import *          # noqa: F401,F403 – star-import on purpose
from video import (
    init_video_writers,
    recording_step,
    recording_get_observation,
)

from object_positions import get_object_positions


# ------------- Helper utilities (do NOT redefine any primitives) -------------
def _quat_from_euler(deg_xyz):
    """Utility: build an xyzw quaternion from Euler degrees (xyz order)."""
    return R.from_euler("xyz", deg_xyz, degrees=True).as_quat()


def _safe_lookup(dct, key):
    """Return dct[key] if it exists, otherwise raise a descriptive error."""
    if key not in dct:
        raise RuntimeError(
            f"[run_oracle_plan] Required object name ‘{key}’ not found in "
            "get_object_positions() dictionary. Available keys:\n"
            f"{list(dct.keys())}"
        )
    return dct[key]


# =============================  Main entry point  =============================
def run_oracle_plan():
    """
    Execute the oracle plan that must follow the exact primitive sequence:
        rotate → move → move → pick → pull → pick → place

    Logical description:
        1) Rotate the gripper 90 deg about the Z-axis so the fingers are
           parallel to the drawer face.
        2) Move to the BOTTOM drawer’s “side” pose.
        3) Move onto the drawer handle (“anchor” pose).
        4) Grasp the bottom-drawer handle.
        5) Pull the drawer straight out to open it.
        6) Pick the rubbish from the table.
        7) Place the rubbish into the bin.
    """
    print("\n==========  [run_oracle_plan]  START ==========")

    # ------------------------------------------------------------------
    #  Environment initialisation
    # ------------------------------------------------------------------
    env, task = setup_environment()
    try:
        # Task reset returns (low-level descriptions, RLBench observation)
        _, obs = task.reset()

        # Initialise optional video writer utilities
        init_video_writers(obs)

        # Wrap task.step / task.get_observation so every frame is captured
        task.step = recording_step(task.step)
        task.get_observation = recording_get_observation(task.get_observation)

        # ------------------------------------------------------------------
        #  Fetch all relevant Cartesian positions (x, y, z)
        # ------------------------------------------------------------------
        positions = get_object_positions()

        bottom_side_pos   = _safe_lookup(positions, "bottom_side_pos")
        bottom_anchor_pos = _safe_lookup(positions, "bottom_anchor_pos")
        rubbish_pos       = _safe_lookup(positions, "rubbish")
        bin_pos           = _safe_lookup(positions, "bin")

        # ------------------------------------------------------------------
        #  Oracle plan (7 steps – must match the specification exactly)
        # ------------------------------------------------------------------
        done = False

        # 1) rotate – align gripper 90° about Z
        target_quat = _quat_from_euler([0, 0, 90])
        print("\n[Plan-Step 1] rotate → 90° about Z")
        obs, reward, done = rotate(env, task, target_quat)
        if done:
            print("[Early-Exit] Task ended during rotate.")
            return

        # 2) move – approach the bottom drawer from the side
        print("\n[Plan-Step 2] move → bottom_side_pos:", bottom_side_pos)
        obs, reward, done = move(env, task, bottom_side_pos)
        if done:
            print("[Early-Exit] Task ended during move-to-side.")
            return

        # 3) move – onto the drawer handle (anchor position)
        print("\n[Plan-Step 3] move → bottom_anchor_pos:", bottom_anchor_pos)
        obs, reward, done = move(env, task, bottom_anchor_pos)
        if done:
            print("[Early-Exit] Task ended during move-to-anchor.")
            return

        # 4) pick – grasp the drawer handle
        print("\n[Plan-Step 4] pick → drawer handle (anchor)")
        obs, reward, done = pick(
            env,
            task,
            target_pos=bottom_anchor_pos,
            approach_distance=0.05,   # a bit closer for small handle
            approach_axis='-y'        # typical drawer handle approach
        )
        if done:
            print("[Early-Exit] Task ended during pick-drawer.")
            return

        # 5) pull – pull drawer straight out along +X by 0.12 m
        print("\n[Plan-Step 5] pull → distance 0.12 m along +X")
        obs, reward, done = pull(env, task, pull_distance=0.12, pull_axis="x")
        if done:
            print("[Early-Exit] Task ended during pull.")
            return

        # 6) pick – pick the rubbish from the table
        print("\n[Plan-Step 6] pick → rubbish")
        obs, reward, done = pick(env, task, target_pos=rubbish_pos)
        if done:
            print("[Early-Exit] Task ended during pick-rubbish.")
            return

        # 7) place – drop the rubbish into the bin
        print("\n[Plan-Step 7] place → bin")
        obs, reward, done = place(env, task, target_pos=bin_pos)
        if done:
            print("[Success] Task completed – rubbish placed in bin.")
        else:
            print("[Info] Plan finished but task signals done=False – "
                  "verify success criterion in the environment.")

    finally:
        # Guarantee simulator shuts down even if an exception is raised
        shutdown_environment(env)

    print("==========  [run_oracle_plan]  END ==========")


# Allow CLI execution
if __name__ == "__main__":
    run_oracle_plan()