from __future__ import annotations

import importlib.util
import tempfile
import unittest
from pathlib import Path
import xml.etree.ElementTree as ET
import zipfile

from build123d import Box, export_step


ROOT = Path(__file__).resolve().parents[1]


def load_module(name: str, path: Path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader
    spec.loader.exec_module(module)
    return module


SURFACEIR = load_module(
    "surfaceir_pipeline_robot_export_test",
    ROOT / "scripts" / "surfaceir_pipeline.py",
)
ROBOT_EXPORT = load_module(
    "robot_description_export_test",
    ROOT / "scripts" / "robot_description_export.py",
)


class RobotDescriptionExportTests(unittest.TestCase):
    def test_portable_urdf_and_mjcf_packages_rebuild_from_designir(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            teacher = root / "teacher.step"
            designir = root / "box.designir.json"
            export_step(Box(10, 20, 30), teacher)
            SURFACEIR.write_json(
                designir,
                SURFACEIR.extract_surfaceir(teacher),
            )
            teacher.unlink()

            for output_format, description_suffix, xml_root in (
                ("urdf", ".urdf", "robot"),
                ("mjcf", ".mjcf.xml", "mujoco"),
            ):
                result = ROBOT_EXPORT.export_robot_package(
                    designir,
                    root / output_format,
                    output_format,
                )
                package = Path(result["package"])
                self.assertTrue(package.is_file())
                with zipfile.ZipFile(package) as archive:
                    names = archive.namelist()
                    description = next(
                        name
                        for name in names
                        if name.endswith(description_suffix)
                    )
                    self.assertTrue(
                        any(name.startswith("meshes/") and name.endswith(".stl") for name in names)
                    )
                    self.assertIn("robot-export.json", names)
                    root_element = ET.fromstring(archive.read(description))
                    self.assertEqual(xml_root, root_element.tag)
                self.assertEqual(
                    "static_single_link",
                    result["robot_structure"],
                )


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