"""Tests for solver.mates.primitives -- mate type definitions."""

from __future__ import annotations

from typing import ClassVar

import numpy as np
import pytest

from solver.mates.primitives import (
    GeometryRef,
    GeometryType,
    Mate,
    MateType,
    dof_removed,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_ref(
    body_id: int,
    geom_type: GeometryType,
    *,
    geometry_id: str = "Geom001",
    origin: np.ndarray | None = None,
    direction: np.ndarray | None = None,
) -> GeometryRef:
    """Factory for GeometryRef with sensible defaults."""
    if origin is None:
        origin = np.zeros(3)
    if direction is None and geom_type in {
        GeometryType.FACE,
        GeometryType.AXIS,
        GeometryType.PLANE,
    }:
        direction = np.array([0.0, 0.0, 1.0])
    return GeometryRef(
        body_id=body_id,
        geometry_type=geom_type,
        geometry_id=geometry_id,
        origin=origin,
        direction=direction,
    )


# ---------------------------------------------------------------------------
# MateType
# ---------------------------------------------------------------------------


class TestMateType:
    """MateType enum construction and DOF values."""

    EXPECTED_DOF: ClassVar[dict[str, int]] = {
        "COINCIDENT": 3,
        "CONCENTRIC": 2,
        "PARALLEL": 2,
        "PERPENDICULAR": 1,
        "TANGENT": 1,
        "DISTANCE": 1,
        "ANGLE": 1,
        "LOCK": 6,
    }

    def test_member_count(self) -> None:
        assert len(MateType) == 8

    @pytest.mark.parametrize("name,dof", EXPECTED_DOF.items())
    def test_default_dof_values(self, name: str, dof: int) -> None:
        assert MateType[name].default_dof == dof

    def test_value_is_tuple(self) -> None:
        assert MateType.COINCIDENT.value == (0, 3)
        assert MateType.COINCIDENT.default_dof == 3

    def test_access_by_name(self) -> None:
        assert MateType["LOCK"] is MateType.LOCK

    def test_no_alias_collision(self) -> None:
        ordinals = [m.value[0] for m in MateType]
        assert len(ordinals) == len(set(ordinals))


# ---------------------------------------------------------------------------
# GeometryType
# ---------------------------------------------------------------------------


class TestGeometryType:
    """GeometryType enum."""

    def test_member_count(self) -> None:
        assert len(GeometryType) == 5

    def test_string_values(self) -> None:
        for gt in GeometryType:
            assert isinstance(gt.value, str)
            assert gt.value == gt.name.lower()

    def test_access_by_name(self) -> None:
        assert GeometryType["FACE"] is GeometryType.FACE


# ---------------------------------------------------------------------------
# GeometryRef
# ---------------------------------------------------------------------------


class TestGeometryRef:
    """GeometryRef dataclass."""

    def test_construction(self) -> None:
        ref = _make_ref(0, GeometryType.AXIS, geometry_id="Axis001")
        assert ref.body_id == 0
        assert ref.geometry_type is GeometryType.AXIS
        assert ref.geometry_id == "Axis001"
        np.testing.assert_array_equal(ref.origin, np.zeros(3))
        assert ref.direction is not None

    def test_default_direction_none(self) -> None:
        ref = GeometryRef(
            body_id=0,
            geometry_type=GeometryType.POINT,
            geometry_id="Point001",
        )
        assert ref.direction is None

    def test_to_dict_round_trip(self) -> None:
        ref = _make_ref(
            1,
            GeometryType.FACE,
            origin=np.array([1.0, 2.0, 3.0]),
            direction=np.array([0.0, 1.0, 0.0]),
        )
        d = ref.to_dict()
        restored = GeometryRef.from_dict(d)
        assert restored.body_id == ref.body_id
        assert restored.geometry_type is ref.geometry_type
        assert restored.geometry_id == ref.geometry_id
        np.testing.assert_array_almost_equal(restored.origin, ref.origin)
        assert restored.direction is not None
        np.testing.assert_array_almost_equal(restored.direction, ref.direction)

    def test_to_dict_with_none_direction(self) -> None:
        ref = GeometryRef(
            body_id=2,
            geometry_type=GeometryType.POINT,
            geometry_id="Point002",
            origin=np.array([5.0, 6.0, 7.0]),
        )
        d = ref.to_dict()
        assert d["direction"] is None
        restored = GeometryRef.from_dict(d)
        assert restored.direction is None


# ---------------------------------------------------------------------------
# Mate
# ---------------------------------------------------------------------------


class TestMate:
    """Mate dataclass."""

    def test_construction(self) -> None:
        ref_a = _make_ref(0, GeometryType.FACE)
        ref_b = _make_ref(1, GeometryType.FACE)
        m = Mate(mate_id=0, mate_type=MateType.COINCIDENT, ref_a=ref_a, ref_b=ref_b)
        assert m.mate_id == 0
        assert m.mate_type is MateType.COINCIDENT

    def test_value_default_zero(self) -> None:
        ref_a = _make_ref(0, GeometryType.FACE)
        ref_b = _make_ref(1, GeometryType.FACE)
        m = Mate(mate_id=0, mate_type=MateType.COINCIDENT, ref_a=ref_a, ref_b=ref_b)
        assert m.value == 0.0

    def test_tolerance_default(self) -> None:
        ref_a = _make_ref(0, GeometryType.FACE)
        ref_b = _make_ref(1, GeometryType.FACE)
        m = Mate(mate_id=0, mate_type=MateType.COINCIDENT, ref_a=ref_a, ref_b=ref_b)
        assert m.tolerance == 1e-6

    def test_to_dict_round_trip(self) -> None:
        ref_a = _make_ref(0, GeometryType.AXIS, origin=np.array([1.0, 0.0, 0.0]))
        ref_b = _make_ref(1, GeometryType.AXIS, origin=np.array([2.0, 0.0, 0.0]))
        m = Mate(
            mate_id=5,
            mate_type=MateType.CONCENTRIC,
            ref_a=ref_a,
            ref_b=ref_b,
            value=0.0,
            tolerance=1e-8,
        )
        d = m.to_dict()
        restored = Mate.from_dict(d)
        assert restored.mate_id == m.mate_id
        assert restored.mate_type is m.mate_type
        assert restored.ref_a.body_id == m.ref_a.body_id
        assert restored.ref_b.body_id == m.ref_b.body_id
        assert restored.value == m.value
        assert restored.tolerance == m.tolerance

    def test_from_dict_missing_optional(self) -> None:
        d = {
            "mate_id": 1,
            "mate_type": "DISTANCE",
            "ref_a": _make_ref(0, GeometryType.POINT).to_dict(),
            "ref_b": _make_ref(1, GeometryType.POINT).to_dict(),
        }
        m = Mate.from_dict(d)
        assert m.value == 0.0
        assert m.tolerance == 1e-6


# ---------------------------------------------------------------------------
# dof_removed
# ---------------------------------------------------------------------------


class TestDofRemoved:
    """Context-dependent DOF removal counts."""

    def test_coincident_face_face(self) -> None:
        ref_a = _make_ref(0, GeometryType.FACE)
        ref_b = _make_ref(1, GeometryType.FACE)
        assert dof_removed(MateType.COINCIDENT, ref_a, ref_b) == 3

    def test_coincident_point_point(self) -> None:
        ref_a = _make_ref(0, GeometryType.POINT)
        ref_b = _make_ref(1, GeometryType.POINT)
        assert dof_removed(MateType.COINCIDENT, ref_a, ref_b) == 3

    def test_coincident_edge_edge(self) -> None:
        ref_a = _make_ref(0, GeometryType.EDGE)
        ref_b = _make_ref(1, GeometryType.EDGE)
        assert dof_removed(MateType.COINCIDENT, ref_a, ref_b) == 2

    def test_coincident_face_point(self) -> None:
        ref_a = _make_ref(0, GeometryType.FACE)
        ref_b = _make_ref(1, GeometryType.POINT)
        assert dof_removed(MateType.COINCIDENT, ref_a, ref_b) == 1

    def test_concentric_axis_axis(self) -> None:
        ref_a = _make_ref(0, GeometryType.AXIS)
        ref_b = _make_ref(1, GeometryType.AXIS)
        assert dof_removed(MateType.CONCENTRIC, ref_a, ref_b) == 2

    def test_lock_any(self) -> None:
        ref_a = _make_ref(0, GeometryType.FACE)
        ref_b = _make_ref(1, GeometryType.POINT)
        assert dof_removed(MateType.LOCK, ref_a, ref_b) == 6

    def test_distance_any(self) -> None:
        ref_a = _make_ref(0, GeometryType.POINT)
        ref_b = _make_ref(1, GeometryType.EDGE)
        assert dof_removed(MateType.DISTANCE, ref_a, ref_b) == 1

    def test_unknown_combo_uses_default(self) -> None:
        """Unlisted geometry combos fall back to default_dof."""
        ref_a = _make_ref(0, GeometryType.EDGE)
        ref_b = _make_ref(1, GeometryType.POINT)
        result = dof_removed(MateType.COINCIDENT, ref_a, ref_b)
        assert result == MateType.COINCIDENT.default_dof


# ---------------------------------------------------------------------------
# Mate.validate
# ---------------------------------------------------------------------------


class TestMateValidation:
    """Mate.validate() compatibility checks."""

    def test_valid_concentric(self) -> None:
        ref_a = _make_ref(0, GeometryType.AXIS)
        ref_b = _make_ref(1, GeometryType.AXIS)
        m = Mate(mate_id=0, mate_type=MateType.CONCENTRIC, ref_a=ref_a, ref_b=ref_b)
        m.validate()  # should not raise

    def test_invalid_concentric_face(self) -> None:
        ref_a = _make_ref(0, GeometryType.FACE)
        ref_b = _make_ref(1, GeometryType.AXIS)
        m = Mate(mate_id=0, mate_type=MateType.CONCENTRIC, ref_a=ref_a, ref_b=ref_b)
        with pytest.raises(ValueError, match="CONCENTRIC"):
            m.validate()

    def test_valid_coincident_face_face(self) -> None:
        ref_a = _make_ref(0, GeometryType.FACE)
        ref_b = _make_ref(1, GeometryType.FACE)
        m = Mate(mate_id=0, mate_type=MateType.COINCIDENT, ref_a=ref_a, ref_b=ref_b)
        m.validate()  # should not raise

    def test_invalid_self_mate(self) -> None:
        ref_a = _make_ref(0, GeometryType.FACE)
        ref_b = _make_ref(0, GeometryType.FACE, geometry_id="Face002")
        m = Mate(mate_id=0, mate_type=MateType.COINCIDENT, ref_a=ref_a, ref_b=ref_b)
        with pytest.raises(ValueError, match="Self-mate"):
            m.validate()

    def test_invalid_parallel_point(self) -> None:
        ref_a = _make_ref(0, GeometryType.POINT)
        ref_b = _make_ref(1, GeometryType.AXIS)
        m = Mate(mate_id=0, mate_type=MateType.PARALLEL, ref_a=ref_a, ref_b=ref_b)
        with pytest.raises(ValueError, match="PARALLEL"):
            m.validate()

    def test_invalid_tangent_axis(self) -> None:
        ref_a = _make_ref(0, GeometryType.AXIS)
        ref_b = _make_ref(1, GeometryType.FACE)
        m = Mate(mate_id=0, mate_type=MateType.TANGENT, ref_a=ref_a, ref_b=ref_b)
        with pytest.raises(ValueError, match="TANGENT"):
            m.validate()

    def test_missing_direction_for_axis(self) -> None:
        ref_a = GeometryRef(
            body_id=0,
            geometry_type=GeometryType.AXIS,
            geometry_id="Axis001",
            origin=np.zeros(3),
            direction=None,  # missing!
        )
        ref_b = _make_ref(1, GeometryType.AXIS)
        m = Mate(mate_id=0, mate_type=MateType.CONCENTRIC, ref_a=ref_a, ref_b=ref_b)
        with pytest.raises(ValueError, match="direction"):
            m.validate()
