"""Offline cooker contracts: no real CAD download or CoACD run required here."""

import copy
import hashlib
import importlib.util
import tempfile
import unittest
from pathlib import Path

import numpy as np

SCRIPT = Path(__file__).resolve().parents[1] / "generate_full_collisions.py"
spec = importlib.util.spec_from_file_location("collision_cooker", SCRIPT)
cooker = importlib.util.module_from_spec(spec)
spec.loader.exec_module(cooker)


def cube(center, size=0.004):
    return (
        np.array([[x, y, z] for x in (-size, size) for y in (-size, size) for z in (-size, size)])
        + center
    )


class CookingTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.source = Path(self.temp.name)
        (self.source / "meshes").mkdir()
        (self.source / "meshes/m.stl").write_bytes(b"source")
        links = "".join(
            f'<link name="{name}"><visual name="{name}_visual"><geometry>'
            '<mesh filename="meshes/m.stl" scale=".001 .001 .001"/>'
            "</geometry></visual></link>"
            for name in ("base", "motor", "arm")
        )
        self.xml = (
            '<robot name="fixture">'
            + links
            + """
          <joint name="fixed" type="fixed"><parent link="base"/><child link="motor"/></joint>
          <joint name="hinge" type="continuous">
            <parent link="motor"/><child link="arm"/><axis xyz="0 0 1"/>
          </joint>
        </robot>"""
        )
        self.config = {
            "rootLink": "base",
            "meshes": {"meshes/m.stl": hashlib.sha256(b"source").hexdigest()},
            "assemblyCores": {
                "hinge": {"radiusM": 0.015, "axialHalfExtentM": 0.02, "reason": "fixture bearing"}
            },
        }
        self.save_source()

    def save_source(self):
        data = self.xml.encode()
        (self.source / "LeKiwi.urdf").write_bytes(data)
        self.config["urdfSha256"] = hashlib.sha256(data).hexdigest()

    def test_covers_fixed_and_moving_visuals_and_preserves_scale(self):
        items = cooker.visuals(self.source, self.config)
        self.assertEqual(
            [x["visual"] for x in items], ["base_visual", "motor_visual", "arm_visual"]
        )
        self.assertTrue(all(x["scale"] == [0.001] * 3 for x in items))

    def test_rejects_changed_source_urdf_and_mesh(self):
        (self.source / "LeKiwi.urdf").write_text("modified")
        with self.assertRaisesRegex(ValueError, "URDF"):
            cooker.visuals(self.source, self.config)
        self.save_source()
        (self.source / "meshes/m.stl").write_bytes(b"modified")
        with self.assertRaisesRegex(ValueError, "SHA-256"):
            cooker.visuals(self.source, self.config)

    def test_rejects_unsupported_scale_and_path_escape(self):
        for scale in ("nan .001 .001", "-.001 .001 .001", ".001 .001", "inf 1 1"):
            original = self.xml
            self.xml = self.xml.replace(".001 .001 .001", scale)
            self.save_source()
            with self.assertRaisesRegex(ValueError, "缩放"):
                cooker.visuals(self.source, self.config)
            self.xml = original
        self.xml = self.xml.replace("meshes/m.stl", "../escape.stl")
        self.save_source()
        with self.assertRaisesRegex(ValueError, "未批准"):
            cooker.visuals(self.source, self.config)

    def test_rejects_missing_or_incomplete_coverage(self):
        self.config["meshes"]["meshes/extra.stl"] = "unknown"
        with self.assertRaisesRegex(ValueError, "覆盖"):
            cooker.visuals(self.source, self.config)
        self.config["rootLink"] = "not_present"
        self.config["meshes"] = {}
        with self.assertRaisesRegex(ValueError, "覆盖"):
            cooker.visuals(self.source, self.config)

    def test_radial_minimum_uses_polygon_edges_not_only_vertices(self):
        bounds = cooker.swept_bounds(cube([2.0, 0, 0], 0.5), np.array([0, 0, 1.0]))
        self.assertAlmostEqual(bounds[2], 1.5)
        self.assertGreater(np.linalg.norm(cube([2.0, 0, 0], 0.5)[:, :2], axis=1).min(), bounds[2])
        self.assertTrue(cooker.intervals_overlap(bounds, (-0.1, 0.1, 1.51, 1.52)))
        self.assertFalse(cooker.intervals_overlap(bounds, (1.0, 2.0, 1.51, 1.52)))
        for angle in np.linspace(-np.pi, np.pi, 37):
            rotation = np.array(
                [[np.cos(angle), -np.sin(angle), 0], [np.sin(angle), np.cos(angle), 0], [0, 0, 1]]
            )
            np.testing.assert_allclose(
                cooker.swept_bounds(cube([2.0, 0, 0], 0.5) @ rotation.T, np.array([0, 0, 1.0])),
                bounds,
                atol=1e-12,
            )

    def test_core_partition_preserves_volume_and_separates_mating_from_structure(self):
        from scipy.spatial import ConvexHull

        coverage = cooker.visuals(self.source, self.config)
        parts = [
            {
                "visual": item["visual"],
                "name": item["link"],
                "vertices": " ".join(map(str, cube([0.012, 0, 0], 0.01).ravel())),
            }
            for item in coverage
        ]
        for item in coverage:
            item["parts"] = 1
        divided = cooker.partition_assembly_cores(self.source, self.config, parts, coverage)
        policy = cooker.joint_policy(self.source, self.config, divided, coverage)[0]
        for visual in (item["visual"] for item in coverage):
            group = [piece for piece in divided if piece["visual"] == visual]
            volume = 0
            for piece in group:
                hull = ConvexHull(np.fromstring(piece["vertices"], sep=" ").reshape(-1, 3))
                volume += hull.volume
                # The old unsplit hull locked the bearing although this point is inside its core.
                if np.max(hull.equations[:, :3] @ [0.012, 0, 0] + hull.equations[:, 3]) < 1e-8:
                    self.assertIn(piece["name"], policy["coreParts"])
                if np.max(hull.equations[:, :3] @ [0.021, 0, 0] + hull.equations[:, 3]) < 1e-8:
                    self.assertNotIn(piece["name"], policy["coreParts"])
            self.assertAlmostEqual(volume / (0.02**3), 1, delta=0.001)
            self.assertGreater(len(group), 1)
        self.assertTrue(policy["pairs"])

    def test_exceptions_are_bounded_core_parts_not_whole_adjacent_bodies(self):
        coverage = cooker.visuals(self.source, self.config)
        parts = [
            {"name": name, "visual": visual, "vertices": " ".join(map(str, cube(center).ravel()))}
            for name, visual, center in [
                ("base_core", "base_visual", [0, 0, 0]),
                ("base_structural", "base_visual", [0.05, 0, 0]),
                ("far_axial", "base_visual", [0.05, 0, 0.2]),
                ("arm_core", "arm_visual", [0, 0, 0]),
                ("arm_structural", "arm_visual", [0.05, 0, 0]),
            ]
        ]
        policy = cooker.joint_policy(self.source, self.config, parts, coverage)[0]
        self.assertEqual(policy["coreParts"], ["base_core", "arm_core"])
        self.assertEqual(policy["pairs"], [["base_structural", "arm_structural"]])
        self.assertGreater(policy["sweptPrunedPairs"], 0)
        self.assertEqual([x["weldJoint"] for x in coverage], ["base", "base", "hinge"])
        bad = copy.deepcopy(self.config)
        bad["assemblyCores"]["hinge"]["radiusM"] = 1
        with self.assertRaisesRegex(ValueError, "尺寸预算"):
            cooker.joint_policy(self.source, bad, parts, coverage)
        with self.assertRaisesRegex(ValueError, "整对"):
            cooker.joint_policy(
                self.source, self.config, [p for p in parts if "core" in p["name"]], coverage
            )


if __name__ == "__main__":
    unittest.main()