from __future__ import annotations

import json
import logging
import os
from typing import Any
from typing import Dict
from typing import List
from typing import Optional

import gevent
from esrf_pathlib import ESRFPath
from ewoksjob.client import get_future
from ewoksjob.client import submit

from ..persistent.parameters import ParameterInfo
from ..processor import BaseProcessor
from ..processor import BlissScanType
from ..resources import resource_filename
from .xrr_plotter import Bm32XrrPlotter

logger = logging.getLogger(__name__)


class Bm32XrrProcessor(
    BaseProcessor,
    parameters=[
        ParameterInfo("queue", category="workflows"),
        ParameterInfo(
            "workflow", category="workflows", doc="Main XRR processing workflow (.json)"
        ),
        ParameterInfo("data_portal_upload", category="data portal"),
        ParameterInfo("plot_name", category="plotting"),
        ParameterInfo(
            "tth_path",
            category="XRR (dataset paths)",
            doc="e.g. '{scan}.1/measurement/tth'",
        ),
        ParameterInfo(
            "det_path",
            category="XRR (dataset paths)",
            doc="e.g. '{scan}.1/measurement/det'",
        ),
        ParameterInfo(
            "mon_path",
            category="XRR (dataset paths)",
            doc="e.g. '{scan}.1/measurement/mon4'",
        ),
        ParameterInfo(
            "energy", category="Corrections/Geometry", doc="X-ray energy (eV)"
        ),
        ParameterInfo(
            "itime", category="Corrections/Geometry", doc="Integration time (seconds)"
        ),
        ParameterInfo(
            "dark_current",
            category="Corrections/Geometry",
            doc="Dark current (scalar or 1D)",
        ),
        ParameterInfo(
            "direct_beam_amp",
            category="Corrections/Geometry",
            doc="Direct beam amplitude (fit prior)",
        ),
        ParameterInfo(
            "beam_size", category="Corrections/Geometry", doc="Beam size on sample (mm)"
        ),
        ParameterInfo(
            "sample_length",
            category="Corrections/Geometry",
            doc="Illuminated sample length (mm)",
        ),
        ParameterInfo(
            "tthmin",
            category="Corrections/Geometry",
            doc="Lower angle cutoff for processing (deg)",
        ),
        ParameterInfo(
            "qvaluesRangesToRemove",
            category="Clean Signal",
            doc="List of q-ranges to remove from the signal (e.g. [ [0.1,0.15], [0.25,0.3] ])",
        ),
        ParameterInfo(
            "qmax",
            category="Clean Signal",
            doc="Maximum q-value to consider (1/Angstrom)",
        ),
        ParameterInfo(
            "peak_prominence",
            category="Inverse FT",
            doc="Peak prominence for find_peaks (Inverse FT)",
        ),
        ParameterInfo(
            "baseline_type", category="Inverse FT", doc="Baseline type (Inverse FT)"
        ),
        ParameterInfo(
            "peak_distance",
            category="Inverse FT",
            doc="Minimum peak distance (Inverse FT)",
        ),
        ParameterInfo(
            "zrange", category="Inverse FT", doc="Z range to consider (Inverse FT)"
        ),
    ],
):
    """Trigger the BM32 XRR workflow from BLISS scan objects."""

    def __init__(
        self,
        config: Optional[Dict[str, Any]] = None,
        defaults: Optional[Dict[str, Any]] = None,
    ) -> None:
        if defaults is None:
            defaults = {}

        defaults.setdefault("workflow", "xrr_processor.json")
        defaults.setdefault("data_portal_upload", False)
        defaults.setdefault("plot_name", "XRR analysis")

        defaults.setdefault("tth_path", "{scan}.1/measurement/tth")
        defaults.setdefault("det_path", "{scan}.1/measurement/det")
        defaults.setdefault("mon_path", "{scan}.1/measurement/mon4")

        defaults.setdefault("itime", None)
        defaults.setdefault("dark_current", 0.0)
        defaults.setdefault("direct_beam_amp", None)
        defaults.setdefault("beam_size", None)
        defaults.setdefault("sample_length", None)
        defaults.setdefault("tthmin", None)
        defaults.setdefault("energy", None)
        defaults.setdefault("qvaluesRangesToRemove", [])
        defaults.setdefault("qmax", 0.8)
        defaults.setdefault("peak_prominence", 1e-10)
        defaults.setdefault("baseline_type", "pchip")
        defaults.setdefault("peak_distance", None)
        defaults.setdefault("zrange", 100)

        super().__init__(config=config, defaults=defaults)
        self._tasks: List[Any] = []
        self.plotter = Bm32XrrPlotter()

    def _info_categories(self) -> Dict[str, dict]:
        categories = super()._info_categories()
        categories.setdefault("status", {})["Workflow tasks"] = self.purge_tasks()
        if self.plotter is not None:
            categories.setdefault("status", {})[
                "Plotting tasks"
            ] = self.plotter.purge_tasks()
        return categories

    def purge_tasks(self) -> int:
        self._tasks = [task for task in self._tasks if task]
        return len(self._tasks)

    def _spawn(self, *args, **kwargs) -> None:
        self._tasks.append(gevent.spawn(*args, **kwargs))
        self.purge_tasks()

    @staticmethod
    def _consume_future_result(future, scan_nb: int) -> None:
        try:
            future.result()
        except Exception:
            logger.exception("BM32 XRR workflow failed for scan %s", scan_nb)

    @staticmethod
    def _scan_number_from_scan(scan: BlissScanType) -> int:
        return int(scan.scan_info["scan_nb"])

    @staticmethod
    def _filename_from_scan(scan: BlissScanType) -> str:
        return scan.scan_info["filename"]

    @staticmethod
    def _get_output_dir(master_filename: str) -> str:
        return str(ESRFPath(master_filename).processed_dataset_path)

    def _set_output_basename(
        self, scan: BlissScanType, bg_scan_nb: Optional[int] = None
    ) -> str:
        scan_saving = scan.scan_saving
        sample = scan_saving.sample_name or "sample"
        dataset = scan_saving.dataset_name or "dataset"
        scan_nb = self._scan_number_from_scan(scan)
        return f"{sample}_{dataset}_{scan_nb}_{bg_scan_nb}"

    def _ensure_ready(self, scan: BlissScanType) -> None:
        self._scan_number_from_scan(scan)
        master_filename = self._filename_from_scan(scan)
        if not master_filename:
            raise ValueError("master filename is required")

    def workflow_destination(
        self, scan: BlissScanType, background: Optional[BlissScanType] = None
    ) -> str:
        master_filename = self._filename_from_scan(scan)
        output_dir = self._get_output_dir(master_filename)
        bg_scan_nb = (
            self._scan_number_from_scan(background) if background is not None else None
        )
        output_basename = self._set_output_basename(scan, bg_scan_nb)
        return os.path.join(output_dir, f"{output_basename}.json")

    def export_destination(
        self, scan: BlissScanType, background: Optional[BlissScanType] = None
    ) -> str:
        master_filename = self._filename_from_scan(scan)
        output_dir = self._get_output_dir(master_filename)
        bg_scan_nb = (
            self._scan_number_from_scan(background) if background is not None else None
        )
        output_basename = self._set_output_basename(scan, bg_scan_nb)
        return os.path.join(output_dir, f"{output_basename}.hdf5")

    def get_inputs(
        self, scan: BlissScanType, background: Optional[BlissScanType] = None
    ) -> List[Dict[str, Any]]:
        scan_nb = self._scan_number_from_scan(scan)
        master_filename = self._filename_from_scan(scan)
        bg_master_filename = (
            self._filename_from_scan(background) if background is not None else ""
        )
        bg_scan_nb = (
            self._scan_number_from_scan(background) if background is not None else None
        )
        output_basename = self._set_output_basename(scan, bg_scan_nb)

        inputs: List[Dict[str, Any]] = [
            {
                "task_identifier": "ewoksbm32.tasks.xrr.load_raw_data.LoadRawData",
                "name": "filename",
                "value": master_filename,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.load_raw_data.LoadRawData",
                "name": "tth_path",
                "value": self.tth_path.format(scan=scan_nb),
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.load_raw_data.LoadRawData",
                "name": "det_path",
                "value": self.det_path.format(scan=scan_nb),
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.load_raw_data.LoadRawData",
                "name": "mon_path",
                "value": self.mon_path.format(scan=scan_nb),
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.load_raw_data.LoadRawData",
                "name": "bg_nx_data_url",
                "value": bg_master_filename or master_filename,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.load_raw_data.LoadRawData",
                "name": "itime",
                "value": self.itime,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.preprocess.Preprocess",
                "name": "energy",
                "value": self.energy,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.preprocess.Preprocess",
                "name": "itime",
                "value": self.itime,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.preprocess.Preprocess",
                "name": "dark_current",
                "value": self.dark_current,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.preprocess.Preprocess",
                "name": "direct_beam_amp",
                "value": self.direct_beam_amp,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.preprocess.Preprocess",
                "name": "beam_size",
                "value": self.beam_size,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.preprocess.Preprocess",
                "name": "sample_length",
                "value": self.sample_length,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.preprocess.Preprocess",
                "name": "tthmin",
                "value": self.tthmin,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.clean_signal.CleanSignal",
                "name": "qvaluesRangesToRemove",
                "value": self.qvaluesRangesToRemove,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.clean_signal.CleanSignal",
                "name": "qmax",
                "value": self.qmax,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.inverse_ft.InverseFT",
                "name": "peak_prominence",
                "value": self.peak_prominence,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.inverse_ft.InverseFT",
                "name": "baseline_type",
                "value": self.baseline_type,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.inverse_ft.InverseFT",
                "name": "peak_distance",
                "value": self.peak_distance,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.inverse_ft.InverseFT",
                "name": "zrange",
                "value": self.zrange,
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.export_data.ExportHDF5",
                "name": "filename",
                "value": self.export_destination(scan, background),
            },
            {
                "task_identifier": "ewoksbm32.tasks.xrr.export_data.ExportHDF5",
                "name": "sample_name",
                "value": output_basename,
            },
        ]

        return [item for item in inputs if item.get("value") is not None]

    def _get_workflow(self) -> Dict[str, Any]:
        with open(resource_filename("bm32", self.workflow), "r") as wf:
            return json.load(wf)

    def get_submit_arguments(
        self, scan: BlissScanType, background: Optional[BlissScanType] = None
    ) -> Dict[str, Any]:
        return {
            "inputs": self.get_inputs(scan, background),
            "outputs": [{"all": True}],
        }

    def _get_workflow_upload_parameters(self, scan: BlissScanType) -> Dict[str, Any]:
        master_filename = self._filename_from_scan(scan)
        output_dir = self._get_output_dir(master_filename)
        dataset_metadata = scan.scan_saving.dataset or {}
        sample_name = scan.scan_saving.sample_name
        esrf_path = ESRFPath(master_filename)
        raw_directory = esrf_path.raw_dataset_path
        beamline = raw_directory.beamline
        proposal = raw_directory.proposal
        raw_directory = str(raw_directory)
        metadata = {"Sample_name": dataset_metadata.get("Sample_name", sample_name)}
        return {
            "beamline": beamline,
            "proposal": proposal,
            "dataset": "XRR_processing",
            "path": output_dir,
            "raw": [raw_directory],
            "metadata": metadata,
        }

    def _trigger_plotting(
        self,
        future,
        scan: BlissScanType,
        background: Optional[BlissScanType] = None,
    ) -> None:
        if self.plotter is None:
            return

        scan_label = f"scan {self._scan_number_from_scan(scan)}"
        if background is not None:
            scan_label += f" (bg {self._scan_number_from_scan(background)})"

        scan_saving = scan.scan_saving
        sample = scan_saving.sample_name or "sample"
        dataset = scan_saving.dataset_name or "dataset"
        display_name = f"{sample} | {dataset} | {scan_label}"
        self.plotter.handle_workflow_result(
            future=future,
            sample_name=display_name,
            plot_name=self.plot_name,
        )

    def run_processing(
        self,
        scan: BlissScanType,
        background: Optional[BlissScanType] = None,
    ) -> None:
        self._ensure_ready(scan)
        workflow = self._get_workflow()
        kwargs = self.get_submit_arguments(scan, background)
        kwargs["convert_destination"] = self.workflow_destination(scan, background)
        if self.data_portal_upload:
            kwargs["upload_parameters"] = self._get_workflow_upload_parameters(scan)

        future = submit(args=(workflow,), kwargs=kwargs, queue=self.queue)
        future = get_future(future.uuid)
        self._trigger_plotting(future, scan, background)
        self._spawn(
            self._consume_future_result, future, self._scan_number_from_scan(scan)
        )
