from __future__ import annotations

import json
import posixpath
from pathlib import Path
from typing import Tuple
from typing import Type
from typing import TypeVar

import h5py
from silx.io.utils import DataUrl

from ewoks3dxrd.models import InputsWithOverwrite

T = TypeVar("T", bound=InputsWithOverwrite)

IMAGED11_DEFAULT_ATTR_UNITS = {
    "Number_of_pixels": "",
    "sum_intensity": "",
    "omega": "deg",
    "f_raw": "px",
    "s_raw": "px",
    "fc": "px",
    "sc": "px",
    "ds": "1/µm",
    "eta": "rad",
    "gx": "1/µm",
    "gy": "1/µm",
    "gz": "1/µm",
    "tth": "rad",
    "xl": "µm",
    "yl": "µm",
    "zl": "µm",
}


def group_exists(filename: str | Path, data_group_path: str) -> bool:
    if isinstance(filename, str):
        filename = Path(filename)
    if not filename.exists():
        return False

    with h5py.File(filename, "r") as file:
        if data_group_path in file:
            return True
    return False


def check_throw_write_data_group_url(
    overwrite: bool, filename: str | Path, data_group_path: str
):
    if not overwrite and group_exists(
        filename=filename, data_group_path=data_group_path
    ):
        raise ValueError(
            f"""Data group '{data_group_path}' already exists in {filename},
            Set `overwrite` to True if you wish to overwrite the existing data group.
            """
        )


def get_data_url_paths(data_url_as_str: str) -> Tuple[str, str]:
    data_url = DataUrl(data_url_as_str)
    file_path = data_url.file_path()
    data_path = data_url.data_path()
    if not Path(file_path).exists():
        raise FileNotFoundError(f"{file_path} not found")
    if data_path is None or not group_exists(file_path, data_path):
        raise KeyError(f"{data_path} not found in {file_path}")

    return file_path, data_path


def get_entry_name(process_group_path: str):
    return posixpath.dirname(process_group_path)


def save_task_inputs(
    group_path: h5py.Group,
    task_inputs: InputsWithOverwrite,
    group_name: str = "task_inputs",
):
    # Saves a Pydantic model as a JSON string in the task Nexus group.
    input_group = group_path.require_group(group_name)
    input_group.attrs["NX_class"] = "NXnote"
    input_group.attrs["module"] = task_inputs.__class__.__module__

    # We exclude 'overwrite'
    # mode='json' handles Path and Enum objects automatically
    serializable_dict = task_inputs.model_dump(exclude={"overwrite"}, mode="json")

    json_str = json.dumps(serializable_dict, indent=4)
    input_group.create_dataset("data", data=json_str)
    input_group.attrs["type"] = "application/json"


def get_task_inputs(
    filename: str | Path,
    process_group_path: str,
    model_class: Type[T],
    group_name: str = "task_inputs",
) -> T:
    # Retrieves and reconstructs the Pydantic model from the Nexus file.
    with h5py.File(filename, "r") as h5_file:
        dataset_path = f"{process_group_path}/{group_name}/data"
        if dataset_path not in h5_file:
            raise KeyError(f"Task inputs not found at {dataset_path}")

        json_str = h5_file[dataset_path][()]
        data_dict = json.loads(json_str)

        return model_class(**data_dict)
