import dataclasses
import os
import tempfile
import xml.etree.ElementTree as ET
from pathlib import Path

import mujoco
import numpy as np
import torch

from src.core.env import BaseEnv
from src.core.robot import RobotConfig
from src.core.runner import BaseRunner, BaseRunnerConfig


@dataclasses.dataclass
class MuJoCoRunnerConfig(BaseRunnerConfig):
    num_envs: int = 16
    device: str = "cpu"
    dt: float = 0.002
    substeps: int = 10


class ActuatorLimits:
    """Software limit-switch: cuts motor ctrl when a joint hits its range.

    The real robot has physical limit switches that kill motor current
    at the travel endpoints.  MuJoCo's built-in joint limits only apply
    a spring force — they don't zero the actuator signal.  This class
    replicates the hardware behavior.
    """

    def __init__(self, model: mujoco.MjModel) -> None:
        jnt_ids = model.actuator_trnid[:model.nu, 0]
        self.jnt_ids = jnt_ids
        self.limited = model.jnt_limited[jnt_ids].astype(bool)
        self.lo = model.jnt_range[jnt_ids, 0]
        self.hi = model.jnt_range[jnt_ids, 1]
        self.gear_sign = np.sign(model.actuator_gear[:model.nu, 0])

    def enforce(self, model: mujoco.MjModel, data: mujoco.MjData) -> None:
        """Zero ctrl that would push past joint limits (call every substep)."""
        if not self.limited.any():
            return
        pos = data.qpos[self.jnt_ids]
        signed_ctrl = self.gear_sign * data.ctrl[:model.nu]
        at_hi = self.limited & (pos >= self.hi) & (signed_ctrl > 0)
        at_lo = self.limited & (pos <= self.lo) & (signed_ctrl < 0)
        data.ctrl[at_hi | at_lo] = 0.0


# ── Public utilities ─────────────────────────────────────────────────


