from __future__ import annotations

import sysconfig
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import numpy as np
import pinocchio as pin

from .kinematics import RM75Kinematics
from .types import JointLimits, physical_joint_limits


def default_dual_source_path() -> Path:
    source_path = (
        Path(__file__).resolve().parents[2]
        / "models"
        / "dual_arm_mujoco_fixed.urdf"
    )
    if source_path.is_file():
        return source_path
    installed_candidates = [
        Path(sysconfig.get_path("data"))
        / "share"
        / "rm75_ik"
        / "models"
        / "dual_arm_mujoco_fixed.urdf"
    ]
    resolved = Path(__file__).resolve()
    if len(resolved.parents) > 4:
        installed_candidates.append(
            resolved.parents[4]
            / "share/rm75_ik/models/dual_arm_mujoco_fixed.urdf"
        )
    for installed_path in installed_candidates:
        if installed_path.is_file():
            return installed_path
    raise FileNotFoundError("dual_arm_mujoco_fixed.urdf was not found")


def _origin_to_se3(element: ET.Element) -> pin.SE3:
    origin = element.find("origin")
    if origin is None:
        return pin.SE3.Identity()
    xyz = np.fromstring(origin.get("xyz", "0 0 0"), sep=" ", dtype=float)
    rpy = np.fromstring(origin.get("rpy", "0 0 0"), sep=" ", dtype=float)
    if xyz.shape != (3,) or rpy.shape != (3,):
        raise ValueError(f"invalid URDF origin on element {element.get('name')!r}")
    return pin.SE3(pin.rpy.rpyToMatrix(*rpy), xyz)


@dataclass(frozen=True)
class DualArmMounts:
    left_base: pin.SE3
    right_base: pin.SE3
    right_visual_origin_delta_m: float


def load_dual_arm_mounts(source_urdf: Optional[Path | str] = None) -> DualArmMounts:
    path = Path(source_urdf) if source_urdf is not None else default_dual_source_path()
    root = ET.parse(path).getroot()
    joints = {joint.get("name"): joint for joint in root.findall("joint")}
    try:
        world_left_joint1 = _origin_to_se3(joints["joint1"])
        world_right_joint1 = _origin_to_se3(joints["joint8"])
    except KeyError as exc:
        raise ValueError("dual-arm source URDF must contain joint1 and joint8") from exc

    base_to_joint1 = pin.SE3(np.eye(3), np.array([0.0, 0.0, 0.2405]))
    left_base = world_left_joint1 * base_to_joint1.inverse()
    right_base = world_right_joint1 * base_to_joint1.inverse()

    right_visual = root.find(
        "./link[@name='robot_base']/visual[@name='base_link_right']"
    )
    if right_visual is None:
        visual_delta = float("nan")
    else:
        visual_pose = _origin_to_se3(right_visual)
        visual_delta = float(
            np.linalg.norm(right_base.translation - visual_pose.translation)
        )
    return DualArmMounts(left_base, right_base, visual_delta)


class DualArmAssembly:
    """Two independent RM75-B chains placed in a common world frame."""

    dof = 14

    def __init__(
        self,
        mounts: DualArmMounts,
        left: RM75Kinematics,
        right: RM75Kinematics,
    ) -> None:
        self.mounts = mounts
        self._kinematics = {"left": left, "right": right}

    @classmethod
    def from_source_urdf(
        cls,
        source_urdf: Optional[Path | str] = None,
        limits: Optional[JointLimits] = None,
    ) -> "DualArmAssembly":
        selected_limits = limits or physical_joint_limits()
        return cls(
            load_dual_arm_mounts(source_urdf),
            RM75Kinematics(limits=selected_limits),
            RM75Kinematics(limits=selected_limits),
        )

    def local_forward(
        self,
        arm: str,
        q_rad: np.ndarray,
        tool: Optional[pin.SE3] = None,
    ) -> pin.SE3:
        try:
            kinematics = self._kinematics[arm]
        except KeyError as exc:
            raise ValueError("arm must be 'left' or 'right'") from exc
        return kinematics.forward(q_rad, tool)

    def forward(
        self,
        arm: str,
        q_rad: np.ndarray,
        tool: Optional[pin.SE3] = None,
    ) -> pin.SE3:
        local = self.local_forward(arm, q_rad, tool)
        if arm == "left":
            return self.mounts.left_base * local
        if arm == "right":
            return self.mounts.right_base * local
        raise ValueError("arm must be 'left' or 'right'")
