import json
from typing import Any
from typing import Dict
from typing import List

import h5py
from pydantic import BaseModel

from ewoks3dxrd.nexus.utils import get_task_inputs
from ewoks3dxrd.nexus.utils import save_task_inputs


class MockSubConfig(BaseModel):
    name: str
    values: List[float]
    metadata: Dict[str, Any]


class MockTaskInputs(BaseModel):
    id: int
    config: MockSubConfig
    tags: List[str]
    overwrite: bool = False  # This should be excluded during save


def test_save_and_get_task_inputs(tmp_path):
    test_file = tmp_path / "test_data.nx"
    original_inputs = MockTaskInputs(
        id=123,
        config=MockSubConfig(
            name="Test Experiment",
            values=[0.1, 0.2, 0.5],
            metadata={"mode": "gaussian", "threshold": 5},
        ),
        tags=["raw", "segmented"],
        overwrite=True,
    )

    with h5py.File(test_file, "w") as h5_file:
        proc_group = h5_file.create_group("my_process")
        save_task_inputs(proc_group, original_inputs)

    with h5py.File(test_file, "r") as h5_file:
        json_data = h5_file["my_process/task_inputs/data"][()]
        parsed_json = json.loads(json_data)

        assert parsed_json["id"] == 123
        assert "overwrite" not in parsed_json  # Ensure exclusion worked
        assert parsed_json["config"]["metadata"]["mode"] == "gaussian"

    reconstructed = get_task_inputs(
        filename=test_file, process_group_path="my_process", model_class=MockTaskInputs
    )

    assert isinstance(reconstructed, MockTaskInputs)
    assert reconstructed.id == 123
    assert reconstructed.config.values == [0.1, 0.2, 0.5]
    assert reconstructed.tags == ["raw", "segmented"]
    assert reconstructed.overwrite is False
