from __future__ import annotations

import re
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Dict, Iterable, Optional, Tuple

import numpy as np
import pinocchio as pin

from .dual_arm import default_dual_source_path, load_dual_arm_mounts
from .kinematics import default_urdf_path
from .types import physical_joint_limits


def _numbers(text: str, expected: int) -> np.ndarray:
    values = np.fromstring(text, sep=" ", dtype=float)
    if values.shape != (expected,):
        raise ValueError(f"expected {expected} numeric values, got {text!r}")
    return values


def _format(values: Iterable[float]) -> str:
    return " ".join(f"{float(value):.12g}" for value in values)


def _quaternion_from_rpy(rpy: np.ndarray) -> np.ndarray:
    quaternion = pin.Quaternion(pin.rpy.rpyToMatrix(*rpy))
    return np.array([quaternion.w, quaternion.x, quaternion.y, quaternion.z])


def _se3_attributes(transform: pin.SE3) -> Dict[str, str]:
    quaternion = pin.Quaternion(transform.rotation)
    return {
        "pos": _format(transform.translation),
        "quat": _format([quaternion.w, quaternion.x, quaternion.y, quaternion.z]),
    }


def _origin_attributes(parent: ET.Element) -> Dict[str, str]:
    origin = parent.find("origin")
    if origin is None:
        return {"pos": "0 0 0", "quat": "1 0 0 0"}
    xyz = _numbers(origin.get("xyz", "0 0 0"), 3)
    rpy = _numbers(origin.get("rpy", "0 0 0"), 3)
    return {"pos": _format(xyz), "quat": _format(_quaternion_from_rpy(rpy))}


def _mesh_asset_name(filename: str, prefix: str) -> str:
    stem = Path(filename).stem.lower()
    return f"{prefix}_{re.sub(r'[^a-z0-9_]+', '_', stem)}"


def _dual_obj_directory(dual_source_path: Path) -> Path:
    source_candidate = dual_source_path.parent / "dual_arm_obj"
    if source_candidate.is_dir():
        return source_candidate
    raise FileNotFoundError(
        f"dual-arm OBJ directory was not found beside {dual_source_path}"
    )


def _single_mesh_directory(single_urdf_path: Path) -> Path:
    candidates = (
        single_urdf_path.parent / "meshes",
        single_urdf_path.parent / "rm75_meshes",
    )
    for candidate in candidates:
        if candidate.is_dir():
            return candidate
    raise FileNotFoundError(
        f"RM75-B STL mesh directory was not found beside {single_urdf_path}"
    )


def _add_inertial(body: ET.Element, link: ET.Element) -> None:
    inertial = link.find("inertial")
    if inertial is None:
        raise ValueError(f"link {link.get('name')!r} has no inertial data")
    mass_element = inertial.find("mass")
    inertia_element = inertial.find("inertia")
    if mass_element is None or inertia_element is None:
        raise ValueError(f"link {link.get('name')!r} has incomplete inertial data")
    origin = inertial.find("origin")
    xyz = _numbers(origin.get("xyz", "0 0 0"), 3) if origin is not None else np.zeros(3)
    rpy = _numbers(origin.get("rpy", "0 0 0"), 3) if origin is not None else np.zeros(3)
    inertia_matrix = np.array(
        [
            [float(inertia_element.get("ixx")), float(inertia_element.get("ixy")), float(inertia_element.get("ixz"))],
            [float(inertia_element.get("ixy")), float(inertia_element.get("iyy")), float(inertia_element.get("iyz"))],
            [float(inertia_element.get("ixz")), float(inertia_element.get("iyz")), float(inertia_element.get("izz"))],
        ]
    )
    rotation = pin.rpy.rpyToMatrix(*rpy)
    inertia_matrix = rotation @ inertia_matrix @ rotation.T
    full_inertia = [
        inertia_matrix[0, 0],
        inertia_matrix[1, 1],
        inertia_matrix[2, 2],
        inertia_matrix[0, 1],
        inertia_matrix[0, 2],
        inertia_matrix[1, 2],
    ]
    ET.SubElement(
        body,
        "inertial",
        {
            "pos": _format(xyz),
            "mass": mass_element.get("value", "0"),
            "fullinertia": _format(full_inertia),
        },
    )


def _add_link_visual(
    body: ET.Element,
    link: ET.Element,
    arm: str,
    single_mesh_keys: Dict[str, str],
) -> None:
    visual = link.find("visual")
    if visual is None:
        return
    mesh = visual.find("geometry/mesh")
    if mesh is None:
        return
    filename = Path(mesh.get("filename", "")).name
    try:
        mesh_name = single_mesh_keys[filename]
    except KeyError as exc:
        raise ValueError(f"missing packaged single-arm mesh {filename!r}") from exc
    ET.SubElement(
        body,
        "geom",
        {
            "name": f"{arm}_{link.get('name')}_visual",
            "type": "mesh",
            "mesh": mesh_name,
            "material": f"{arm}_arm",
            "contype": "0",
            "conaffinity": "0",
            "group": "1",
            **_origin_attributes(visual),
        },
    )


