"""Defines helper functions for MJCF accessors."""

import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, cast, get_args

from onshape.formats.common import ActuatorMetadata, JointMetadata

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from urdf2mjcf.convert import ConversionMetadata as ConversionMetadataRef


@dataclass
class CollisionParams:
    condim: int = field(default=6)
    contype: int = field(default=0)
    conaffinity: int = field(default=1)
    solref: list[float] = field(default_factory=lambda: [0.005, 1.0])
    solimp: list[float] = field(default_factory=lambda: [0.99, 0.999, 0.00001])
    friction: list[float] = field(default_factory=lambda: [0.8, 0.02, 0.01])


@dataclass
class SiteMetadata:
    name: str = field()
    body_name: str = field()
    site_type: str | None = field(default=None)
    pos: list[float] | None = field(default=None)  # (x, y, z)
    size: list[float] | None = field(default=None)  # (size_x, size_y, size_z)


@dataclass
class ImuSensor:
    body_name: str = field()
    pos: list[float] | None = field(default=None)  # (x, y, z)
    rpy: list[float] | None = field(default=None)  # (roll, pitch, yaw)
    acc_noise: float | None = field(default=None)
    gyro_noise: float | None = field(default=None)
    mag_noise: float | None = field(default=None)


@dataclass
class ForceSensor:
    body_name: str = field()
    site_name: str = field()
    name: str | None = field(default=None)
    noise: float | None = field(default=None)


@dataclass
class TouchSensor:
    body_name: str = field()
    site_name: str = field()
    name: str | None = field(default=None)
    noise: float | None = field(default=None)


@dataclass
class ExplicitFloorContacts:
    contact_links: list[str] = field(default_factory=lambda: [])
    class_name: str = field(default="collision")


@dataclass
class CollisionGeometry:
    name: str = field()
    collision_type: str = field()
    sphere_radius: float = field(default=0.01)
    axis_order: tuple[int, int, int] = field(default=(0, 1, 2))
    flip_axis: bool = field(default=False)
    offset_x: float = field(default=0.0)
    offset_y: float = field(default=0.0)
    offset_z: float = field(default=0.0)


@dataclass
class WeldConstraint:
    """Represents a weld constraint between two bodies.

    By default, body1 is set to "world" to create a fixed constraint in global space.
    """

    body1: str = field()
    body2: str = field()  # default_factory=lambda: "world")
    solimp: list[float] = field(default_factory=lambda: [0.9, 0.95, 0.005])  # clamped to 0.005 for higher stiffness
    solref: list[float] = field(
        default_factory=lambda: [0.005, 1.0]
    )  # critical damping 1.0, very quick time constant to return to static


@dataclass
class ConversionMetadata:
    suffix: str | None = field(default=None)
    freejoint: bool = field(default=True)
    collision_params: CollisionParams = field(default_factory=lambda: CollisionParams())
    sites: list[SiteMetadata] = field(default_factory=lambda: [])
    imus: list[ImuSensor] = field(default_factory=lambda: [])
    force_sensors: list[ForceSensor] = field(default_factory=lambda: [])
    touch_sensors: list[TouchSensor] = field(default_factory=lambda: [])
    collision_geometries: list[CollisionGeometry] = field(default_factory=lambda: [])
    explicit_contacts: ExplicitFloorContacts | None = field(default_factory=ExplicitFloorContacts)
    weld_constraints: list[WeldConstraint] = field(default_factory=lambda: [])
    remove_redundancies: bool = field(default=True)
    floating_base: bool = field(default=True)
    maxhullvert: int = field(default=64)
    angle: str = field(default="radian")
    floor_name: str = field(default="floor")
    add_floor: bool = field(default=False)
    backlash: float | None = field(default=None)
    backlash_damping: float = field(default=0.01)
    height_offset: float = field(default=0.0)
    visualize_collision_meshes: bool = field(default=True)
    joint_name_to_metadata: dict[str, JointMetadata] | None = field(default=None)
    actuator_type_to_metadata: dict[str, ActuatorMetadata] | None = field(default=None)