def load_mujoco_model(robot: RobotConfig) -> mujoco.MjModel:
    """Load a URDF (or MJCF) and apply robot.yaml settings.

    Single model-loading entry point for all MuJoCo-based code:
    training runners, MJX, and system identification.

    Two-step approach required because MuJoCo's URDF parser ignores
    ``<actuator>`` in the ``<mujoco>`` extension block:

      1. Load the URDF -> MuJoCo converts it to internal MJCF
      2. Export the MJCF XML, inject actuators + joint overrides, reload

    This keeps the URDF clean (re-exportable from CAD) -- all hardware
    tuning lives in ``robot.yaml``.
    """
    abs_path = robot.urdf_path.resolve()
    model_dir = abs_path.parent
    is_urdf = abs_path.suffix.lower() == ".urdf"

    # -- Step 1: Load URDF with meshdir injection --
    if is_urdf:
        tree = ET.parse(abs_path)
        root = tree.getroot()

        # MuJoCo's URDF parser strips directory prefixes from mesh
        # filenames, so we inject a <mujoco><compiler meshdir="..."/>
        # block.  The original URDF stays clean and simulator-agnostic.
        meshdir = None
        for mesh_el in root.iter("mesh"):
            fn = mesh_el.get("filename", "")
            parent = str(Path(fn).parent)
            if parent and parent != ".":
                meshdir = parent
                break
        if meshdir:
            mj_ext = ET.SubElement(root, "mujoco")
            ET.SubElement(mj_ext, "compiler", attrib={
                "meshdir": meshdir,
                "balanceinertia": "true",
            })

        # Write to a temp file (unique name for multiprocessing safety).
        fd, tmp_path = tempfile.mkstemp(
            suffix=".urdf", prefix="_mj_", dir=str(model_dir),
        )
        os.close(fd)
        try:
            tree.write(tmp_path, xml_declaration=True, encoding="unicode")
            model_raw = mujoco.MjModel.from_xml_path(tmp_path)
        finally:
            Path(tmp_path).unlink(missing_ok=True)
    else:
        model_raw = mujoco.MjModel.from_xml_path(str(abs_path))

    # If robot.yaml has no actuators/joints, we're done.
    if not robot.actuators and not robot.joints:
        return model_raw

    # -- Step 2: Export MJCF, inject actuators + joint overrides --
    fd, tmp_path = tempfile.mkstemp(
        suffix=".xml", prefix="_mj_", dir=str(model_dir),
    )
    os.close(fd)
    try:
        mujoco.mj_saveLastXML(tmp_path, model_raw)
        mjcf_str = Path(tmp_path).read_text()
        root = ET.fromstring(mjcf_str)

        # -- Inject actuators --
        if robot.actuators:
            act_elem = ET.SubElement(root, "actuator")
            for act in robot.actuators:
                attribs = {
                    "name": f"{act.joint}_{act.type}",
                    "joint": act.joint,
                    "gear": str(act.gear_avg),
                    "ctrlrange": f"{act.ctrl_range[0]} {act.ctrl_range[1]}",
                }

                # dyntype is only available on <general>, not on
                # shortcut elements like <motor>/<position>/<velocity>.
                use_general = act.filter_tau > 0

                if use_general:
                    attribs["dyntype"] = "filter"
                    attribs["dynprm"] = str(act.filter_tau)
                    attribs["gaintype"] = "fixed"
                    if act.type == "position":
                        attribs["biastype"] = "affine"
                        attribs["gainprm"] = str(act.kp)
                        attribs["biasprm"] = f"0 -{act.kp} -{act.kv}"
                    elif act.type == "velocity":
                        attribs["biastype"] = "affine"
                        attribs["gainprm"] = str(act.kp)
                        attribs["biasprm"] = f"0 0 -{act.kp}"
                    else:  # motor
                        attribs["biastype"] = "none"
                    ET.SubElement(act_elem, "general", attrib=attribs)
                else:
                    if act.type == "position":
                        attribs["kp"] = str(act.kp)
                        if act.kv > 0:
                            attribs["kv"] = str(act.kv)
                    elif act.type == "velocity":
                        attribs["kp"] = str(act.kp)
                    ET.SubElement(act_elem, act.type, attrib=attribs)

        # -- Apply joint overrides from robot.yaml --
        # For actuated joints with a motor model, MuJoCo damping/frictionloss
        # are set to 0 — the motor model handles them via qfrc_applied.
        joint_damping: dict[str, float] = {}
        joint_frictionloss: dict[str, float] = {}
        for act in robot.actuators:
            if act.has_motor_model:
                joint_damping[act.joint] = 0.0
                joint_frictionloss[act.joint] = 0.0
        joint_armature: dict[str, float] = {}
        for name, jcfg in robot.joints.items():
            if jcfg.damping is not None:
                joint_damping[name] = jcfg.damping
            if jcfg.armature is not None:
                joint_armature[name] = jcfg.armature
            if jcfg.frictionloss is not None:
                joint_frictionloss[name] = jcfg.frictionloss

        for body in root.iter("body"):
            for jnt in body.findall("joint"):
                name = jnt.get("name")
                if name in joint_damping:
                    jnt.set("damping", str(joint_damping[name]))
                if name in joint_armature:
                    jnt.set("armature", str(joint_armature[name]))
                if name in joint_frictionloss:
                    jnt.set("frictionloss", str(joint_frictionloss[name]))

        # Disable self-collision on all geoms.
        for geom in root.iter("geom"):
            geom.set("contype", "0")
            geom.set("conaffinity", "0")

        modified_xml = ET.tostring(root, encoding="unicode")
        Path(tmp_path).write_text(modified_xml)
        return mujoco.MjModel.from_xml_path(tmp_path)
    finally:
        Path(tmp_path).unlink(missing_ok=True)


# -- Runner -----------------------------------------------------------