def _add_arm(
    worldbody: ET.Element,
    arm: str,
    mount: pin.SE3,
    single_root: ET.Element,
    single_mesh_keys: Dict[str, str],
    gripper_mesh_name: str,
) -> None:
    links = {link.get("name"): link for link in single_root.findall("link")}
    joints_by_parent: Dict[str, list[ET.Element]] = {}
    for joint in single_root.findall("joint"):
        parent_element = joint.find("parent")
        if parent_element is None:
            continue
        joints_by_parent.setdefault(parent_element.get("link", ""), []).append(joint)

    limits = physical_joint_limits()
    base_link = links["base_link"]
    base_body = ET.SubElement(
        worldbody,
        "body",
        {"name": f"{arm}_base_link", **_se3_attributes(mount)},
    )
    _add_link_visual(base_body, base_link, arm, single_mesh_keys)

    def append_children(parent_body: ET.Element, parent_link_name: str) -> None:
        for joint in joints_by_parent.get(parent_link_name, []):
            child_element = joint.find("child")
            if child_element is None:
                raise ValueError(f"joint {joint.get('name')!r} has no child")
            child_name = child_element.get("link", "")
            child_link = links[child_name]
            body = ET.SubElement(
                parent_body,
                "body",
                {
                    "name": f"{arm}_{child_name}",
                    **_origin_attributes(joint),
                },
            )
            joint_index = int(joint.get("name", "joint_0").split("_")[-1]) - 1
            axis_element = joint.find("axis")
            axis = "0 0 1" if axis_element is None else axis_element.get("xyz", "0 0 1")
            ET.SubElement(
                body,
                "joint",
                {
                    "name": f"{arm}_joint_{joint_index + 1}",
                    "type": "hinge",
                    "axis": axis,
                    "range": _format(
                        [limits.lower[joint_index], limits.upper[joint_index]]
                    ),
                    "limited": "true",
                    "damping": "0",
                },
            )
            _add_inertial(body, child_link)
            _add_link_visual(body, child_link, arm, single_mesh_keys)
            if child_name == "link_7":
                ET.SubElement(
                    body,
                    "site",
                    {
                        "name": f"{arm}_flange",
                        "type": "sphere",
                        "size": "0.008",
                        "rgba": "0.1 0.9 0.2 1" if arm == "left" else "0.1 0.5 0.95 1",
                        "group": "2",
                    },
                )
                ET.SubElement(
                    body,
                    "geom",
                    {
                        "name": f"{arm}_gripper_visual",
                        "type": "mesh",
                        "mesh": gripper_mesh_name,
                        "pos": "0 0 0.092",
                        "quat": _format(_quaternion_from_rpy(np.array([-np.pi, np.pi, -np.pi]))),
                        "material": f"{arm}_gripper",
                        "contype": "0",
                        "conaffinity": "0",
                        "group": "1",
                    },
                )
            append_children(body, child_name)

    append_children(base_body, "base_link")


