from __future__ import annotations

import importlib.util
import json
import tempfile
import unittest
from pathlib import Path

WORKER_PATH = (
    Path(__file__).resolve().parents[1] / "windows-node" / "origin-worker" / "worker.py"
)
SPEC = importlib.util.spec_from_file_location("zcbot_origin_worker", WORKER_PATH)
assert SPEC and SPEC.loader
worker = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(worker)


class OriginWorkerUnitTests(unittest.TestCase):
    def test_csv_and_json_inputs_are_read_without_origin(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            csv_path = root / "input.csv"
            csv_path.write_text("x,y\n1,2\n3,4\n", encoding="utf-8")
            self.assertEqual(worker._read_rows(csv_path, None), (["x", "y"], [["1", "2"], ["3", "4"]]))

            json_path = root / "input.json"
            json_path.write_text(json.dumps([{"x": 1, "y": 2}, {"x": 3, "y": 4}]), encoding="utf-8")
            self.assertEqual(worker._read_rows(json_path, None), (["x", "y"], [[1, 2], [3, 4]]))

    def test_manifest_uses_stable_id_and_streaming_digest(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "plot-spec.json"
            path.write_text("{}", encoding="utf-8")
            manifest = worker._manifest(path, "application/json")
            self.assertEqual(manifest["artifact_id"], "plot_spec")
            self.assertEqual(manifest["sha256"], worker._file_sha256(path))
            self.assertEqual(manifest["size_bytes"], 2)

    def test_axis_title_includes_units(self) -> None:
        self.assertEqual(worker._axis_title({"title": "Stress", "unit": "MPa"}, "Y"), "Stress (MPa)")
        self.assertEqual(worker._axis_title(None, "Time"), "Time")

    def test_keyed_input_directory_requires_exactly_one_file(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            keyed = root / "input" / "sample"
            keyed.mkdir(parents=True)
            path = keyed / "data.csv"
            path.write_text("x,y\n1,2\n", encoding="utf-8")
            self.assertEqual(worker._input_file(root, "sample"), path)
            (keyed / "extra.csv").write_text("x,y\n3,4\n", encoding="utf-8")
            with self.assertRaisesRegex(ValueError, "INPUT_FILE_COUNT_INVALID:sample"):
                worker._input_file(root, "sample")

    def test_series_are_resolved_against_each_input(self) -> None:
        input_data = {
            "first": (["time", "strength"], [[1, 10]]),
            "second": (["temperature", "value"], [[20, 30]]),
        }
        series = [
            {"input": "first", "x": "time", "y": "strength", "label": "7 d"},
            {"input": "second", "x": "temperature", "y": "value", "label": "28 d"},
        ]
        resolved, labels = worker._resolve_series(input_data, series)
        self.assertEqual(resolved, [
            {"input": "first", "label": "7 d", "x": 0, "y": 1},
            {"input": "second", "label": "28 d", "x": 0, "y": 1},
        ])
        self.assertEqual(labels, {("first", 1): "7 d", ("second", 1): "28 d"})

    def test_shared_y_column_rejects_conflicting_labels(self) -> None:
        input_data = {"sample": (["x1", "x2", "y"], [[1, 2, 3]])}
        with self.assertRaisesRegex(ValueError, "SERIES_LABEL_CONFLICT"):
            worker._resolve_series(input_data, [
                {"input": "sample", "x": "x1", "y": "y", "label": "First"},
                {"input": "sample", "x": "x2", "y": "y", "label": "Second"},
            ])

    def test_series_support_xyz_and_y_error_roles(self) -> None:
        resolved, labels = worker._resolve_series(
            {"sample": (["x", "y", "z", "sd"], [[0, 1, 2, 0.1]])},
            [{
                "input": "sample", "x": "x", "y": "y", "z": "z",
                "y_error": "sd", "label": "测量值",
            }],
        )
        self.assertEqual(resolved, [{
            "input": "sample", "label": "测量值", "x": 0, "y": 1,
            "z": 2, "y_error": 3,
        }])
        self.assertEqual(labels, {("sample", 1): "测量值"})

    def test_heatmap_matrix_accepts_complete_unordered_grid(self) -> None:
        matrix, xy_map = worker._heatmap_matrix(
            [[1, 20, 4], [0, 10, 1], [1, 10, 2], [0, 20, 3]],
            {"x": 0, "y": 1, "z": 2},
        )
        self.assertEqual(matrix, [[1.0, 2.0], [3.0, 4.0]])
        self.assertEqual(xy_map, (0.0, 1.0, 10.0, 20.0))

    def test_heatmap_matrix_rejects_invalid_grid(self) -> None:
        with self.assertRaisesRegex(ValueError, "HEATMAP_GRID_INCOMPLETE"):
            worker._heatmap_matrix(
                [[0, 10, 1], [1, 10, 2], [0, 20, 3]],
                {"x": 0, "y": 1, "z": 2},
            )
        with self.assertRaisesRegex(ValueError, "HEATMAP_COORDINATES_DUPLICATED"):
            worker._heatmap_matrix(
                [[0, 10, 1], [0, 10, 2]], {"x": 0, "y": 1, "z": 2}
            )
        with self.assertRaisesRegex(ValueError, "HEATMAP_GRID_NOT_REGULAR"):
            worker._heatmap_matrix(
                [
                    [0, 10, 1], [1, 10, 2], [3, 10, 3],
                    [0, 20, 4], [1, 20, 5], [3, 20, 6],
                ],
                {"x": 0, "y": 1, "z": 2},
            )


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