class MuJoCoRunner(BaseRunner[MuJoCoRunnerConfig]):
    def __init__(self, env: BaseEnv, config: MuJoCoRunnerConfig):
        super().__init__(env, config)

    @property
    def num_envs(self) -> int:
        return self.config.num_envs

    @property
    def device(self) -> torch.device:
        return torch.device(self.config.device)

    def _sim_initialize(self, config: MuJoCoRunnerConfig) -> None:
        self._model = load_mujoco_model(self.env.robot)
        self._model.opt.timestep = config.dt
        self._data: list[mujoco.MjData] = [
            mujoco.MjData(self._model) for _ in range(config.num_envs)
        ]
        self._nq = self._model.nq
        self._nv = self._model.nv
        self._limits = ActuatorLimits(self._model)

        # Build motor model: list of (actuator_config, joint_qvel_index) for
        # actuators that have asymmetric motor dynamics.
        self._motor_actuators: list[tuple] = []
        for act in self.env.robot.actuators:
            if act.has_motor_model:
                jnt_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_JOINT, act.joint)
                qvel_idx = self._model.jnt_dofadr[jnt_id]
                self._motor_actuators.append((act, qvel_idx))

    def _sim_step(self, actions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        actions_np: np.ndarray = actions.cpu().numpy()

        # Apply per-actuator ctrl transform (deadzone + gear compensation)
        for act_idx, (act, _) in enumerate(self._motor_actuators):
            for env_idx in range(self.num_envs):
                actions_np[env_idx, act_idx] = act.transform_ctrl(
                    float(actions_np[env_idx, act_idx])
                )

        qpos_batch = np.zeros((self.num_envs, self._nq), dtype=np.float32)
        qvel_batch = np.zeros((self.num_envs, self._nv), dtype=np.float32)

        # Per-env domain-randomization scales (all 1.0 when DR is disabled).
        fr_scale = self._dr_scales["friction_scale"].cpu().numpy()
        dp_scale = self._dr_scales["damping_scale"].cpu().numpy()
        tq_scale = self._dr_scales["torque_scale"].cpu().numpy()

        for i, data in enumerate(self._data):
            # torque_scale emulates motor-constant / battery-voltage variation.
            data.ctrl[:] = actions_np[i] * tq_scale[i]
            for _ in range(self.config.substeps):
                # Apply asymmetric motor forces via qfrc_applied
                for act, qvel_idx in self._motor_actuators:
                    vel = data.qvel[qvel_idx]
                    ctrl = data.ctrl[0]  # TODO: generalise for multi-actuator
                    data.qfrc_applied[qvel_idx] = act.compute_motor_force(
                        vel, ctrl,
                        friction_scale=fr_scale[i],
                        damping_scale=dp_scale[i],
                    )
                self._limits.enforce(self._model, data)
                mujoco.mj_step(self._model, data)

            qpos_batch[i] = data.qpos
            qvel_batch[i] = data.qvel

        return (
            torch.from_numpy(qpos_batch).to(self.device),
            torch.from_numpy(qvel_batch).to(self.device),
        )

    def _sim_reset(self, env_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        ids = env_ids.cpu().numpy()

        # Env-defined initial-state distribution (shared with the MJX runner).
        qpos_lo, qpos_hi, qvel_lo, qvel_hi = self.env.initial_state_ranges(
            self._nq, self._nv,
        )

        for env_id in ids:
            data = self._data[env_id]
            mujoco.mj_resetData(self._model, data)

            data.qpos[:] += np.random.uniform(qpos_lo, qpos_hi)
            data.qvel[:] += np.random.uniform(qvel_lo, qvel_hi)
            data.ctrl[:] = 0.0

        # Full-batch return (see BaseRunner._sim_reset contract).
        qpos_batch = np.stack([d.qpos for d in self._data]).astype(np.float32)
        qvel_batch = np.stack([d.qvel for d in self._data]).astype(np.float32)
        return (
            torch.from_numpy(qpos_batch).to(self.device),
            torch.from_numpy(qvel_batch).to(self.device),
        )

    def _render_frame(self, env_idx: int = 0) -> np.ndarray:
        """Offscreen render of a single environment."""
        if not hasattr(self, "_offscreen_renderer"):
            self._offscreen_renderer = mujoco.Renderer(
                self._model, width=640, height=480,
            )
        mujoco.mj_forward(self._model, self._data[env_idx])
        self._offscreen_renderer.update_scene(self._data[env_idx])
        return self._offscreen_renderer.render().copy()
