import json
import tempfile
import unittest
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

from app.processing import JobContext
from app.processing.modules.gripper_device import GripperDeviceModule
from app.workflow_dispatch import _episode_input_groups
from app.workflow_types import migrate_graph_types


class GripperDeviceModuleTests(unittest.TestCase):
    def _context(self, root: Path) -> JobContext:
        return JobContext(
            node_id="gripper",
            node_type="gripper_device",
            config={"source_key": "gripper_rgb"},
            job={},
            input_root=root,
            output_root=root / "outputs",
            incoming={},
            progress=lambda _value: None,
            node_states={},
        )

    def test_publishes_one_composite_device_with_paired_force_output(self):
        with tempfile.TemporaryDirectory() as temp:
            root = Path(temp)
            (root / "meta").mkdir()
            (root / "videos" / "observation.images.gripper_rgb" / "chunk-000").mkdir(parents=True)
            (root / "videos" / "observation.images.gripper_rgb" / "chunk-000" / "episode_000000.mp4").touch()
            (root / "meta" / "info.json").write_text(json.dumps({
                "devices": [{
                    "key": "gripper:test",
                    "kind": "gripper",
                    "name": "UMI Gripper",
                    "slots": ["gripper_rgb", "gripper_force_left", "gripper_force_right"],
                }],
            }), encoding="utf-8")
            (root / "data" / "chunk-000").mkdir(parents=True)
            pq.write_table(pa.table({
                "observation.gripper_state": [[0.5, 1.0, 0.0]],
                "observation.slam_pose": [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]],
                "observation.gripper_left_force_matrix": [list(range(3))],
                "observation.gripper_right_force_matrix": [list(range(3))],
            }), root / "data" / "chunk-000" / "episode_000000.parquet")

            outputs = GripperDeviceModule().run(self._context(root))

            self.assertEqual(set(outputs), {
                "rgb_video", "gripper_state", "slam_trajectory",
                "tactile_force_matrices",
            })
            self.assertEqual(outputs["slam_trajectory"].metadata["names"],
                             ["x", "y", "z", "qx", "qy", "qz", "qw"])
            self.assertEqual(outputs["tactile_force_matrices"].metadata["columns"], [
                "observation.gripper_left_force_matrix",
                "observation.gripper_right_force_matrix",
            ])
            self.assertEqual(outputs["tactile_force_matrices"].metadata["shape"],
                             [2, 250, 250, 3])

    def test_dispatch_groups_gripper_stream_as_one_device(self):
        episode = {
            "camera_names": ["gripper_rgb", "gripper_stereo_left", "gripper_stereo_right"],
            "device_names": {
                "gripper_rgb": "UMI Gripper",
                "gripper_stereo_left": "UMI Gripper",
                "gripper_stereo_right": "UMI Gripper",
            },
            "devices": [],
            "sensors": [],
        }
        groups = _episode_input_groups(episode)
        self.assertEqual([group["input_type"] for group in groups], ["gripper_device"])
        self.assertEqual(groups[0]["source_keys"], [
            "gripper_rgb", "gripper_stereo_left", "gripper_stereo_right",
        ])

    def test_migrates_old_force_ports_to_one_connection(self):
        graph, changed = migrate_graph_types({
            "nodes": [{
                "id": "gripper", "data": {
                    "nodeType": "gripper_device",
                    "outputs": [
                        {"key": "tactile_left", "label": "Left Force Matrix"},
                        {"key": "tactile_right", "label": "Right Force Matrix"},
                    ],
                },
            }, {"id": "target", "data": {"nodeType": "annotation"}}],
            "edges": [
                {"id": "left", "source": "gripper", "sourceHandle": "tactile_left",
                 "target": "target", "targetHandle": "data"},
                {"id": "right", "source": "gripper", "sourceHandle": "tactile_right",
                 "target": "target", "targetHandle": "data"},
            ],
        })
        self.assertTrue(changed)
        self.assertEqual(graph["nodes"][0]["data"]["outputs"][-1]["key"],
                         "tactile_force_matrices")
        self.assertEqual(len(graph["edges"]), 1)
        self.assertEqual(graph["edges"][0]["sourceHandle"], "tactile_force_matrices")


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