"""MuJoCo simulator interface for LLM-driven demonstration generation."""

from __future__ import annotations

import logging
from pathlib import Path

import numpy as np

from orbit.generator.base import BaseSimulator

logger = logging.getLogger(__name__)


class MuJoCoSimulator(BaseSimulator):
    """Interface to MuJoCo for LLM-driven demonstration generation.

    Requires: pip install mujoco>=3.0 (orbit-robotics[sim])
    """

    def __init__(self, model_path: str):
        try:
            import mujoco
        except ImportError:
            raise ImportError(
                "MuJoCo is required for simulation. "
                "Install with: pip install orbit-robotics[sim]"
            )

        self._mujoco = mujoco
        model_path = str(Path(model_path).resolve())
        self.model = mujoco.MjModel.from_xml_path(model_path)
        self.data = mujoco.MjData(self.model)
        self._initial_qpos = self.data.qpos.copy()
        self._initial_qvel = self.data.qvel.copy()

    def get_state(self) -> dict:
        """Return current simulator state."""
        return {
            "qpos": self.data.qpos.copy(),
            "qvel": self.data.qvel.copy(),
            "site_xpos": self.data.site_xpos.copy(),
            "time": self.data.time,
        }

    def execute_action(self, action: np.ndarray) -> dict:
        """Execute one action step and return new state."""
        n_ctrl = min(len(action), self.model.nu)
        self.data.ctrl[:n_ctrl] = action[:n_ctrl]
        self._mujoco.mj_step(self.model, self.data)
        return self.get_state()

    def reset(self) -> dict:
        """Reset to initial state."""
        self._mujoco.mj_resetData(self.model, self.data)
        self.data.qpos[:] = self._initial_qpos
        self.data.qvel[:] = self._initial_qvel
        self._mujoco.mj_forward(self.model, self.data)
        return self.get_state()

    def get_api_description(self) -> str:
        """Return API description for the LLM agent."""
        return (
            f"MuJoCo simulator with {self.model.nq} position DOFs, "
            f"{self.model.nv} velocity DOFs, {self.model.nu} actuators.\n"
            f"Control range: {self.model.actuator_ctrlrange.tolist()}\n\n"
            f"Available API:\n"
            f"  sim.get_state() -> dict with keys: qpos, qvel, site_xpos, time\n"
            f"  sim.execute_action(action: np.ndarray) -> new state dict\n"
            f"  sim.reset() -> initial state dict\n"
            f"  action shape: ({self.model.nu},)\n"
        )

    def check_success(self, task: str, state: dict) -> bool:
        """Basic success check — override for task-specific checks."""
        # Default: can't determine success without task-specific logic
        return False
