# run_skeleton_task.py (Completed Implementation)

import inspect
import math
import time
from typing import Dict, Any, List

import numpy as np
from pyrep.objects.shape import Shape

from env import setup_environment, shutdown_environment

# Import every predefined skill exactly as-is
from skill_code import *     # noqa: F403,F401  (wild-card import is required by the grading system)

from video import init_video_writers, recording_step, recording_get_observation

# Helper module that the project provides – returns a dict {name: (x, y, z)}
from object_positions import get_object_positions


# -------------------  Internal Utility Helpers  ------------------- #
def _safe_skill_call(skill_fn, env, task, *dynamic_args, **dynamic_kwargs):
    """
    Call a predefined skill function while being robust to signature mismatches.
    Only the arguments that actually appear in the function signature will be
    forwarded, which prevents most run-time crashes due to wrong arity.
    """
    sig = inspect.signature(skill_fn)
    bound_args = {}

    # First, env & task are always attempted
    for name, arg in (("env", env), ("task", task)):
        if name in sig.parameters:
            bound_args[name] = arg

    # Next, forward positional extras if there is room (*args or matching names)
    positional_iter = iter(dynamic_args)
    for p_name, p in sig.parameters.items():
        if p_name in bound_args:
            continue
        if p.kind in (inspect.Parameter.POSITIONAL_ONLY,
                      inspect.Parameter.POSITIONAL_OR_KEYWORD):
            try:
                bound_args[p_name] = next(positional_iter)
            except StopIteration:
                # No more extras provided – stop filling positionals
                break

    # Finally, forward keyword extras that match
    for k, v in dynamic_kwargs.items():
        if k in sig.parameters:
            bound_args[k] = v

    try:
        print(f"[Skill] Calling {skill_fn.__name__} with args: {bound_args}")
        return skill_fn(**bound_args)   # noqa: F405  (skill_fn is dynamic)
    except Exception as exc:            # pylint: disable=broad-except
        print(f"[Skill] WARNING: call to {skill_fn.__name__} failed – {exc}")
        return None


def _compute_quat_about_z(angle_rad: float) -> np.ndarray:
    """
    Return a quaternion (x,y,z,w) that represents a rotation of 'angle_rad'
    about the global Z-axis.
    """
    half = angle_rad * 0.5
    return np.array([0.0, 0.0, math.sin(half), math.cos(half)], dtype=np.float32)


# -------------------------  Main Routine  ------------------------- #
def run_skeleton_task():
    print("===== Starting Skeleton Task =====")

    # === Environment Setup ===
    env, task = setup_environment()

    try:
        # Reset task, obtain initial description & first observation
        descriptions, obs = task.reset()
        print("[Env] Task descriptions:", descriptions)

        # Optional video recording support
        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)

        # === Gather Object Positions ===
        positions: Dict[str, np.ndarray] = get_object_positions()
        print("[Env] Detected objects from helper module:", list(positions.keys()))

        # Basic diagnostics to help discover missing/unknown predicates
        # From feedback we know `lock-known` is referenced but absent in the real domain.
        print("[Diagnostics] The domain may be missing predicate 'lock-known'. "
              "We will try to interact with any drawers to infer lock state.")

        # ******************************************************
        # *   SIMPLE EXPLORATION STRATEGY                      *
        # *   1) Locate an object that looks like a drawer     *
        # *   2) Rotate the gripper by 90° (skill: rotate)     *
        # *   3) Attempt to pull it open (skill: pull)         *
        # *   4) Report any success / failure                  *
        # ******************************************************
        drawer_candidates: List[str] = [name for name in positions if "drawer" in name.lower()]

        if not drawer_candidates:
            print("[Exploration] No obvious drawer object names – skipping exploration.")
        else:
            target_drawer_name = drawer_candidates[0]
            drawer_pos = positions[target_drawer_name]
            print(f"[Exploration] Investigating drawer: {target_drawer_name} at {drawer_pos}")

            # Step 1) Move to the drawer area if a move skill that only
            # requires (env, task) exists.  If the signature needs more,
            # _safe_skill_call will gracefully skip it.
            _safe_skill_call(move, env, task, target_drawer_name)   # noqa: F405

            # Step 2) Rotate gripper by +90° about Z so that the fingers
            # are perpendicular to the drawer handle (common requirement).
            ninety_deg_quat = _compute_quat_about_z(math.pi / 2.0)
            _safe_skill_call(rotate, env, task, ninety_deg_quat)    # noqa: F405

            # Step 3) Pull the drawer – we do not know its exact signature
            # so we forward the drawer name as a dynamic positional argument.
            pull_result = _safe_skill_call(pull, env, task, target_drawer_name)  # noqa: F405

            # Step 4) Simple success heuristic: the pull skill often returns
            # (obs, reward, done).  If `done` becomes True or reward > 0, we
            # assume success.  Otherwise, we log failure so that the user can
            # analyse the missing predicate (`lock-known`).
            if isinstance(pull_result, tuple) and len(pull_result) >= 3:
                _obs, _reward, _done = pull_result[:3]
                if _done or (_reward is not None and _reward > 0.0):
                    print("[Exploration] Pull appears to have succeeded – "
                          "drawer is probably unlocked/open!")
                else:
                    print("[Exploration] Drawer never opened – lock status unknown "
                          "(this is consistent with a missing 'lock-known' predicate).")
            else:
                print("[Exploration] Could not determine outcome of pull – "
                      "return value:", pull_result)

        # ===== End of user plan / exploration =====
        print("===== Exploration phase complete =====")

    finally:
        # Ensure that we always shut down the simulation cleanly
        shutdown_environment(env)

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


if __name__ == "__main__":
    run_skeleton_task()