def convert_to_mjcf_metadata(metadata: ConversionMetadata) -> "ConversionMetadataRef":
    try:
        import urdf2mjcf

        logger.debug("urdf2mjcf version: %s", urdf2mjcf.__version__)

    except ImportError as e:
        raise ImportError(
            "urdf2mjcf is required to run this script. Install it with `pip install "
            "'onshape[mujoco]'` to install the required dependencies."
        ) from e

    from urdf2mjcf.model import (
        ActuatorMetadata as ActuatorMetadataRef,
        Angle,
        CollisionGeometry as CollisionGeometryRef,
        CollisionParams as CollisionParamsRef,
        CollisionType,
        ConversionMetadata as ConversionMetadataRef,
        ExplicitFloorContacts as ExplicitFloorContactsRef,
        ForceSensor as ForceSensorRef,
        ImuSensor as ImuSensorRef,
        JointMetadata as JointMetadataRef,
        SiteMetadata as SiteMetadataRef,
        TouchSensor as TouchSensorRef,
        WeldConstraint as WeldConstraintRef,
    )

    if metadata.angle not in get_args(Angle):
        raise ValueError(f"Invalid angle type: {metadata.angle}. Must be one of {get_args(Angle)}")

    for cg in metadata.collision_geometries:
        if not hasattr(CollisionType, cg.collision_type.upper()):
            raise ValueError(f"Bad collision type: {cg.collision_type}. Must be in {CollisionType.__members__.keys()}")

    joint_name_to_metadata = (
        {name: JointMetadataRef.from_dict(vars(param)) for name, param in metadata.joint_name_to_metadata.items()}
        if metadata.joint_name_to_metadata
        else None
    )

    actuator_type_to_metadata = (
        {typ: ActuatorMetadataRef.from_dict(vars(param)) for typ, param in metadata.actuator_type_to_metadata.items()}
        if metadata.actuator_type_to_metadata
        else None
    )

    return ConversionMetadataRef(
        freejoint=metadata.freejoint,
        collision_params=CollisionParamsRef(
            condim=metadata.collision_params.condim,
            contype=metadata.collision_params.contype,
            solref=metadata.collision_params.solref,
            friction=metadata.collision_params.friction,
        ),
        sites=[
            SiteMetadataRef(
                name=site.name,
                body_name=site.body_name,
                site_type=cast(urdf2mjcf.model.SiteType | None, site.site_type),
                pos=site.pos,
                size=site.size,
            )
            for site in metadata.sites
        ],
        imus=[
            ImuSensorRef(
                body_name=imu.body_name,
                pos=imu.pos,
                rpy=imu.rpy,
                acc_noise=imu.acc_noise,
                gyro_noise=imu.gyro_noise,
                mag_noise=imu.mag_noise,
            )
            for imu in metadata.imus
        ],
        force_sensors=[
            ForceSensorRef(
                body_name=fs.body_name,
                site_name=fs.site_name,
                name=fs.name,
                noise=fs.noise,
            )
            for fs in metadata.force_sensors
        ],
        touch_sensors=[
            TouchSensorRef(
                body_name=ts.body_name,
                site_name=ts.site_name,
                name=ts.name,
                noise=ts.noise,
            )
            for ts in metadata.touch_sensors
        ],
        collision_geometries=[
            CollisionGeometryRef(
                name=cg.name,
                collision_type=getattr(CollisionType, cg.collision_type.upper()),
                sphere_radius=cg.sphere_radius,
                axis_order=cg.axis_order,
                flip_axis=cg.flip_axis,
                offset_x=cg.offset_x,
                offset_y=cg.offset_y,
                offset_z=cg.offset_z,
            )
            for cg in metadata.collision_geometries
        ],
        explicit_contacts=(
            None
            if metadata.explicit_contacts is None
            else ExplicitFloorContactsRef(
                contact_links=metadata.explicit_contacts.contact_links,
                class_name=metadata.explicit_contacts.class_name,
            )
        ),
        weld_constraints=[
            WeldConstraintRef(
                body1=weld.body1,
                body2=weld.body2,
                solimp=weld.solimp,
                solref=weld.solref,
            )
            for weld in metadata.weld_constraints
        ],
        remove_redundancies=metadata.remove_redundancies,
        floating_base=metadata.floating_base,
        maxhullvert=metadata.maxhullvert,
        angle=cast(Angle, metadata.angle),
        floor_name=metadata.floor_name,
        add_floor=metadata.add_floor,
        backlash=metadata.backlash,
        backlash_damping=metadata.backlash_damping,
        height_offset=metadata.height_offset,
        visualize_collision_meshes=metadata.visualize_collision_meshes,
        joint_name_to_metadata=joint_name_to_metadata,
        actuator_type_to_metadata=actuator_type_to_metadata,
    )
