import time
import numpy as np

from env import setup_environment, shutdown_environment

# Import every predefined skill exactly as provided.
# NOTE: The wildcard is intentional and required by the skeleton.
from skill_code import *                     # noqa: F401,F403
from video import (                          # noqa: F401
    init_video_writers,
    recording_step,
    recording_get_observation,
)
from object_positions import get_object_positions


def _print_exception(e, prefix="[ERROR]"):
    """Small helper for uniform exception logging."""
    print(f"{prefix} {type(e).__name__}: {e}")


def _safe_call(skill_fn, *args, **kwargs):
    """
    Wrapper that executes a skill function, catching and printing any exception
    (but allowing the overall task to continue).  Returns (obs, reward, done).
    """
    try:
        return skill_fn(*args, **kwargs)          # noqa: E501 — we rely on skills’ exact signatures
    except Exception as exc:                      # pylint: disable=broad-except
        _print_exception(exc)
        # Fall back to original observation if something went wrong
        # (The task object is always the second positional argument.)
        task = args[1]
        return task.get_observation(), 0.0, False


def perform_exploration(env, task):
    """
    Very small “exploration” phase whose sole purpose is to make sure we
    understand / satisfy the ‘rotated’ predicate mentioned in the feedback.
    We simply rotate the gripper to exactly +90 deg about the x-axis once.
    """
    # A quaternion for +90° around x-axis (RLBench uses (x,y,z,w))
    target_quat = np.array([0.70710678, 0.0, 0.0, 0.70710678], dtype=np.float32)

    print("[Exploration] Ensuring ‘rotated’ predicate is achievable …")
    obs, reward, done = _safe_call(rotate, env, task, target_quat)

    if done:                                     # Task ended early (unlikely here).
        print("[Exploration] Task ended during exploration.")
    else:
        print("[Exploration] Completed rotation exploration step.")
    return obs, reward, done


def oracle_plan(env, task, positions):
    """
    Placeholder ‘oracle’ plan.
    Because no specific goal/problem is provided, we only demonstrate a minimal,
    syntactically-valid usage of the predefined skills.
    """
    done = False

    # 1) Exploration step (needed by feedback requirement).
    obs, reward, done = perform_exploration(env, task)
    if done:
        return

    # 2) Example manipulation — if we know at least one object name and location.
    #    We do *not* know real object names, so demonstrate with a best-effort
    #    approach: pick the first entry from the positions dictionary (if any).
    try:
        if positions:
            obj_name, obj_pos = next(iter(positions.items()))
            print(f"[Plan] Attempting to pick the first discovered object: {obj_name}")

            obs, reward, done = _safe_call(
                pick,
                env,
                task,
                target_pos=np.asarray(obj_pos, dtype=np.float32),
                approach_distance=0.15,
                max_steps=150,
                threshold=0.02,
                approach_axis="z",
                timeout=10.0,
            )
            if done:
                return

            # Place the object back down at roughly the same spot
            print("[Plan] Placing the object back to its original location …")
            obs, reward, done = _safe_call(
                place,
                env,
                task,
                target_pos=np.asarray(obj_pos, dtype=np.float32),
                approach_distance=0.20,
                max_steps=150,
                threshold=0.02,
                approach_axis="z",
                timeout=10.0,
            )
    except StopIteration:
        print("[Plan] No objects returned by get_object_positions — skipping pick/place.")

    # Additional plan steps (move, pull, etc.) would be added here once a full
    # PDDL-derived plan is provided.

    if done:
        print("[Plan] Task signalled termination during oracle plan.")
    else:
        print("[Plan] Oracle plan completed (simulation continues).")


def run_skeleton_task():
    """Generic skeleton runner for our RLBench task."""
    print("===== Starting Skeleton Task =====")

    env, task = setup_environment()          # Environment boot up
    try:
        descriptions, obs = task.reset()     # Reset task
        print("[Init] Task descriptions:", descriptions)

        # Video recording set-up (safe even if video is disabled).
        init_video_writers(obs)

        # Wrap task functions so every step frame is recorded.
        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)

        # Retrieve world-state helper info (object positions, etc.).
        positions = {}
        try:
            positions = get_object_positions()
            print("[Init] Retrieved object positions:", positions)
        except Exception as exc:             # pylint: disable=broad-except
            _print_exception(exc, "[WARNING] Failed to obtain object positions.")

        # Execute high-level oracle plan
        oracle_plan(env, task, positions)

    finally:
        # Always perform clean shutdown
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()