import logging
from functools import partial
from logging import Logger
from typing import Callable, Union

from pydantic import BaseModel

from dp3.common.attrspec import AttrType
from dp3.common.config import ModelSpec, PlatformConfig, read_config_dir
from dp3.common.datapoint import DataPointType
from dp3.common.scheduler import Scheduler
from dp3.common.state import SharedFlag
from dp3.common.task import DataPointTask
from dp3.snapshots.snapshooter import SnapShooter
from dp3.task_processing.task_executor import TaskExecutor


def write_datapoints_into_record(model_spec: ModelSpec, tasks: list[DataPointTask], record: dict):
    eid = record["eid"]
    for task in tasks:
        if task.eid != eid:
            continue
        for dp in task.data_points:
            attr_spec = model_spec.attributes[dp.etype, dp.attr]
            if attr_spec.t == AttrType.PLAIN:
                if attr_spec.t in AttrType.PLAIN | AttrType.OBSERVATIONS and attr_spec.is_iterable:
                    v = [
                        elem.model_dump() if isinstance(elem, BaseModel) else elem for elem in dp.v
                    ]
                else:
                    v = dp.v.model_dump() if isinstance(dp.v, BaseModel) else dp.v

                record[dp.attr] = v


def on_entity_creation_in_snapshots(
    model_spec: ModelSpec,
    run_flag: SharedFlag,
    original_hook: Callable[[str, DataPointTask], list[DataPointTask]],
    etype: str,
    record: dict,
) -> list[DataPointTask]:
    """Wrapper for on_entity_creation hooks to enable running as a snapshot callback."""
    if not run_flag.isset():
        return []
    eid = record["eid"]
    mock_task = DataPointTask(etype=etype, eid=eid, data_points=[])
    tasks = original_hook(eid, mock_task)
    write_datapoints_into_record(model_spec, tasks, record)
    return tasks


def on_attr_change_in_snapshots(
    model_spec: ModelSpec,
    run_flag: SharedFlag,
    original_hook: Callable[[str, DataPointTask], Union[list[DataPointTask], None]],
    etype: str,
    record: dict,
) -> list[DataPointTask]:
    """Wrapper for on_entity_creation hooks to enable running as a snapshot callback."""
    if not run_flag.isset():
        return []
    eid = record["eid"]
    mock_task = DataPointTask(etype=etype, eid=eid, data_points=[])
    tasks = original_hook(eid, mock_task)
    if isinstance(tasks, list):
        write_datapoints_into_record(model_spec, tasks, record)
    return tasks


def unset_flag(flag: SharedFlag) -> list:
    flag.unset()
    return []


def reload_module_config(
    log: Logger, platform_config: PlatformConfig, modules: dict, module: str
) -> None:
    """
    Reloads configuration of a module.

    Args:
        log: log to write messages to
        platform_config: Platform configuration
        modules: Dictionary of loaded modules by their names
        module: Name of the module to reload

    Returns:
        Module's configuration
    """
    log.debug(f"Reloading config of module '{module}'")

    if module not in modules:
        log.warning(f"Could not find module '{module}', cannot reload config. Is it loaded?")
        return

    config = read_config_dir(platform_config.config_base_path, recursive=True)
    module_config = config.get(f"modules.{module}", {})
    modules[module].load_config(platform_config, module_config)
    modules[module].refresh.set()
    log.info(f"Config of module '{module}' reloaded successfully.")


