from __future__ import annotations

import importlib.util
import json
import struct
from pathlib import Path


def _load_export_module():
    candidates = [
        Path(__file__).resolve().parents[2] / "render-worker" / "scripts" / "export_step_to_gltf.py",
        Path("/compose/render-worker/scripts/export_step_to_gltf.py"),
    ]
    module_path = next((path for path in candidates if path.exists()), None)
    assert module_path is not None
    spec = importlib.util.spec_from_file_location("test_export_step_to_gltf", module_path)
    assert spec is not None
    assert spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _write_minimal_glb(path: Path, payload: dict) -> None:
    json_bytes = json.dumps(payload, separators=(",", ":")).encode()
    pad = (4 - len(json_bytes) % 4) % 4
    json_bytes += b" " * pad
    chunk = struct.pack("<II", len(json_bytes), 0x4E4F534A) + json_bytes
    header = struct.pack("<III", 0x46546C67, 2, 12 + len(chunk))
    path.write_bytes(header + chunk)


def _read_glb_json(path: Path) -> dict:
    data = path.read_bytes()
    json_len = struct.unpack_from("<I", data, 12)[0]
    return json.loads(data[20 : 20 + json_len])


def test_atomic_export_helpers_publish_temp_glb_over_existing_output(tmp_path: Path):
    module = _load_export_module()
    output_path = tmp_path / "existing.glb"
    output_path.write_bytes(b"old")

    temp_path = module._prepare_atomic_export_path(output_path)
    assert temp_path != output_path
    assert temp_path.parent == output_path.parent
    assert not temp_path.exists()

    temp_path.write_bytes(b"new")
    module._finalize_atomic_export(temp_path, output_path)

    assert output_path.read_bytes() == b"new"
    assert not temp_path.exists()


def test_inject_glb_extras_preserves_exact_leaf_mesh_part_keys_when_available(tmp_path: Path):
    module = _load_export_module()
    glb_path = tmp_path / "instance.glb"
    payload = {
        "asset": {"version": "2.0"},
        "scene": 0,
        "scenes": [{"nodes": [0]}],
        "nodes": [
            {"name": "Assembly", "children": [1, 2, 3]},
            {
                "name": "KERO_Z-575693-QP-DRH_ISB_1_AF21",
                "translation": [0.1, 0.2, 0.3],
                "rotation": [0.0, 0.0, 0.0, 1.0],
            },
            {
                "name": "KERO_Z-575693-QP-DRH_ISB_1_1",
                "mesh": 0,
                "translation": [0.1, 0.2, 0.3],
                "rotation": [0.0, 0.0, 0.0, 1.0],
            },
            {
                "name": "KERO_Z-575693-QP-DRH_ISB_1_AF6_",
                "translation": [0.4, 0.5, 0.6],
                "rotation": [0.0, 0.0, 0.0, 1.0],
            },
        ],
        "meshes": [{"primitives": []}],
    }
    _write_minimal_glb(glb_path, payload)

    module._inject_glb_extras(
        glb_path,
        {"partKeyMap": {}},
        part_key_map={
            "KERO_Z-575693-QP-DRH_ISB_1": "kero_z_575693_qp_drh_isb_1",
            "KERO_Z-575693-QP-DRH_ISB_1_1": "kero_z_575693_qp_drh_isb_1_1",
            "KERO_Z-575693-QP-DRH_ISB_1_AF6_": "kero_z_575693_qp_drh_isb_1_af6",
        },
        part_key_occurrences={
            "KERO_Z-575693-QP-DRH_ISB_1_1": ["kero_z_575693_qp_drh_isb_1_1"],
        },
    )

    result = _read_glb_json(glb_path)
    assert result["nodes"][1]["extras"]["partKey"] == "kero_z_575693_qp_drh_isb_1"
    assert result["nodes"][2]["extras"]["partKey"] == "kero_z_575693_qp_drh_isb_1_1"
    assert result["nodes"][3]["extras"]["partKey"] == "kero_z_575693_qp_drh_isb_1"


