from pathlib import Path

import numpy
from ewokscore import Task
from silx.io.url import DataUrl
from tomoscan.esrf.scan.nxtomoscan import NXtomoScan


class ReduceDarkFlat(
    Task,
    input_names=["nx_path"],
    optional_input_names=[
        "dark_reduction_method",
        "flat_reduction_method",
        "overwrite",
        "output_dtype",
        "return_info",
        "dark_flats_dir",
    ],
    output_names=[
        "reduced_darks_path",
        "reduced_flats_path",
    ],
):
    def run(self):
        """
        Reduce the dark and flat frames of the input NX file.

        param nx_path: Path to the input NX file.
        param dark_reduction_method: Method to reduce dark frames ('mean' or 'median').
        param flat_reduction_method: Method to reduce flat frames ('mean' or 'median').
        param overwrite: Whether to overwrite existing reduced files.
        param output_dtype: Data type for the output reduced frames.
        param return_info: Whether to return additional info from reduction.
        param dark_flats_dir: Directory to save reduced darks and flats. If provided, reduction is skipped.
        """

        nx_path = Path(self.inputs.nx_path).resolve()
        base_name = nx_path.stem
        dark_flats_dir = self.get_input_value("dark_flats_dir", None)

        if dark_flats_dir is not None:
            darks_in_dir = list(Path(dark_flats_dir).glob("*_darks.hdf5"))
            flats_in_dir = list(Path(dark_flats_dir).glob("*_flats.hdf5"))
            if len(darks_in_dir) == 0 or len(flats_in_dir) == 0:
                raise FileNotFoundError(
                    f"No reduced darks or flats found in the specified directory: {dark_flats_dir}"
                )
            self.outputs.reduced_darks_path = darks_in_dir[0]
            self.outputs.reduced_flats_path = flats_in_dir[0]
            return

        d_reduction_method = self.get_input_value("dark_reduction_method", "mean")
        f_reduction_method = self.get_input_value("flat_reduction_method", "median")
        overwrite = self.get_input_value("overwrite", True)
        output_dtype = self.get_input_value("output_dtype", numpy.float32)
        return_info = self.get_input_value("return_info", False)

        scan = NXtomoScan(str(nx_path), entry="entry0000")

        reduced_dark = scan.compute_reduced_darks(
            reduced_method=d_reduction_method,
            overwrite=overwrite,
            output_dtype=output_dtype,
            return_info=return_info,
        )
        reduced_flat = scan.compute_reduced_flats(
            reduced_method=f_reduction_method,
            overwrite=overwrite,
            output_dtype=output_dtype,
            return_info=return_info,
        )

        references_dir = nx_path.parent.parent / "references"
        references_dir.mkdir(parents=True, exist_ok=True)

        dark_file = references_dir / f"{base_name}_darks.hdf5"
        flat_file = references_dir / f"{base_name}_flats.hdf5"

        dark_urls = (
            DataUrl(
                file_path=str(dark_file),
                data_path="{entry}/darks/{index}",
                scheme=NXtomoScan.SCHEME,
            ),
        )
        dark_metadata_urls = (
            DataUrl(
                file_path=str(dark_file),
                data_path="{entry}/darks/",
                scheme=NXtomoScan.SCHEME,
            ),
        )
        flat_urls = (
            DataUrl(
                file_path=str(flat_file),
                data_path="{entry}/flats/{index}",
                scheme=NXtomoScan.SCHEME,
            ),
        )
        flat_metadata_urls = (
            DataUrl(
                file_path=str(flat_file),
                data_path="{entry}/flats/",
                scheme=NXtomoScan.SCHEME,
            ),
        )

        scan.save_reduced_darks(
            reduced_dark,
            overwrite=overwrite,
            output_urls=dark_urls,
            metadata_output_urls=dark_metadata_urls,
        )
        scan.save_reduced_flats(
            reduced_flat,
            overwrite=overwrite,
            output_urls=flat_urls,
            metadata_output_urls=flat_metadata_urls,
        )

        self.outputs.reduced_darks_path = str(dark_file)
        self.outputs.reduced_flats_path = str(flat_file)
