"""
   Wrapper around the job class to build a workflow element (production step + job)
"""

__RCSID__ = "$Id$"

# generic imports
from copy import deepcopy
import json
from typing import Literal


# DIRAC imports
import DIRAC
from DIRAC.ProductionSystem.Client.ProductionStep import ProductionStep
from CTADIRAC.Core.Utilities.tool_box import get_dataset_MQ
from DIRAC.Resources.Catalog.FileCatalogClient import FileCatalogClient


class WorkflowElementTypes:
    allowed_simulation_types: list[str] = ["mcsimulation"]
    allowed_processing_types: list[str] = [
        "ctapipeprocessing",
        "evndispprocessing",
        "merging",
        "ctapipetrainenergy",
        "ctapipetrainclassifier",
        "ctapipeapplymodels",
    ]
    we_types: list[str] = allowed_simulation_types + allowed_processing_types


class WorkflowElementDefinition:
    """Defines variables depending on WorkflowElement type"""

    def __init__(self, we_type: str) -> None:
        self.allowed_simulation_types: list[
            str
        ] = WorkflowElementTypes.allowed_simulation_types
        self.allowed_processing_types: list[
            str
        ] = WorkflowElementTypes.allowed_processing_types
        self.cpu_time: float = 259200.0
        self.mandatory_keys: set[str] = {"MCCampaign", "configuration_id", "version"}
        self.constrained_input_keys: set[str] = {
            "pointing_dir",
            "zenith_angle",
            "moon",
            "sct",
            "div_ang",
        }
        self.constrained_job_keys: set[str] = {"catalogs", "moon", "group_size"}
        self.file_meta_fields: set[str] = {"nsb", "div_ang"}
        self.mandatory_job_config_keys: dict = {}
        self.we_type: Literal(WorkflowElementTypes.we_types) = we_type
        self.prod_step = ProductionStep()
        self.prod_step.Type = "DataReprocessing"
        self.initialize()

    def initialize(self) -> None:
        """Initialize variables depending on WE type"""
        if self.we_type == "mcsimulation":
            from CTADIRAC.Interfaces.API.MCPipeNSBJob import MCPipeNSBJob

            self.job = MCPipeNSBJob()
            self.prod_step.Type = "MCSimulation"
            self.prod_step.Name = "MCSimulation"
            self.constrained_job_keys.update(
                [
                    "version",
                    "sct",
                    "particle",
                    "pointing_dir",
                    "magic",
                    "div_ang",
                ]
            )
            self.constrained_job_keys.remove("group_size")
        elif self.we_type == "ctapipeprocessing":
            from CTADIRAC.Interfaces.API.CtapipeProcessJob import CtapipeProcessJob

            self.job = CtapipeProcessJob(cpuTime=self.cpu_time)
            self.job.input_limit = None
            self.prod_step.Name = "CtapipeProcessing"
            self.mandatory_job_config_keys = {"data_level", "output_extension"}

            self.file_meta_fields.add("split")
        elif self.we_type == "evndispprocessing":
            from CTADIRAC.Interfaces.API.EvnDispSingJob import EvnDispSingJob

            self.job = EvnDispSingJob(cpuTime=self.cpu_time)
            self.job.input_limit = None
            self.prod_step.Name = "EvnDisp"
        elif self.we_type == "merging":
            from CTADIRAC.Interfaces.API.CtapipeMergeJob import CtapipeMergeJob

            self.job = CtapipeMergeJob(cpuTime=self.cpu_time)
            self.job.input_limit = None
            self.job.merged = 0
            self.prod_step.Type = "Merging"
            self.prod_step.Name = "Merging"
            self.mandatory_job_config_keys = {"output_extension"}
            self.file_meta_fields.add("split")
        elif self.we_type == "ctapipetrainenergy":
            from CTADIRAC.Interfaces.API.CtapipeTrainEnergyJob import (
                CtapipeTrainEnergyJob,
            )

            self.job = CtapipeTrainEnergyJob(cpuTime=self.cpu_time)
            self.job.input_limit = None
            self.job.merged = 0
            self.prod_step.Name = "CtapipeTrainEnergy"
        elif self.we_type == "ctapipetrainclassifier":
            from CTADIRAC.Interfaces.API.CtapipeTrainClassifierJob import (
                CtapipeTrainClassifierJob,
            )

            self.job = CtapipeTrainClassifierJob(cpuTime=self.cpu_time)
            self.job.input_limit = None
            self.job.merged = 0
            self.prod_step.Name = "CtapipeTrainClassifier"
        elif self.we_type == "ctapipeapplymodels":
            from CTADIRAC.Interfaces.API.CtapipeApplyModelsJob import (
                CtapipeApplyModelsJob,
            )

            self.job = CtapipeApplyModelsJob(cpuTime=self.cpu_time)
            self.job.input_limit = None
            self.job.merged = 0
            self.prod_step.Name = "CtapipeApplyModels"
            self.file_meta_fields.add("split")
        else:
            from DIRAC.Interfaces.API.Job import Job

            self.job = Job()
            self.mandatory_keys = {}
        self.job.setOutputSandbox(["*Log.txt"])


