"""Automatic pyfai integration for every scan (when enabled)"""

import os
import json
from typing import Dict, Optional, Sequence

try:
    from bliss.scanning import scan_meta
    from bliss import current_session
except ImportError:
    scan_meta = None
    current_session = None

from ..persistent import WithPersistentParameters
from ..utils.directories import get_processed_dir


class XrpdProcessor(
    WithPersistentParameters,
    parameters=[
        "enabled",
        "detector_name",
        "counter_names",
        "pyfai_config",
        "integration_options",
        "retry_timeout",
        "monitor_name",
        "reference",
    ],
):
    def __init__(self) -> None:
        super().__init__()
        if current_session is None:
            raise ImportError("blissdata")
        self._set_parameter_default("enabled", False)
        self._set_parameter_default(
            "integration_options",
            {
                "method": "csr_ocl_gpu",
                "integrator_name": "sigma_clip_ng",
                "extra_options": {"max_iter": 3, "thres": 0},
                "nbpt_rad": 1024,
                "unit": "q_nm^-1",
            },
        )
        self._sync_scan_metadata()

    def enable(
        self,
        detector_name: str,
        counter_names: Optional[Sequence[str]],
        pyfai_config: Optional[str] = None,
        integration_options: Optional[Dict] = None,
    ):
        self.enabled = True
        self.detector_name = detector_name
        self.counter_names = counter_names
        if pyfai_config is not None:
            self.pyfai_config = pyfai_config
        if integration_options is not None:
            self.integration_options = integration_options
        self._sync_scan_metadata()

    def disable(self):
        self.enabled = False
        self._sync_scan_metadata()

    def _sync_scan_metadata(self):
        scan_meta_obj = scan_meta.get_user_scan_meta()
        if self.enabled:
            scan_meta_obj.add_categories({"workflows"})
            scan_meta_obj.workflows.set("@NX_class", {"@NX_class": "NXcollection"})
            scan_meta_obj.workflows.set("nxprocess1", self.get_nxprocess1_content)
        else:
            scan_meta_obj.remove_categories({"workflows"})

    def get_nxprocess1_content(self, scan):
        if not scan.scan_info.get("save"):
            return
        if not scan.scan_info.get("filename"):
            return
        if not scan.scan_info.get("scan_nb"):
            return
        detector_name = self.detector_name
        assert detector_name, "detector name is missing"
        detector_info = scan.scan_info.get("channels", dict()).get(
            f"{detector_name}:image"
        )
        if not detector_info:
            return None
        return {
            "xrpd_data_reduction": {
                "@NX_class": "NXprocess",
                "program": "ewoks",
                "configuration": self.get_configuration(scan),
            }
        }

    def get_configuration(self, scan):
        return {
            "@NX_class": "NXnote",
            "type": "application/json",
            "data": json.dumps(self.get_workflow()),
            "options": json.dumps(self.get_options(scan)),
        }

    def get_options(self, scan):
        return {"inputs": self.get_inputs(scan)}

    def get_inputs(self, scan):
        filename = scan.scan_info.get("filename")
        scan_nb = scan.scan_info.get("scan_nb")
        output_filename = os.path.join(
            get_processed_dir(filename), os.path.basename(filename)
        )
        pyfai_config = self.pyfai_config
        assert pyfai_config, "pyfai configuration is missing"
        detector_name = self.detector_name
        assert detector_name, "detector name is missing"

        integration_options = self.integration_options
        counter_names = self.counter_names
        retry_timeout = self.retry_timeout
        monitor_name = self.monitor_name
        reference = self.reference

        lst = [
            {"id": "config", "name": "filename", "value": pyfai_config},
            {"id": "integrate", "name": "filename", "value": filename},
            {"id": "integrate", "name": "scan", "value": scan_nb},
            {"id": "integrate", "name": "detector_name", "value": detector_name},
            {"id": "integrate", "name": "output_filename", "value": output_filename},
            {"id": "integrate", "name": "monitor_name", "value": monitor_name},
            {"id": "integrate", "name": "reference", "value": reference},
        ]
        if integration_options:
            lst.append(
                {
                    "id": "config",
                    "name": "integration_options",
                    "value": integration_options.to_dict(),
                }
            )
        if counter_names:
            lst.append(
                {"id": "integrate", "name": "counter_names", "value": counter_names}
            )
        if retry_timeout:
            lst.append(
                {"id": "integrate", "name": "retry_timeout", "value": retry_timeout}
            )
        return lst

    def get_workflow(self):
        return {
            "graph": {"id": "xrpd_data_reduction"},
            "nodes": [
                {
                    "id": "config",
                    "task_type": "class",
                    "task_identifier": "ewoksxrpd.tasks.pyfaiconfig.PyFaiConfig",
                },
                {
                    "id": "integrate",
                    "task_type": "class",
                    "task_identifier": "ewoksxrpd.tasks.integrate.IntegrateBlissScan",
                },
            ],
            "links": [
                {
                    "source": "config",
                    "target": "integrate",
                    "data_mapping": [
                        {"source_output": "energy", "target_input": "energy"},
                        {"source_output": "detector", "target_input": "detector"},
                        {"source_output": "geometry", "target_input": "geometry"},
                        {"source_output": "mask", "target_input": "mask"},
                        {
                            "source_output": "integration_options",
                            "target_input": "integration_options",
                        },
                    ],
                }
            ],
        }
