# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/02_modules_core.ipynb (unless otherwise specified).

__all__ = ['get_auto_module_id', 'ModuleState', 'DharpaModule', 'EmptyObject', 'log', 'find_all_module_classes']

# Cell
import collections
import copy

from dataclasses import dataclass, field

from traitlets import HasTraits, All, Bool, Instance, Dict
from abc import ABCMeta, abstractmethod
import typing
from typing import Mapping, Tuple, Union
from ipywidgets import Widget
from ..utils import get_module_name_from_class, get_subclass_map
import logging


log = logging.getLogger("dharpa-toolbox")

_AUTO_MODULE_ID = {}

def get_auto_module_id(module_cls: typing.Type):

    nr = _AUTO_MODULE_ID.setdefault(module_cls, 0)
    _AUTO_MODULE_ID[module_cls] = nr+1

    name = get_module_name_from_class(module_cls)

    return f"{name}_{nr}"


class ModuleState(HasTraits):

    config = Dict(allow_none=False)

    inputs = Instance(klass=HasTraits, allow_none=False)
    outputs = Instance(klass=HasTraits, allow_none=False)

    stale = Bool(default_value=True)
    busy = Bool(default_value=False)


class DharpaModule(metaclass=ABCMeta):

    def __init__(self, **config: typing.Any):

        _id = config.pop("id", None)
        if _id is None:
            _id = get_auto_module_id(self.__class__)

        if "." in _id:
            raise ValueError("Module id can't contain '.'")
        if not isinstance(_id, str):
            raise TypeError(f"Can't create module, module id must be of type string: {_id}")

        self._id: str = _id
        self._config_raw = None

        # self._input_mapping: Mapping[str, Tuple[str, str]] = {}
        # if input_mapping:
        #     for k, v in input_mapping.items():
        #         if isinstance(v, str):
        #             self._input_mapping[k] = (v, k)
        #         elif isinstance(v, collections.Iterable) and len(v) == 2:
        #             self._input_mapping[k] = (v[0], v[1])
        #         else:
        #             raise ValueError(f"Invalid value for input reference: {v} -- must be a string or 2-item iterable")
        self._inputs_current: typing.Dict[str, typing.Any] = {}
        self._input_staging: Mapping[str, typing.Any] = {}

        self._state: ModuleState = None
        self.set_config(**config)

    # @abstractmethod
    # def _get_inputs_config(self, **config) -> ModuleInputValues:
    #     pass
    #
    # @abstractmethod
    # def _get_outputs_config(self, **config) -> ModuleOutputValues:
    #     pass

    @abstractmethod
    def _process(self, **inputs) -> Mapping[str, typing.Any]:
        pass

    @abstractmethod
    def _create_inputs(self, **config) -> HasTraits:
        pass

    @abstractmethod
    def _create_outputs(self, **config) -> HasTraits:
        pass

    def set_config(self, **config: typing.Any):

        if "id" in config.keys() and config["id"] != self.id:
            raise ValueError(f"Changing module id is not allowed: {self.id} != {config['id']}")

        self._config_raw = config
        processed_config = self._preprocess_config(**self._config_raw)

        self._input_staging = {}
        self._input_current = {}

        if self._state is None:

            self._state: ModuleState = ModuleState(config=processed_config, inputs=self._create_inputs(**processed_config), outputs=self._create_outputs(**processed_config), stale=True, busy=False)
        else:
            self._state.inputs.unobserve_all()
            self._state.outputs.unobserve_all()

            self._state.config = processed_config

            self._state.inputs = self._create_inputs(**self._state.config)
            self._state.outputs = self._create_outputs(**self._state.config)

        self._state.inputs.observe(self._input_updated, names=All)
        self._state.stale = True

    def set_config_value(self, key: str, value: typing.Any):

        if key == "id":
            raise ValueError("Changing module id is not allowed.")


        self._config[key] = value
        self.set_config(self._config)

    def _preprocess_config(self, **config: typing.Any):

        return config

    @property
    def inputs(self) -> HasTraits:

        return self._state.inputs

    def set_input(self, input_name: str, value: typing.Any) -> None:

        print("XXXXXXXXXXXXXXXXXXXXXXXXXXXXX")
        print(self)
        print(self._state.inputs.trait_names())

        self._state.inputs.set_trait(input_name, value)


    @property
    def outputs(self) -> HasTraits:

        return self._state.outputs

    def _input_updated(self, change) -> typing.Any:

        log.debug(f"Input updated for module ({self.id}): {change}")

        if change.name not in self._input_staging.keys():
            self._input_staging[change.name] = {"old": change.old, "new": change.new}
        else:
            self._input_staging[change.name]["new"] = change.new
        self._state.stale = True

    def process(self):

        self._state.busy = True# if not self._stale_state:
        #     return

        try:

            current = {}
            for k in self._state.inputs.trait_names():
                v = getattr(self._state.inputs, k)
                current[k] = v

            result = self._process(**current)

            # print("RESULT")
            # print(result)
            self._inputs_current.clear()
            self._inputs_current.update(current)


            self._input_staging.clear()

            with self.outputs.hold_trait_notifications():
                if result:
                    for k, v in result.items():
                        print(f"{k} = {v}")

                        if k not in self.outputs.trait_names():
                            continue
                        self.outputs.set_trait(k, v)
                self._state.stale = False

        # except Exception as e:
        #     raise
        finally:
            self._state.busy = False

    @property
    def id(self) -> str:
        return self._id

    @property
    def stale(self) -> bool:
        return self._state.stale

    @property
    def current_state(self) -> typing.Dict[str, typing.Any]:

        result = {"inputs": {}, "outputs": {}}
        for tn in self.inputs.trait_names():
            stale = False
            result["inputs"][tn] = {}
            if tn in self._inputs_current.keys():
                result["inputs"][tn]["current_value"] = self._inputs_current[tn]
            else:
                stale = True
                result["inputs"][tn]["current_value"] = "-- not set --"

            if tn in self._input_staging.keys():
                stale = True
                result["inputs"][tn]["new_value"] = self._input_staging[tn]["new"]
            else:
                result["inputs"][tn]["new_value"] = "-- not set --"
            result["inputs"][tn]["stale"] = stale

        for tn in self.outputs.trait_names():
            result["outputs"][tn] = getattr(self.outputs, tn)


        result["stale"] = self._state.stale
        return result

    @property
    def input_names(self) -> typing.Iterable[str]:

        return self._state.inputs.trait_names()

    @property
    def output_names(self) -> typing.Iterable[str]:
        return self._state.outputs.trait_names()

    def __eq__(self, other):

        if self.__class__ != other.__class__:
            return False

        return self.id == other.id

    def __hash__(self):

        return hash(self.id)

    def __repr__(self):

        return f"{self.__class__.__name__}(id='{self.id}' input_names={self.inputs.trait_names()} output_names={self.outputs.trait_names()}) config={self._config_raw}"

    def __str__(self):

        return f"module: '{self.id}'"


class EmptyObject(HasTraits):

    pass


# Cell

def find_all_module_classes():
    modules_to_load = ["dharpa_toolbox.modules.core", "dharpa_toolbox.modules.files", "dharpa_toolbox.modules.text"]

    all_module_clases = get_subclass_map(DharpaModule, preload_modules=modules_to_load, key_func=get_module_name_from_class)
    return all_module_clases