def build_normalized_dual_mjcf(
    single_urdf_path: Optional[Path | str] = None,
    dual_source_path: Optional[Path | str] = None,
) -> Tuple[str, Dict[str, bytes]]:
    """Build a canonical 14-DOF dual-arm MJCF and its in-memory mesh assets."""

    single_path = (
        Path(single_urdf_path) if single_urdf_path is not None else default_urdf_path()
    )
    dual_path = (
        Path(dual_source_path)
        if dual_source_path is not None
        else default_dual_source_path()
    )
    single_root = ET.parse(single_path).getroot()
    dual_root = ET.parse(dual_path).getroot()
    single_mesh_dir = _single_mesh_directory(single_path)
    dual_obj_dir = _dual_obj_directory(dual_path)

    mujoco_root = ET.Element("mujoco", {"model": "rm75_normalized_dual_stage2"})
    ET.SubElement(
        mujoco_root,
        "compiler",
        {
            "angle": "radian",
            "autolimits": "true",
            "inertiafromgeom": "false",
            "balanceinertia": "true",
        },
    )
    ET.SubElement(
        mujoco_root,
        "option",
        {"timestep": "0.002", "gravity": "0 0 -9.81", "integrator": "implicitfast"},
    )
    ET.SubElement(mujoco_root, "statistic", {"center": "0 0 0.65", "extent": "1.4"})
    visual = ET.SubElement(mujoco_root, "visual")
    ET.SubElement(
        visual,
        "global",
        {
            "azimuth": "90",
            "elevation": "-18",
            "offwidth": "1280",
            "offheight": "720",
        },
    )
    ET.SubElement(visual, "rgba", {"haze": "0.15 0.18 0.2 1"})

    assets_element = ET.SubElement(mujoco_root, "asset")
    ET.SubElement(assets_element, "material", {"name": "left_arm", "rgba": "0.82 0.84 0.86 1"})
    ET.SubElement(assets_element, "material", {"name": "right_arm", "rgba": "0.7 0.74 0.78 1"})
    ET.SubElement(assets_element, "material", {"name": "left_gripper", "rgba": "0.15 0.75 0.85 1"})
    ET.SubElement(assets_element, "material", {"name": "right_gripper", "rgba": "0.9 0.48 0.18 1"})
    ET.SubElement(assets_element, "material", {"name": "platform", "rgba": "0.28 0.31 0.34 1"})

    assets: Dict[str, bytes] = {}
    single_mesh_keys: Dict[str, str] = {}
    for mesh_path in sorted(single_mesh_dir.glob("*.STL")):
        asset_key = f"rm75_meshes/{mesh_path.name}"
        mesh_name = _mesh_asset_name(mesh_path.name, "rm75")
        assets[asset_key] = mesh_path.read_bytes()
        single_mesh_keys[mesh_path.name] = mesh_name
        ET.SubElement(assets_element, "mesh", {"name": mesh_name, "file": asset_key})

    dual_mesh_names: Dict[str, str] = {}
    for mesh_path in sorted(dual_obj_dir.glob("*.obj")):
        asset_key = f"dual_arm_obj/{mesh_path.name}"
        mesh_name = _mesh_asset_name(mesh_path.name, "dual")
        assets[asset_key] = mesh_path.read_bytes()
        dual_mesh_names[mesh_path.name] = mesh_name
        ET.SubElement(assets_element, "mesh", {"name": mesh_name, "file": asset_key})
    if len(dual_mesh_names) != 19:
        raise ValueError(f"expected 19 dual-arm OBJ assets, found {len(dual_mesh_names)}")

    worldbody = ET.SubElement(mujoco_root, "worldbody")
    ET.SubElement(
        worldbody,
        "light",
        {
            "name": "key_light",
            "pos": "0 -1.2 2.8",
            "dir": "0 0.4 -1",
            "diffuse": "0.9 0.9 0.9",
        },
    )
    ET.SubElement(
        worldbody,
        "camera",
        {
            "name": "overview",
            "pos": "0 -2.6 1.25",
            "xyaxes": "1 0 0 0 0.3 0.953939",
            "fovy": "45",
        },
    )
    ET.SubElement(
        worldbody,
        "geom",
        {
            "name": "floor",
            "type": "plane",
            "size": "2.5 2.5 0.05",
            "rgba": "0.12 0.14 0.16 1",
            "contype": "0",
            "conaffinity": "0",
            "group": "0",
        },
    )

    robot_base = dual_root.find("./link[@name='robot_base']")
    if robot_base is None:
        raise ValueError("dual-arm source URDF has no robot_base link")
    for visual_element in robot_base.findall("visual"):
        mesh_element = visual_element.find("geometry/mesh")
        if mesh_element is None:
            continue
        filename = Path(mesh_element.get("filename", "")).name
        ET.SubElement(
            worldbody,
            "geom",
            {
                "name": f"platform_{visual_element.get('name', filename)}",
                "type": "mesh",
                "mesh": dual_mesh_names[filename],
                "material": "platform",
                "contype": "0",
                "conaffinity": "0",
                "group": "1",
                **_origin_attributes(visual_element),
            },
        )

    mounts = load_dual_arm_mounts(dual_path)
    _add_arm(
        worldbody,
        "left",
        mounts.left_base,
        single_root,
        single_mesh_keys,
        dual_mesh_names["dual_arm_gripper1_vis_1.obj"],
    )
    _add_arm(
        worldbody,
        "right",
        mounts.right_base,
        single_root,
        single_mesh_keys,
        dual_mesh_names["dual_arm_gripper2_vis_1.obj"],
    )
    for arm, color in (
        ("left", "0.95 0.12 0.12 0.8"),
        ("right", "0.95 0.75 0.08 0.8"),
    ):
        marker = ET.SubElement(
            worldbody,
            "body",
            {"name": f"{arm}_target_marker", "mocap": "true", "pos": "0 0 0"},
        )
        ET.SubElement(
            marker,
            "geom",
            {
                "name": f"{arm}_target_marker_geom",
                "type": "sphere",
                "size": "0.018",
                "rgba": color,
                "contype": "0",
                "conaffinity": "0",
                "group": "2",
            },
        )

    ET.indent(mujoco_root, space="  ")
    return ET.tostring(mujoco_root, encoding="unicode"), assets