"""Exo+arm MuJoCo render — RGB + segmentation mask.

One function: :func:`render_arm_with_exo`. Handles texturedir patch,
``cam_fovy`` from K, and the post-render principal-point warp.

Per-resolution model+renderer cache: building the model is ~200 ms; we
keep one per (W, H). The GL context is created in whichever thread
calls _build_ctx first and stays bound to that thread (mujoco EGL is
thread-local). If you need to render from a worker thread, build once
on the main thread first, or accept the second context creation.
"""
from __future__ import annotations

import os
import re
import sys
from pathlib import Path

import cv2
import numpy as np

# Headless-friendly GL backend. Must be set before mujoco is imported.
os.environ.setdefault("MUJOCO_GL", "egl")


_CACHE: dict[tuple[int, int], dict] = {}


def _exo_dir() -> str:
    return str(Path(__file__).resolve().parents[1] / "raiden_fork" /
               "third_party" / "exo_redo")


def _build_ctx(W: int, H: int) -> dict:
    if (W, H) in _CACHE:
        return _CACHE[(W, H)]
    if _exo_dir() not in sys.path:
        sys.path.insert(0, _exo_dir())
    import mujoco
    from ExoConfigs.yam_exo import YAM_BASE_ONLY_CONFIG
    from exo_utils import get_link_poses_from_robot, position_exoskeleton_meshes

    cfg = YAM_BASE_ONLY_CONFIG
    xml = cfg.xml.replace(
        '<compiler angle="radian"',
        f'<compiler angle="radian" balanceinertia="true" '
        f'texturedir="{_exo_dir()}/"',
    )
    xml = re.sub(r'<include file="[^"]*background[^"]*"\s*/>', '', xml)
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)
    mujoco.mj_forward(model, data)
    link_poses = get_link_poses_from_robot(cfg, model, data)
    position_exoskeleton_meshes(cfg, model, data, link_poses)

    arm_idx = []
    for jn in ("joint1", "joint2", "joint3", "joint4", "joint5", "joint6"):
        jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, jn)
        if jid >= 0:
            arm_idx.append(int(model.jnt_qposadr[jid]))
    grip_idx = []
    for jn in ("left_finger", "right_finger", "joint7", "joint8"):
        jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, jn)
        if jid >= 0:
            rng = model.jnt_range[jid]
            sign = +1.0 if abs(rng[1]) >= abs(rng[0]) else -1.0
            stroke = float(max(abs(rng[0]), abs(rng[1])))
            grip_idx.append((int(model.jnt_qposadr[jid]), sign, stroke))

    cam_id = model.cam('estimated_camera').id
    gl = mujoco.GLContext(W, H)
    gl.make_current()
    renderer = mujoco.Renderer(model, height=H, width=W)
    ctx = dict(model=model, data=data, renderer=renderer, cam_id=cam_id,
               arm_idx=arm_idx, grip_idx=grip_idx, gl=gl)
    _CACHE[(W, H)] = ctx
    return ctx


def render_arm_with_exo(
    joints7: np.ndarray,
    T_cam_in_rbase: np.ndarray,
    K: np.ndarray,
    W: int,
    H: int,
) -> tuple[np.ndarray, np.ndarray]:
    """Render arm+exo from a cam at ``T_cam_in_rbase``. Returns (rgb, mask).

    rgb: (H, W, 3) uint8. mask: (H, W) uint8 0/1 where the model is visible.
    Both are warped to compensate for MuJoCo's center-principal-point
    assumption (using K's cx, cy).
    """
    import mujoco
    ctx = _build_ctx(W, H)
    model, data, renderer, cam_id = (
        ctx["model"], ctx["data"], ctx["renderer"], ctx["cam_id"]
    )

    # Field of view from K (scaled to render H).
    model.cam_fovy[cam_id] = float(np.degrees(2 * np.arctan(H / (2 * K[1, 1]))))

    # Joint state (6 arm + gripper normalized [0,1]).
    for idx, q in zip(ctx["arm_idx"], joints7[:len(ctx["arm_idx"])]):
        data.qpos[idx] = float(q)
    if ctx["grip_idx"] and len(joints7) >= 7:
        g = float(np.clip(joints7[6], 0.0, 1.0))
        for qidx, sign, stroke in ctx["grip_idx"]:
            data.qpos[qidx] = sign * g * stroke
    mujoco.mj_forward(model, data)

    # Camera pose (convention matches raiden.calibration.exo_calibrate).
    T_link_in_cam = np.linalg.inv(np.asarray(T_cam_in_rbase, dtype=np.float64))
    data.cam_xpos[cam_id] = np.linalg.inv(T_link_in_cam)[:3, 3]
    flip = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
    data.cam_xmat[cam_id] = (flip @ T_link_in_cam[:3, :3]).T.reshape(-1)

    # RGB pass.
    renderer.update_scene(data, camera=cam_id)
    rgb = renderer.render()

    # Segmentation pass → mask.
    renderer.enable_segmentation_rendering()
    renderer.update_scene(data, camera=cam_id)
    seg = renderer.render()
    renderer.disable_segmentation_rendering()
    mask = (seg[:, :, 0].astype(np.int32) != -1).astype(np.uint8)

    # Principal-point warp.
    cx, cy = float(K[0, 2]), float(K[1, 2])
    dx, dy = int(round(cx - W / 2)), int(round(cy - H / 2))
    if dx or dy:
        M = np.float32([[1, 0, dx], [0, 1, dy]])
        rgb = cv2.warpAffine(rgb, M, (W, H), flags=cv2.INTER_LINEAR)
        mask = cv2.warpAffine(mask, M, (W, H), flags=cv2.INTER_NEAREST)

    return rgb, mask