"""User API for XRPD processing on the Bliss repl"""

import os
import json
import shutil
import logging
from typing import Optional, List, Dict

try:
    from bliss.scanning import scan_meta
    from bliss import current_session
    from bliss.common.logtools import user_print
except ImportError:
    scan_meta = None
    current_session = None
    user_print = print

from ewoksjob.client import submit

from ..persistent import WithPersistentParameters, ParameterInfo, autocomplete_property
from ..resources.xrpd import RESOURCE_ROOT
from ..utils import directories
from .plotter import Plotter

logger = logging.getLogger(__name__)


class XrpdProcessor(
    WithPersistentParameters,
    parameters=[
        ParameterInfo("_enabled"),
        ParameterInfo("workflow_with_saving", category="workflows"),
        ParameterInfo("workflow_without_saving", category="workflows"),
        ParameterInfo("worker", category="workflows"),
        ParameterInfo("lima_names", category="data access"),
        ParameterInfo("data_from_memory", category="data access"),
        ParameterInfo("retry_timeout", category="data access"),
        ParameterInfo("retry_period", category="data access"),
        ParameterInfo("flush_period", category="data access"),
        ParameterInfo("number_of_scans", category="plotting"),
        ParameterInfo("trigger_from_bliss", category="plotting"),
        ParameterInfo("monitor_name", category="PyFai"),
        ParameterInfo("reference", category="PyFai"),
    ],
):
    def __init__(self) -> None:
        if current_session is None:
            raise ImportError("bliss")
        super().__init__()

        self._set_parameter_default("_enabled", False)
        default_workflow = os.path.join(
            RESOURCE_ROOT, "integrate_scan_with_saving.json"
        )
        self._set_parameter_default("workflow_with_saving", default_workflow)
        default_workflow = os.path.join(
            RESOURCE_ROOT, "integrate_scan_without_saving.json"
        )
        self._set_parameter_default("workflow_without_saving", default_workflow)
        self._set_parameter_default("lima_names", list())
        self._set_parameter_default("data_from_memory", True)
        self._set_parameter_default("trigger_from_bliss", True)
        self._set_parameter_default("retry_period", 1)
        self._set_parameter_default("flush_period", 5)
        self._set_parameter_default("number_of_scans", 4)

        self._plotter = Plotter(number_of_scans=self.number_of_scans)
        self._plotter.replot()
        self._sync_scan_metadata()

    @property
    def plotter(self) -> Plotter:
        return self._plotter

    def clear_plots(self) -> None:
        return self.plotter.clear()

    def replot(self) -> None:
        return self.plotter.replot()

    def purge_plotting_tasks(self) -> int:
        return self.plotter.purge_tasks()

    def stop_plotting_tasks(self) -> int:
        return self.plotter.kill_tasks()

    @autocomplete_property
    def number_of_scans(self) -> Optional[int]:
        return self._get_parameter("number_of_scans")

    @number_of_scans.setter
    def number_of_scans(self, value: int):
        self.plotter.number_of_scans = value
        self._set_parameter("number_of_scans", self.plotter.number_of_scans)

    def _info_categories(self) -> Dict[str, dict]:
        categories = super()._info_categories()
        ntasks = self.plotter.purge_tasks()
        categories["status"] = {"Enabled": self._enabled, "Plotting tasks": ntasks}
        return categories

    def enable(self, *detectors) -> None:
        self._enabled = True
        self.set_lima_names(*detectors)
        self._sync_scan_metadata()

    def disable(self) -> None:
        self._enabled = False
        self._sync_scan_metadata()

    def _sync_scan_metadata(self):
        scan_meta_obj = scan_meta.get_user_scan_meta()
        if not self._enabled:
            try:
                # Before Bliss 1.11
                keys = list(scan_meta_obj._metadata.keys())
                for key in keys:
                    if key.name.lower() == "workflows":
                        scan_meta_obj._metadata.pop(key)
            except Exception:
                pass
            # Since Bliss 1.11
            scan_meta_obj.remove_categories({"workflows"})
            try:
                _ = list(scan_meta_obj._metadata_keys())
            except AttributeError:
                pass
            return
        if "workflows" not in scan_meta_obj.categories_names():
            scan_meta_obj.add_categories({"workflows"})
        scan_meta_obj.workflows.timing = scan_meta.META_TIMING.START
        scan_meta_obj.workflows.set("@NX_class", {"@NX_class": "NXcollection"})
        scan_meta_obj.workflows.set("processing", self.on_new_scan)

    def set_lima_names(self, *detectors) -> None:
        lima_names = set()
        for detector in detectors:
            lima_names.add(detector.name)
        if lima_names:
            self.lima_names = sorted(lima_names)

    def on_new_scan(self, scan) -> Optional[dict]:
        """Executed at the start of every scan"""
        if not self.scan_requires_processing(scan):
            return
        self.ensure_workflow_accessible(scan)
        if not self.check_monitor(scan):
            user_print(
                f"Monitor '{self.monitor_name}' is not included in the scan. Remove it from the pyfai options."
            )
            return

        metadata = dict()
        for lima_name in self.get_lima_names(scan):
            # Submit arguments
            workflow = self.get_workflow(scan)
            kwargs = self.get_submit_arguments(scan, lima_name)
            if scan.scan_info.get("save"):
                kwargs["convert_destination"] = self.workflow_destination(
                    scan, lima_name
                )

            if self.trigger_from_bliss:
                # Trigger workflow from the current process.
                future = submit(args=(workflow,), kwargs=kwargs, queue=self.worker)
                filename = self.get_filename(scan)
                scan_nb = scan.scan_info.get("scan_nb")

                scan_name = os.path.splitext(os.path.basename(filename))[0]
                scan_name = f"{scan_name}: {scan_nb}.1 {scan.name}"
                if scan.scan_info.get("save"):
                    if len(self.lima_names) > 1:
                        output_filename = self.save_nexus_filename(
                            scan, lima_name=lima_name
                        )
                    else:
                        output_filename = self.save_nexus_filename(scan)
                    output_url = f"{output_filename}::/{scan_nb}.1/{lima_name}_integrate/integrated"
                else:
                    output_url = None
                self.plotter.handle_workflow_result(
                    future,
                    lima_name,
                    scan_name,
                    output_url=output_url,
                    retry_timeout=self.retry_timeout,
                    retry_period=self.retry_period,
                )
            else:
                # Save metadata to Redis (is saved in the raw data file as well).
                # Workflows may be triggered by an external process as a result.
                kwargs["workflow"] = workflow
                metadata[f"process_{lima_name}"] = {
                    "@NX_class": "NXprocess",
                    "program": "ewoks",
                    "configuration": {
                        "@NX_class": "NXnote",
                        "type": "application/json",
                        "data": json.dumps(kwargs),
                    },
                }

        return metadata

    def get_submit_arguments(self, scan, lima_name) -> dict:
        return {
            "inputs": self.get_inputs(scan, lima_name),
            "outputs": [{"all": False}],
        }

    def get_workflow(self, scan) -> str:
        """Get the workflow filename for the scan"""
        if scan.scan_info.get("save"):
            return self.workflow_with_saving
        else:
            return self.workflow_without_saving

    def set_workflow(self, scan, filename) -> None:
        """Set the workflow filename for the scan"""
        if scan.scan_info.get("save"):
            self.workflow_with_saving = filename
        else:
            self.workflow_without_saving = filename

    def get_filename(self, scan) -> str:
        filename = scan.scan_info.get("filename")
        if filename:
            return filename
        return current_session.scan_saving.filename

    def ensure_workflow_accessible(self, scan) -> None:
        """Copy the workflow file so that the worker will be able to access it"""
        wfilename = self.get_workflow(scan)
        if not os.path.exists(wfilename):
            wfilename = os.path.join(RESOURCE_ROOT, "integrate_scan_with_saving.json")
        if not wfilename.startswith(RESOURCE_ROOT):
            return  # visitor or inhouse folder
        dfilename = self.get_filename(scan)
        root = directories.get_workflows_dir(dfilename)
        os.makedirs(root, exist_ok=True)
        shutil.copy(wfilename, root)
        wfilename = os.path.join(root, os.path.basename(wfilename))
        self.set_workflow(scan, wfilename)

    def get_lima_names(self, scan) -> List[str]:
        channels = scan.scan_info.get("channels", dict())
        return sorted(
            (
                lima_name
                for lima_name in self.lima_names
                if f"{lima_name}:image" in channels
            )
        )

    def check_monitor(self, scan) -> bool:
        monitor_name = self.monitor_name
        if not monitor_name:
            return True
        channels = scan.scan_info.get("channels", dict())
        return any(monitor_name in name for name in channels)

    def scan_requires_processing(self, scan) -> bool:
        return bool(self.get_lima_names(scan))

    def scan_processed_directory(self, scan) -> str:
        filename = self.get_filename(scan)
        root = directories.get_processed_dir(filename)
        collection = os.path.basename(directories.get_collection_dir(filename))
        dataset = os.path.basename(directories.get_dataset_dir(filename))
        return os.path.join(root, collection, dataset)

    def workflow_destination(self, scan, lima_name) -> str:
        filename = self.get_filename(scan)
        scan_nb = scan.scan_info.get("scan_nb")
        root = self.scan_processed_directory(scan)
        stem = os.path.splitext(os.path.basename(filename))[0]
        basename = f"{stem}_{scan_nb}_{lima_name}.json"
        return os.path.join(root, basename)

    def save_nexus_filename(self, scan, lima_name: Optional[str] = None) -> str:
        filename = self.get_filename(scan)
        root = self.scan_processed_directory(scan)
        basename = os.path.basename(filename)
        if lima_name:
            name, ext = os.path.splitext(basename)
            basename = f"{name}_{lima_name}{ext}"
        return os.path.join(root, basename)

    def save_nexus_url(self, scan, lima_name: Optional[str] = None) -> str:
        scan_nb = scan.scan_info.get("scan_nb")
        filename = self.save_nexus_filename(scan, lima_name=lima_name)
        return f"silx://{filename}::{scan_nb}.1"

    def get_inputs(self, scan, lima_name: str) -> List[dict]:
        inputs = self.get_config_inputs(lima_name)
        inputs += self.get_integrate_inputs(scan, lima_name, "IntegrateBlissScan")
        inputs += self.get_integrate_inputs(
            scan, lima_name, "IntegrateBlissScanWithoutSaving"
        )
        inputs += self.get_save_inputs(scan, lima_name, "SaveNexusIntegrated")
        return inputs

    def get_integrate_inputs(
        self, scan, lima_name: str, task_identifier: str
    ) -> List[dict]:
        filename = self.get_filename(scan)
        scan_nb = scan.scan_info.get("scan_nb")
        inputs = [
            {
                "task_identifier": task_identifier,
                "name": "filename",
                "value": filename,
            },
            {
                "task_identifier": task_identifier,
                "name": "output_filename",
                "value": self.save_nexus_filename(scan),
            },
            {
                "task_identifier": task_identifier,
                "name": "scan",
                "value": scan_nb,
            },
            {
                "task_identifier": task_identifier,
                "name": "monitor_name",
                "value": self.monitor_name,
            },
            {
                "task_identifier": task_identifier,
                "name": "reference",
                "value": self.reference,
            },
            {
                "task_identifier": task_identifier,
                "name": "maximum_persistent_workers",
                "value": len(self.lima_names),
            },
            {
                "task_identifier": task_identifier,
                "name": "detector_name",
                "value": lima_name,
            },
        ]
        if len(self.lima_names) > 1:
            inputs += [
                {
                    "task_identifier": task_identifier,
                    "name": "external_output_filename",
                    "value": self.save_nexus_filename(scan, lima_name=lima_name),
                }
            ]
        if self.data_from_memory:
            scan_memory_url = f"{scan.root_node.db_name}:{scan._node_name}"
            inputs.append(
                {
                    "task_identifier": task_identifier,
                    "name": "scan_memory_url",
                    "value": scan_memory_url,
                }
            )
        if self.retry_timeout:
            inputs.append(
                {
                    "task_identifier": task_identifier,
                    "name": "retry_timeout",
                    "value": self.retry_timeout,
                }
            )
        if self.flush_period:
            inputs.append(
                {
                    "task_identifier": task_identifier,
                    "name": "flush_period",
                    "value": self.flush_period,
                }
            )
        return inputs

    def get_save_inputs(self, scan, lima_name: str, task_identifier: str) -> List[dict]:
        inputs = [
            {
                "task_identifier": task_identifier,
                "name": "url",
                "value": self.save_nexus_url(scan),
            }
        ]
        if len(self.lima_names) > 1:
            inputs += [
                {
                    "task_identifier": task_identifier,
                    "name": "external_url",
                    "value": self.save_nexus_url(scan, lima_name=lima_name),
                }
            ]
        return inputs

    def get_config_inputs(self, lima_name: str) -> List[dict]:
        inputs = []
        filename = self.get_config_filename(lima_name)
        if filename:
            inputs.append(
                {
                    "task_identifier": "PyFaiConfig",
                    "name": "filename",
                    "value": filename,
                }
            )
        integration_options = self.get_integration_options(lima_name)
        if integration_options:
            try:
                integration_options = integration_options.to_dict()
            except AttributeError:
                pass
            inputs.append(
                {
                    "task_identifier": "PyFaiConfig",
                    "name": "integration_options",
                    "value": integration_options,
                }
            )
        return inputs

    def get_config_filename(self, lima_name: str) -> Optional[str]:
        raise NotImplementedError

    def get_integration_options(self, lima_name: str) -> Optional[dict]:
        raise NotImplementedError