def test_inject_glb_extras_keeps_unique_leaf_mesh_part_keys_without_semantic_siblings(tmp_path: Path):
    module = _load_export_module()
    glb_path = tmp_path / "leaf.glb"
    payload = {
        "asset": {"version": "2.0"},
        "scene": 0,
        "scenes": [{"nodes": [0]}],
        "nodes": [
            {"name": "Assembly", "children": [1]},
            {
                "name": "UNIQUE_PART_1_1",
                "mesh": 0,
                "translation": [0.0, 0.0, 0.0],
                "rotation": [0.0, 0.0, 0.0, 1.0],
            },
        ],
        "meshes": [{"primitives": []}],
    }
    _write_minimal_glb(glb_path, payload)

    module._inject_glb_extras(
        glb_path,
        {"partKeyMap": {}},
        part_key_map={
            "UNIQUE_PART_1_1": "unique_part_1_1",
        },
    )

    result = _read_glb_json(glb_path)
    assert result["nodes"][1]["extras"]["partKey"] == "unique_part_1_1"


def test_inject_glb_extras_falls_back_to_semantic_siblings_when_exact_mesh_key_is_missing_even_if_instance_transforms_differ(tmp_path: Path):
    module = _load_export_module()
    glb_path = tmp_path / "instance-mismatch.glb"
    payload = {
        "asset": {"version": "2.0"},
        "scene": 0,
        "scenes": [{"nodes": [0]}],
        "nodes": [
            {"name": "Assembly", "children": [1, 2, 3]},
            {
                "name": "KERO_Z-575693-QP-DRH_ISB_1_AF21",
                "translation": [0.1, 0.2, 0.3],
                "rotation": [0.0, 0.0, 0.0, 1.0],
            },
            {
                "name": "KERO_Z-575693-QP-DRH_ISB_1_AF22",
                "translation": [-0.1, -0.2, 0.3],
                "rotation": [0.0, 0.0, 0.0, 1.0],
            },
            {
                "name": "KERO_Z-575693-QP-DRH_ISB_1_1",
                "mesh": 0,
                "translation": [0.9, 0.8, 0.3],
                "rotation": [0.0, 0.0, 0.70710678, 0.70710678],
            },
        ],
        "meshes": [{"primitives": []}],
    }
    _write_minimal_glb(glb_path, payload)

    module._inject_glb_extras(
        glb_path,
        {"partKeyMap": {}},
        part_key_map={
            "KERO_Z-575693-QP-DRH_ISB_1": "kero_z_575693_qp_drh_isb_1",
            "KERO_Z-575693-QP-DRH_ISB_1_AF21": "kero_z_575693_qp_drh_isb_1_af21",
            "KERO_Z-575693-QP-DRH_ISB_1_AF22": "kero_z_575693_qp_drh_isb_1_af22",
        },
    )

    result = _read_glb_json(glb_path)
    assert result["nodes"][3]["extras"]["partKey"] == "kero_z_575693_qp_drh_isb_1"


def test_inject_glb_extras_assigns_distinct_occurrence_keys_to_repeated_leaf_meshes(tmp_path: Path):
    module = _load_export_module()
    glb_path = tmp_path / "repeated.glb"
    payload = {
        "asset": {"version": "2.0"},
        "scene": 0,
        "scenes": [{"nodes": [0]}],
        "nodes": [
            {"name": "Assembly", "children": [1, 2, 3]},
            {"name": "KERO_Z-575693-QP-DRH_ISB_1_1", "mesh": 0},
            {"name": "KERO_Z-575693-QP-DRH_ISB_1_1", "mesh": 1},
            {"name": "KERO_Z-575693-QP-DRH_ISB_1_1", "mesh": 2},
        ],
        "meshes": [
            {"primitives": []},
            {"primitives": []},
            {"primitives": []},
        ],
    }
    _write_minimal_glb(glb_path, payload)

    module._inject_glb_extras(
        glb_path,
        {"partKeyMap": {}},
        part_key_map={
            "KERO_Z-575693-QP-DRH_ISB_1_1": "kero_z_575693_qp_drh_isb_1_1",
        },
        part_key_occurrences={
            "KERO_Z-575693-QP-DRH_ISB_1_1": [
                "kero_z_575693_qp_drh_isb_1_1",
                "kero_z_575693_qp_drh_isb_1_1_2",
                "kero_z_575693_qp_drh_isb_1_1_3",
            ],
        },
    )

    result = _read_glb_json(glb_path)
    assert result["nodes"][1]["extras"]["partKey"] == "kero_z_575693_qp_drh_isb_1_1"
    assert result["nodes"][2]["extras"]["partKey"] == "kero_z_575693_qp_drh_isb_1_1_2"
    assert result["nodes"][3]["extras"]["partKey"] == "kero_z_575693_qp_drh_isb_1_1_3"