class WorkflowElement(WorkflowElementDefinition):
    """Composite class for workflow element (production step + job)"""

    def __init__(self, parent_prod_step: int, we_type: str) -> None:
        WorkflowElementDefinition.__init__(self, we_type)
        self.prod_step.ParentStep = parent_prod_step
        self.fc = FileCatalogClient()

    def set_constrained_job_attribute(self, key, value) -> None:
        """Set job attribute with constraints"""
        if self.we_type in self.allowed_simulation_types:
            self.set_constrained_mc_job_attribute(key, value)
        elif self.we_type in self.allowed_processing_types:
            self.set_constrained_processing_job_attribute(key, value)
        else:
            setattr(self.job, key, value)

    def set_constrained_mc_job_attribute(self, key, value) -> None:
        """Set simulation job attribute with constraints"""
        if key == "catalogs":
            # remove whitespaces between catalogs if there are some and separate between commas
            setattr(self.job, key, json.dumps(value.replace(", ", ",").split(sep=",")))
        elif key == "version":
            if "-sc" not in str(
                value
            ):  # if the version has not already been set by 'sct'
                self.job.version = str(value)
        elif key == "particle":
            self.job.set_particle(value)
        elif key == "pointing_dir":
            self.job.set_pointing_dir(value)
        elif key == "sct":
            self.job.set_sct(value)
        elif key == "moon":
            self.job.set_moon(value.replace(", ", ",").split(sep=","))
        elif key == "magic":
            self.job.set_magic(value)
        elif key == "div_ang":
            self.job.set_div_ang(value)

    def set_constrained_processing_job_attribute(self, key, value) -> None:
        """Set processing job attribute with constraints"""
        if key == "catalogs":
            # remove whitespaces between catalogs if there are some and separate between commas
            setattr(self.job, key, json.dumps(value.replace(", ", ",").split(sep=",")))
        elif key == "group_size":
            setattr(self.job, key, value)
            self.prod_step.GroupSize = self.job.group_size
        elif key == "moon":
            if value == "dark":
                self.job.output_file_metadata["nsb"] = 1
            elif value == "half":
                self.job.output_file_metadata["nsb"] = 5
            elif value == "full":
                self.job.output_file_metadata["nsb"] = 19

    def set_constrained_input_query(self, key, value) -> None:
        """Set input meta query with constraints"""
        if self.we_type in self.allowed_processing_types:
            self.set_constrained_processing_input_query(key, value)
        else:
            self.prod_step.Inputquery[key] = value

    def set_constrained_processing_input_query(self, key, value) -> None:
        """Set processing input meta query with constraints"""
        if key == "pointing_dir":
            if value == "North":
                self.prod_step.Inputquery["phiP"] = 180
            elif value == "South":
                self.prod_step.Inputquery["phiP"] = 0
        elif key == "zenith_angle":
            self.prod_step.Inputquery["thetaP"] = float(value)
        elif key == "sct":
            self.prod_step.Inputquery["sct"] = str(value)
        elif key == "moon":
            if value == "dark":
                self.prod_step.Inputquery["nsb"] = 1
            elif value == "half":
                self.prod_step.Inputquery["nsb"] = 5
            elif value == "full":
                self.prod_step.Inputquery["nsb"] = 19
        elif key == "div_ang":
            self.prod_step.Inputquery["div_ang"] = str(value)

    def build_job_attributes_from_job_config(self, workflow_step) -> None:
        """Build job attributes with job_config values"""
        for key in self.mandatory_job_config_keys:
            if key not in workflow_step["job_config"]:
                DIRAC.gLogger.error(f"{key} is mandatory")
                DIRAC.exit(-1)
        for key, value in workflow_step["job_config"].items():
            if value is not None:
                if key in self.constrained_job_keys:
                    self.set_constrained_job_attribute(key, value)
                elif key in self.file_meta_fields:
                    self.job.output_file_metadata[key] = value
                else:
                    setattr(self.job, key, value)
            elif key in self.mandatory_keys:
                DIRAC.gLogger.error(f"{key} is mandatory")
                DIRAC.exit(-1)
        for key in self.mandatory_keys:
            if key not in self.job.__dict__:
                DIRAC.gLogger.error(f"{key} is mandatory")
                DIRAC.exit(-1)

    def build_job_attributes_from_common(self, workflow_config) -> None:
        """Build job attributes concerning the whole production"""
        for key, value in workflow_config["Common"].items():
            if value is not None:
                setattr(self.job, key, value)
            elif key in self.mandatory_keys:
                DIRAC.gLogger.error(f"{key} is mandatory")
                DIRAC.exit(-1)

    def build_job_attributes_from_input(self) -> None:
        """Set job attributes from input meta query"""
        for key, value in self.prod_step.Inputquery.items():
            if key in self.file_meta_fields:
                try:
                    self.job.output_file_metadata[key] = value["="]
                except (KeyError, TypeError):
                    self.job.output_file_metadata[key] = value
                except Exception as e:
                    print(f"Unexpected exception for key {key}: {e}")
                    raise
            else:
                self.set_attribute(self.job, key, value, "=")

    def set_attribute(self, obj, key, value, sym) -> None:
        try:
            setattr(obj, key, value[f"{sym}"])
        except (KeyError, TypeError):
            setattr(obj, key, value)
        except Exception as e:
            print(f"Unexpected exception for key {key}: {e}")
            raise

    def build_job_attributes(self, workflow_config, workflow_step) -> None:
        """Set job attributes"""
        self.build_job_attributes_from_input()
        self.build_job_attributes_from_common(workflow_config)
        self.build_job_attributes_from_job_config(workflow_step)

    def build_input_data(self, workflow_step) -> None:
        """Build input data"""
        self.prod_step.Inputquery = {}
        if self.we_type not in self.allowed_simulation_types:
            if self.prod_step.ParentStep:
                self.prod_step.Inputquery = self.prod_step.ParentStep.Outputquery
            else:
                if workflow_step["input_meta_query"].get("dataset"):
                    self.prod_step.Inputquery = get_dataset_MQ(
                        workflow_step["input_meta_query"]["dataset"]
                    )

            self.refine_input_MQ(workflow_step)

    def refine_input_MQ(self, workflow_step) -> None:
        """refine input MQ by adding extra specification if needed"""
        for key, value in workflow_step["input_meta_query"].items():
            if value is not None:
                if key in self.constrained_input_keys:
                    self.set_constrained_input_query(key, value)
                elif (key != "dataset") and (key != "parentID"):
                    self.prod_step.Inputquery[key] = value

    def build_job_input_data(self, mode) -> None:
        """Limit the nb of input data to process (for testing purpose)"""
        if self.we_type not in self.allowed_simulation_types:
            res = self.fc.findFilesByMetadata(dict(self.prod_step.Inputquery))
            if not res["OK"]:
                DIRAC.gLogger.error(res["Message"])
                DIRAC.exit(-1)
            input_data_limit = self.job.input_limit
            input_data = res["Value"][:input_data_limit]
            if mode.lower() in ["wms", "local"]:
                self.job.setInputData(input_data)
            if not input_data and mode.lower() == "wms":
                DIRAC.gLogger.error("No job submitted: job must have input data")
                DIRAC.exit(-1)
            if mode.lower() == "ps" and self.job.input_limit:
                f = open("test_input_data.list", "w")
                for lfn in input_data:
                    f.write(lfn + "\n")
                f.close()
                DIRAC.gLogger.notice(
                    "\t\tInput limit found: %d files dumped to test_input_data.list"
                    % len(input_data)
                )
                self.prod_step.Inputquery = {}

    def build_element_config(self) -> None:
        """Set job and production step attributes specific to the configuration"""
        if self.we_type in ["evndispprocessing", "merging"]:
            self.prod_step.GroupSize = self.job.group_size
        self.job.set_executable_sequence(debug=False)
        self.prod_step.Body = self.job.workflow.toXML()

    def build_job_output_data(self, workflow_step) -> None:
        """Build job output meta data"""
        if self.we_type in self.allowed_processing_types:
            metadata = deepcopy(self.prod_step.Inputquery)
            for key, value in workflow_step["job_config"].items():
                metadata[key] = value
            if self.we_type == "merging":
                merged = self.get_merging_level()
                metadata["merged"] = merged
            self.job.set_output_metadata(metadata)
        elif self.we_type in self.allowed_simulation_types:
            self.job.set_output_metadata()

    def get_merging_level(self) -> int:
        """Get the merging level from parent step or from the user query"""
        if self.prod_step.ParentStep:
            merged = self.prod_step.ParentStep.Outputquery["merged"]
        else:
            merged = self.job.merged
        return merged

    def build_output_data(self) -> None:
        """Build output data from the job metadata and the metadata added on the files"""
        self.prod_step.Outputquery = deepcopy(self.job.output_metadata)
        for key, value in self.job.output_file_metadata.items():
            if isinstance(value, list):
                if len(value) > 1:
                    self.prod_step.Outputquery[key] = {"in": value}
                else:
                    self.prod_step.Outputquery[key] = value[0]
            else:
                self.prod_step.Outputquery[key] = value
