from __future__ import annotations

import logging
import os

import fabio
import h5py
import pint
from silx.io.url import DataUrl
from silx.io.utils import h5py_read_dataset
from silx.io.utils import open as hdf5_open

from tomoscan.framereducer.reducedframesinfos import ReducedFramesInfos

_logger = logging.getLogger(__name__)

_ureg = pint.get_application_registry()


def format_data_path(data_path):
    """
    Format the data path by removing the "{index}" or "{index_zfill4}" pattern if it exists and return the formatted data path and a boolean indicating if the "{index_zfill4}" pattern was found.
    """
    index_zfill4_pattern = False
    if data_path.endswith("{index_zfill4}"):
        index_zfill4_pattern = True
        data_path = data_path[: -len("{index_zfill4}")]
    if data_path.endswith("{index}"):
        data_path = data_path[: -len("{index}")]
    if data_path.endswith("/"):
        data_path = data_path[:-1]
    return data_path, index_zfill4_pattern


def get_unit(attrs) -> pint.Unit | None:
    """
    Get the unit from the attributes of a dataset or group. The unit can be specified with the "unit" or "units" attribute. If the unit is "kev", it will be converted to keV. If the unit is not specified, None will be returned.
    """
    if "unit" in attrs:
        unit = attrs["unit"]
    elif "units":
        unit = attrs["units"]
    if unit == "kev":
        unit = _ureg.keV
    if unit is not None:
        return pint.Unit(unit)

    return None


def read_reduced_frames_metadata(url: DataUrl, res_metadata) -> None:
    """
    Read reduced frame metadata and append it to 'res_metadata'.
    """
    metadata_file = url.file_path()
    data_path = url.data_path()

    if url.scheme() in ("hdf5", "silx"):
        if not os.path.exists(metadata_file):
            return
        with hdf5_open(metadata_file) as h5s:
            if data_path not in h5s:
                return
            parent_group = h5s[data_path]
            if ReducedFramesInfos.COUNT_TIME_KEY in parent_group:
                count_time = h5py_read_dataset(
                    parent_group[ReducedFramesInfos.COUNT_TIME_KEY]
                )
                unit = get_unit(
                    attrs=parent_group[ReducedFramesInfos.COUNT_TIME_KEY].attrs,
                )
                if unit is None:
                    unit = _ureg.second
                res_metadata.count_time = (count_time * unit).to(_ureg.second).magnitude

            if ReducedFramesInfos.MACHINE_ELECT_CURRENT_KEY in parent_group:
                machine_current = h5py_read_dataset(
                    parent_group[ReducedFramesInfos.MACHINE_ELECT_CURRENT_KEY]
                )

                unit = get_unit(
                    attrs=parent_group[
                        ReducedFramesInfos.MACHINE_ELECT_CURRENT_KEY
                    ].attrs,
                )
                if unit is None:
                    unit = _ureg.ampere
                res_metadata.machine_current = (
                    (machine_current * unit).to(_ureg.ampere).magnitude
                )
    elif url.scheme() == "fabio":
        # note: for fabio reading the metadata is done from 'read_reduced_frame' to avoid reading twice the file.
        _logger.warning("Reading EDF metadata is done through read_reduced_frame")
    else:
        raise ValueError(
            f"scheme {url.scheme()} is not handled. Should be fabio, silx of hdf5"
        )


def read_reduced_frames(
    url: DataUrl, read_as_url: bool, res_data, res_metadata
) -> None:
    """
    Read one reduced frame and append frame data to 'res_data' and frame metadata to 'res_metadata'.
    url must contains:
    * scheme: "fabio", "hdf5" or "silx"
    * file_path: path to the file containing the reduced frame
    * data_path: for hdf5/silx, path to the dataset or group containing the reduced frame
    * data_slice: for fabio, index of the reduced frame in the file; for hdf5/silx, index of the reduced frame in the dataset or group

    :param url: url of the reduced frame to read.
    :param read_as_url: if True, the frame data will be stored as a DataUrl instead of being read.
    :param res_data: dict where the frame data will be stored. The key is the frame index and the value is the frame data or a DataUrl if read_as_url is True.
    :param res_metadata: ReducedFramesInfos where the frame metadata will be stored. Will only be take into account if scheme is "fabio".
    """
    scheme = url.scheme()

    if scheme == "fabio":
        return _read_reduced_frame_fabio(
            url=url,
            read_as_url=read_as_url,
            res_data=res_data,
            res_metadata=res_metadata,
        )

    elif scheme in ("hdf5", "silx"):
        return _read_reduced_frame_silx(
            url=url,
            read_as_url=read_as_url,
            res_data=res_data,
        )
    else:
        raise ValueError(
            f"scheme {scheme} is not handled. Should be fabio, silx of hdf5"
        )


def _read_reduced_frame_fabio(url: DataUrl, read_as_url: bool, res_data, res_metadata):
    frame_index = url.data_slice()
    frame_file_path = url.file_path()

    if not os.path.exists(frame_file_path):
        return
    try:
        with fabio.open(frame_file_path) as handler:
            if handler.nframes > 1:
                _logger.warning(
                    f"{frame_file_path} is expected to have one frame. Only the first one will be picked"
                )
            if frame_index in res_data:
                _logger.error(f"two frames found with the same index {frame_index}")
            if read_as_url:
                res_data[frame_index] = DataUrl(
                    file_path=frame_file_path, scheme="fabio"
                )
            else:
                res_data[frame_index] = handler.data
            if "SRCUR" in handler.header:
                res_metadata.machine_current.append(float(handler.header["SRCUR"]))
            if "CountTime" in handler.header:
                res_metadata.count_time.append(float(handler.header["CountTime"]))

    except OSError as e:
        _logger.error(e)


def _read_reduced_frame_silx(url: DataUrl, read_as_url: bool, res_data):
    frame_file_path = url.file_path()
    data_path = url.data_path()
    data_path, index_zfill4_pattern = format_data_path(data_path)

    if not os.path.exists(frame_file_path):
        return
    with hdf5_open(frame_file_path) as h5s:
        dataset_or_group = h5s[data_path]
        if isinstance(dataset_or_group, h5py.Dataset):
            idx = None
            if dataset_or_group.name.isnumeric():
                try:
                    idx = int(dataset_or_group)
                except ValueError:
                    idx = None
            if read_as_url:
                res_data[idx] = DataUrl(
                    file_path=frame_file_path,
                    data_path=data_path,
                    scheme="silx",
                )
            else:
                res_data[idx] = dataset_or_group[()]
        else:
            assert isinstance(dataset_or_group, h5py.Group), (
                f"expect a group not {type(dataset_or_group)}"
            )
            # browse children
            for name, item in dataset_or_group.items():
                if isinstance(item, h5py.Dataset):
                    if name.isnumeric():
                        if index_zfill4_pattern and not len(name) == 4:
                            continue
                        else:
                            try:
                                idx = int(name)
                            except ValueError:
                                _logger.info(f"fail to cast {name} as an integer")
                                continue
                            if read_as_url:
                                res_data[idx] = DataUrl(
                                    file_path=frame_file_path,
                                    data_path=data_path + "/" + name,
                                    scheme="silx",
                                )
                            else:
                                res_data[idx] = dataset_or_group[name][()]