class CallbackRegistrar:
    """Interface for callback registration."""

    def __init__(
        self, scheduler: Scheduler, task_executor: TaskExecutor, snap_shooter: SnapShooter
    ):
        self._scheduler = scheduler
        self._task_executor = task_executor
        self._snap_shooter = snap_shooter

        self.log = logging.getLogger(self.__class__.__name__)
        self.model_spec = task_executor.model_spec
        self.attr_spec_t_to_on_attr = {
            AttrType.PLAIN: "on_new_plain",
            AttrType.OBSERVATIONS: "on_new_observation",
            AttrType.TIMESERIES: "on_new_ts_chunk",
        }

    def scheduler_register(
        self,
        func: Callable,
        *,
        func_args: Union[list, tuple] = None,
        func_kwargs: dict = None,
        year: Union[int, str] = None,
        month: Union[int, str] = None,
        day: Union[int, str] = None,
        week: Union[int, str] = None,
        day_of_week: Union[int, str] = None,
        hour: Union[int, str] = None,
        minute: Union[int, str] = None,
        second: Union[int, str] = None,
        timezone: str = "UTC",
    ) -> int:
        """
        Register a function to be run at specified times.

        Pass cron-like specification of when the function should be called,
        see [docs](https://apscheduler.readthedocs.io/en/latest/modules/triggers/cron.html)
        of apscheduler.triggers.cron for details.

        Args:
            func: function or method to be called
            func_args: list of positional arguments to call func with
            func_kwargs: dict of keyword arguments to call func with
            year: 4-digit year
            month: month (1-12)
            day: day of month (1-31)
            week: ISO week (1-53)
            day_of_week: number or name of weekday (0-6 or mon,tue,wed,thu,fri,sat,sun)
            hour: hour (0-23)
            minute: minute (0-59)
            second: second (0-59)
            timezone: Timezone for time specification (default is UTC).
        Returns:
             job ID
        """
        return self._scheduler.register(
            func,
            func_args=func_args,
            func_kwargs=func_kwargs,
            year=year,
            month=month,
            day=day,
            week=week,
            day_of_week=day_of_week,
            hour=hour,
            minute=minute,
            second=second,
            timezone=timezone,
        )

    def register_task_hook(self, hook_type: str, hook: Callable):
        """Registers one of available task hooks

        See: [`TaskGenericHooksContainer`][dp3.task_processing.task_hooks.TaskGenericHooksContainer]
        in `task_hooks.py`
        """
        self._task_executor.register_task_hook(hook_type, hook)

    def register_allow_entity_creation_hook(
        self,
        hook: Callable[[str, DataPointTask], bool],
        entity: str,
    ):
        """
        Registers passed hook to allow entity creation.

        Binds hook to specified entity (though same hook can be bound multiple times).

        Args:
            hook: `hook` callable should expect eid and Task as arguments and return a bool.
            entity: specifies entity type
        """
        self._task_executor.register_entity_hook("allow_entity_creation", hook, entity)

    def register_on_entity_creation_hook(
        self,
        hook: Callable[[str, DataPointTask], list[DataPointTask]],
        entity: str,
        refresh: SharedFlag = None,
        may_change: list[list[str]] = None,
    ):
        """
        Registers passed hook to be called on entity creation.

        Binds hook to specified entity (though same hook can be bound multiple times).

        Allows registration of refreshing on configuration changes, if `refresh` is specified.
        In that case, `may_change` must be specified.

        Args:
            hook: `hook` callable should expect eid and Task as arguments and
                return a list of DataPointTask objects to perform.
            entity: specifies entity type
            refresh: If specified, registered hook will be called on configuration changes.
                Pass `self.refresh` from `BaseModule` subclasses.
            may_change: each item should specify an attribute that `hook` may change,
                for specification format see `register_correlation_hook`
        """
        self._task_executor.register_entity_hook("on_entity_creation", hook, entity)
        if refresh is not None:
            if may_change is None:
                raise ValueError("'may_change' must be specified if 'refresh' is specified")
            self._snap_shooter.register_correlation_hook(
                partial(on_entity_creation_in_snapshots, self.model_spec, refresh, hook),
                entity,
                [],
                may_change,
            )
            self._snap_shooter.register_run_finalize_hook(partial(unset_flag, refresh))

    def register_entity_hook(self, hook_type: str, hook: Callable, entity: str):
        """Registers one of available task entity hooks

        !!! warning "Deprecated"

            This method is deprecated, use `register_on_entity_creation_hook`
            or `register_allow_entity_creation_hook` instead.

        See: [`TaskEntityHooksContainer`][dp3.task_processing.task_hooks.TaskEntityHooksContainer]
        in `task_hooks.py`
        """
        self._task_executor.register_entity_hook(hook_type, hook, entity)

    def register_on_new_attr_hook(
        self,
        hook: Callable[[str, DataPointType], Union[None, list[DataPointTask]]],
        entity: str,
        attr: str,
        refresh: SharedFlag = None,
        may_change: list[list[str]] = None,
    ):
        """
        Registers passed hook to be called on new attribute datapoint.

        Args:
            hook: `hook` callable should expect eid, Task and attribute value as arguments.
                Can optionally return a list of DataPointTasks to perform.
            entity: specifies entity type
            attr: specifies attribute name
            refresh: If specified, registered hook will be called on configuration changes.
                Pass `self.refresh` from `BaseModule` subclasses.
            may_change: each item should specify an attribute that `hook` may change,
                for specification format see `register_correlation_hook`

        Raises:
            ValueError: If entity and attr do not specify a valid attribute, a ValueError is raised.
        """
        try:
            hook_type = self.attr_spec_t_to_on_attr[self.model_spec.attributes[entity, attr].t]
        except KeyError as e:
            raise ValueError(
                f"Cannot register hook for attribute {entity}/{attr}, are you sure it exists?"
            ) from e
        self._task_executor.register_attr_hook(hook_type, hook, entity, attr)

        if refresh is None:
            return
        if may_change is None:
            raise ValueError("'may_change' must be specified if 'refresh' is specified")

        self._snap_shooter.register_correlation_hook(
            partial(on_attr_change_in_snapshots, self.model_spec, refresh, hook),
            entity,
            [[attr]],
            may_change,
        )
        self._snap_shooter.register_run_finalize_hook(partial(unset_flag, refresh))

    def register_attr_hook(self, hook_type: str, hook: Callable, entity: str, attr: str):
        """
        Registers one of available task attribute hooks

        !!! warning "Deprecated"

            This method is deprecated, use `register_on_new_attr_hook` instead.

        See: [`TaskAttrHooksContainer`][dp3.task_processing.task_hooks.TaskAttrHooksContainer]
        in `task_hooks.py`
        """
        self._task_executor.register_attr_hook(hook_type, hook, entity, attr)

    def register_timeseries_hook(
        self,
        hook: Callable[[str, str, list[dict]], list[DataPointTask]],
        entity_type: str,
        attr_type: str,
    ):
        """
        Registers passed timeseries hook to be called during snapshot creation.

        Binds hook to specified `entity_type` and `attr_type` (though same hook can be bound
        multiple times).

        Args:
            hook: `hook` callable should expect entity_type, attr_type and attribute
                history as arguments and return a list of `DataPointTask` objects.
            entity_type: specifies entity type
            attr_type: specifies attribute type

        Raises:
            ValueError: If entity_type and attr_type do not specify a valid timeseries attribute,
                a ValueError is raised.
        """
        self._snap_shooter.register_timeseries_hook(hook, entity_type, attr_type)

    def register_correlation_hook(
        self,
        hook: Callable[[str, dict], Union[None, list[DataPointTask]]],
        entity_type: str,
        depends_on: list[list[str]],
        may_change: list[list[str]],
    ):
        """
        Registers passed hook to be called during snapshot creation.

        Binds hook to specified entity_type (though same hook can be bound multiple times).

        `entity_type` and attribute specifications are validated, `ValueError` is raised on failure.

        Args:
            hook: `hook` callable should expect entity type as str
                and its current values, including linked entities, as dict.
                Can optionally return a list of DataPointTask objects to perform.
            entity_type: specifies entity type
            depends_on: each item should specify an attribute that is depended on
                in the form of a path from the specified entity_type to individual attributes
                (even on linked entities).
            may_change: each item should specify an attribute that `hook` may change.
                specification format is identical to `depends_on`.

        Raises:
            ValueError: On failure of specification validation.
        """
        self._snap_shooter.register_correlation_hook(hook, entity_type, depends_on, may_change)

    def register_snapshot_init_hook(self, hook: Callable[[], list[DataPointTask]]):
        """
        Registers passed hook to be called before a run of snapshot creation begins.

        Args:
            hook: `hook` callable should expect no arguments and
                return a list of DataPointTask objects to perform.
        """
        self._snap_shooter.register_run_init_hook(hook)

    def register_snapshot_finalize_hook(self, hook: Callable[[], list[DataPointTask]]):
        """
        Registers passed hook to be called after a run of snapshot creation ends.

        Args:
            hook: `hook` callable should expect no arguments and
                return a list of DataPointTask objects to perform.
        """
        self._snap_shooter.register_run_finalize_hook(hook)
