# =============================================================================
# Copyright (c) 2024 by Laszlo Anka. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
import time
import traceback
from abc import ABC
from typing import Any, Optional, Type

from pypz.core.channels.io import ChannelReader, ChannelWriter
from pypz.core.channels.status import ChannelStatus
from pypz.core.commons.parameters import OptionalParameter, RequiredParameter
from pypz.core.commons.utils import current_time_millis
from pypz.core.specs.plugin import (
    ExtendedPlugin,
    InputPortPlugin,
    OutputPortPlugin,
    ResourceHandlerPlugin,
)
from pypz.executors.operator.states import (
    StateOperationInit,
    StateOperationRunning,
    StateOperationShutdown,
)

ParamKeyChannelLocationConfig = "channelLocation"
ParamKeyChannelConfig = "channelConfig"
ParamKeySequentialModeEnabled = "sequentialModeEnabled"
ParamKeyPortOpenTimeoutMs = "portOpenTimeoutMs"
ParamKeySyncConnectionsOpen = "syncConnectionsOpen"


class ChannelInputPort(InputPortPlugin, ResourceHandlerPlugin, ExtendedPlugin, ABC):
    """
    This class represents an abstract InputPortPlugin, which utilizes the channels
    to realize transfer functionalities. It means that once there is a proper
    channel implementation, a fully functional InputPortPlugin can be made by
    providing that implementation in the abstract new_input_channel() method.

    :param name: name of the instance, if not provided, it will be attempted to deduce from the variable's name
    :param schema: the schema of the port plugin, which will be used to send/retrieve data
    :param group_mode: if set to True, the all the input ports in the group shall receive all messages
    :param channel_reader_type: the type of the channel reader to be used
    """

    # ======================= parameter descriptors =======================

    _channel_location = RequiredParameter(
        str,
        alt_name=ParamKeyChannelLocationConfig,
        description="Location of the channel resource",
    )
    _channel_config = OptionalParameter(
        dict,
        alt_name=ParamKeyChannelConfig,
        description="Configuration of the channel as dictionary",
    )
    _sequential_mode_enabled = OptionalParameter(
        bool,
        alt_name=ParamKeySequentialModeEnabled,
        description="If set to True, then this port will wait with the "
        "processing start until all the connected output "
        "ports are finished",
    )
    _port_open_timeout_ms = OptionalParameter(
        int,
        alt_name=ParamKeyPortOpenTimeoutMs,
        description="Specifies, how long the port shall wait for incoming"
        "connections",
    )
    _sync_connections_open = OptionalParameter(
        bool,
        alt_name=ParamKeySyncConnectionsOpen,
        description="If set to True, the port will wait for every expected"
        "output ports to be connected",
    )

    # ======================= ctor =======================

    def __init__(
        self,
        name: str = None,
        schema: Any = None,
        group_mode: bool = False,
        channel_reader_type: Type[ChannelReader] = None,
        *args,
        **kwargs,
    ):
        super().__init__(name, schema, group_mode, *args, **kwargs)

        self.channel_reader_type: Type[ChannelReader] = channel_reader_type
        """
        Type of the channel reader to be created by this plugin.
        """

        self._channel_reader: Optional[ChannelReader] = None
        """
        The only channel reader maintained by this plugin.
        """

        self._interrupted: bool = False
        """
        Helper flag to register, if the execution has been interrupted.
        """

        self._delete_resources: bool = True
        """
        This flag signalizes, if resource deletion can be performed. It will
        be cleared, if there is an error on OperationRunning state to allow
        process restart.
        """

        self._need_to_check_connections_opened: bool = True
        """
        Helper flag to signalize, if the channel shall be opened. This is necessary
        to maintain state in _on_port_open().
        """

        self._expected_output_count: Optional[int] = None
        """
        During _on_port_open the expected number of connected outputs will be calculated.
        It then will be used in cases, where we need to wait for all outputs.
        """

        # Expected parameters
        # ===================

        self._channel_location = None
        self._channel_config = {}
        self._sequential_mode_enabled = False
        self._port_open_timeout_ms = 0
        self._port_open_timeout_start = 0
        self._sync_connections_open = False

    # ======================= method implementations =======================

    def _pre_execution(self) -> None:
        if self.channel_reader_type is None:
            raise AttributeError(
                f"[{self.get_full_name()}] Channel reader type is not specified"
            )

        if self._channel_reader is None:
            channel_name = (
                self.get_full_name()
                if self.get_group_principal() is None
                else self.get_group_principal().get_full_name()
            )
            self._channel_reader = self.channel_reader_type(channel_name, self)
            self._channel_reader.set_location(self._channel_location)
            self._channel_reader.invoke_configure_channel(self._channel_config)

    def _post_execution(self) -> None:
        pass

    def _on_resource_creation(self) -> bool:
        if self.is_principal() and (not self._channel_reader.is_resource_created()):
            if self._channel_reader.invoke_resource_creation():
                self._channel_reader.get_logger().debug("Resource created.")
            else:
                self._channel_reader.get_logger().debug(
                    "Waiting for resource creation ..."
                )
                return False
        return True

    def _on_resource_deletion(self) -> bool:
        if self.is_principal() and (not self._channel_reader.is_resource_deleted()):
            # Do not delete resources on error
            if not self._delete_resources:
                return True

            if self._channel_reader.invoke_resource_deletion():
                self._channel_reader.get_logger().debug("Resource deleted.")
            else:
                self._channel_reader.get_logger().debug(
                    "Waiting for resource deletion ..."
                )
                return False
        return True

    def commit_current_read_offset(self):
        self._channel_reader.invoke_commit_current_read_offset()

    def _on_port_open(self) -> bool:
        if not self._channel_reader.is_channel_open():
            if self._expected_output_count is None:
                self._expected_output_count = sum(
                    output.get_group_size() for output in self.get_connected_ports()
                )

            if self._channel_reader.invoke_open_channel():
                self._channel_reader.get_logger().debug("Channel opened")

                # Resetting, since we reuse this variable at waiting for all inputs to get connected
                self._port_open_timeout_start = 0
            else:
                if 0 < self._port_open_timeout_ms:
                    if 0 == self._port_open_timeout_start:
                        self._port_open_timeout_start = current_time_millis()
                    elif self._port_open_timeout_ms < (
                        current_time_millis() - self._port_open_timeout_start
                    ):
                        raise TimeoutError(
                            f"Timeout exceeded {self._port_open_timeout_ms} [ms]"
                        )

                self._channel_reader.get_logger().debug(
                    "Waiting for channel to open..."
                )
                return False

        if (
            ChannelReader.NotInitialized
            == self._channel_reader.get_read_record_offset()
        ):
            self._channel_reader.set_initial_record_offset_auto()
            self._channel_reader.get_logger().debug(
                f"Initialized offset: "
                f"{self._channel_reader.get_read_record_offset()}"
            )

        if not self._channel_reader.is_channel_started():
            self._channel_reader.start_channel()

        self._channel_reader.invoke_sync_status_update()

        # Waiting for outputs to be connected
        # ===================================

        if self._need_to_check_connections_opened:
            if 0 < self._expected_output_count:
                opened_output_count = len(
                    self._channel_reader.retrieve_connected_channel_unique_names(
                        lambda flt: (self.get_group_name() is None)
                        or (flt.get_channel_group_name() != self.get_group_name())
                    )
                )

                if (0 == opened_output_count) or (
                    self._sync_connections_open
                    and (self._expected_output_count > opened_output_count)
                ):
                    if 0 < self._port_open_timeout_ms:
                        if 0 == self._port_open_timeout_start:
                            self._port_open_timeout_start = current_time_millis()
                        elif self._port_open_timeout_ms < (
                            current_time_millis() - self._port_open_timeout_start
                        ):
                            raise TimeoutError(
                                f"Timeout exceeded {self._port_open_timeout_ms} [ms]"
                            )
                    return False

            self._need_to_check_connections_opened = False

        # Waiting for outputs to be finished
        # ==================================
        # Note that normally it does not make sense to slow down the execution by making it sequential,
        # since it is used a couple of times for demo/comparison purposes.
        # Notice that at this point all other relevant checks has been made
        # so there is at least 1 output connected and open

        if self._sequential_mode_enabled:
            if (
                self._channel_reader.is_any_connected_channel_healthy_and_not_stopped_and_not_closed()
            ):
                self.get_logger().debug("Waiting for output channel(s) to finish ...")
                return False

        return True

    def retrieve(self) -> Any:
        return self._channel_reader.invoke_read_records()

    def can_retrieve(self) -> bool:
        if not self._channel_reader.is_channel_open():
            raise PermissionError(
                f"Method allowed only, if channel is open ({self.get_full_name()})"
            )

        finished_output_count = len(
            self._channel_reader.retrieve_connected_channel_unique_names(
                lambda flt: (
                    (self.get_group_name() is None)
                    or (flt.get_channel_group_name() != self.get_group_name())
                )
                and (
                    (not flt.is_channel_healthy())
                    or flt.is_channel_stopped()
                    or flt.is_channel_closed()
                )
            )
        )

        return self._channel_reader.has_records() or (
            self._expected_output_count > finished_output_count
        )

    def _on_port_close(self) -> bool:
        if self._channel_reader.is_channel_open():
            if (
                self._channel_reader.can_close()
                and self._channel_reader.invoke_close_channel()
            ):
                self._channel_reader.get_logger().debug("Channel closed.")
            else:
                self._channel_reader.get_logger().debug(
                    "Waiting for channel to close..."
                )
                return False

        if self.is_principal():
            self.get_logger().debug("Closing port gracefully ...")
            time.sleep(5)
        return True

    def _on_interrupt(self, system_signal: int = None) -> None:
        self.get_logger().warning(f"Interrupted by system signal: {system_signal}")
        self._interrupted = True

    def _on_error(self, source: Any, exception: Exception) -> None:
        if issubclass(
            source, (StateOperationInit, StateOperationRunning, StateOperationShutdown)
        ):
            self._delete_resources = False

        if self._channel_reader.is_channel_open():
            self._channel_reader.invoke_sync_send_status_message(ChannelStatus.Error)


class ChannelOutputPort(OutputPortPlugin, ResourceHandlerPlugin, ExtendedPlugin, ABC):
    """
    This class represents an abstract OutputPortPlugin, which utilizes the channels
    to realize transfer functionalities. It means that once there is a proper
    channel implementation, a fully functional OutputPortPlugin can be made by
    providing that implementation in the abstract new_output_channel() method.

    :param name: name of the instance, if not provided, it will be attempted to deduce from the variable's name
    :param schema: the schema of the port plugin, which will be used to send/retrieve data
    :param channel_writer_type: the type of the channel writer to be used
    """

    # ======================= parameter descriptors =======================

    _channel_location = RequiredParameter(
        str,
        alt_name=ParamKeyChannelLocationConfig,
        description="Location of the channel resource",
    )
    _channel_config = OptionalParameter(
        dict,
        alt_name=ParamKeyChannelConfig,
        description="Configuration of the channel as dictionary",
    )
    _port_open_timeout_ms = OptionalParameter(
        int,
        alt_name=ParamKeyPortOpenTimeoutMs,
        description="Specifies, how long the port shall wait for incoming"
        "connections",
    )

    # ======================= ctor =======================

    def __init__(
        self,
        name: str = None,
        schema: Any = None,
        channel_writer_type: Type[ChannelWriter] = None,
        *args,
        **kwargs,
    ):
        super().__init__(name, schema, *args, **kwargs)

        self.channel_writer_type: Type[ChannelWriter] = channel_writer_type
        """
        Type of the channel writer to be created by this plugin.
        """

        self._channel_writers: Optional[set[ChannelWriter]] = None
        """
        The channel writers to be maintained by this plugin. Note that a channel
        writer will be created for each connection from this output port plugin.
        """

        self._interrupted: bool = False
        """
        Helper flag to register, if the execution has been interrupted.
        """

        self._delete_resources: bool = True
        """
        This flag signalizes, if resource deletion can be performed. It will
        be cleared, if there is an error on OperationRunning state to allow
        process restart.
        """

        self._resource_deletion_errors: set[ChannelWriter] = set()
        """
        Set to store the channels that had errors on resource deletion. It is necessary
        to be able not to reschedule those channels in re-execution of the state.
        """

        self._port_close_errors: set[ChannelWriter] = set()
        """
        Set to store the channels that had errors on port closure. It is necessary
        to be able not to reschedule those channels in re-execution of the state.
        """

        # Expected parameters
        # ===================

        self._channel_location = None
        self._channel_config = {}
        self._port_open_timeout_ms = 0
        self._port_open_timeout_start = 0

    # ======================= method implementations =======================

    def _pre_execution(self) -> None:
        if self.channel_writer_type is None:
            raise AttributeError(
                f"[{self.get_full_name()}] Channel writer type is not specified"
            )

        if self._channel_writers is None:
            self._channel_writers = set()
            for connection in self.get_connected_ports():
                new_output_channel = self.channel_writer_type(
                    connection.get_full_name(), self
                )
                new_output_channel.set_location(self._channel_location)
                new_output_channel.invoke_configure_channel(self._channel_config)
                self._channel_writers.add(new_output_channel)

    def _post_execution(self) -> None:
        pass

    def _on_resource_creation(self) -> bool:
        if not self.is_principal():
            return True

        all_resources_created = True

        for channel in self._channel_writers:
            if not channel.is_resource_created():
                if channel.invoke_resource_creation():
                    channel.get_logger().debug("Resource created.")
                else:
                    channel.get_logger().debug("Waiting for resource creation ...")
                    all_resources_created = False

        if not all_resources_created:
            return False

        if 0 == len(self._channel_writers):
            self.get_logger().warning("No channel reader(s) connected")

        return True

    def _on_resource_deletion(self) -> bool:
        if not self.is_principal():
            return True

        all_resources_deleted = True

        for channel in self._channel_writers:
            if (not channel.is_resource_deleted()) and (
                channel not in self._resource_deletion_errors
            ):
                try:
                    if self._delete_resources:
                        if channel.invoke_resource_deletion():
                            channel.get_logger().debug("Resource deleted.")
                        else:
                            channel.get_logger().debug(
                                "Waiting for resource deletion ..."
                            )
                            all_resources_deleted = False
                except:  # noqa: E722
                    # We catch the fact that error occurred, but not interrupting the deletion
                    # process of the other channels
                    self.get_logger().error(traceback.format_exc())
                    self._resource_deletion_errors.add(channel)

        if not all_resources_deleted:
            return False

        if 0 < len(self._resource_deletion_errors):
            raise InterruptedError(
                f"Error occurred at resource deletion in the following channels: "
                f"{[channel.get_channel_name() for channel in self._resource_deletion_errors]}"
            )

        return True

    def _on_port_open(self) -> bool:
        all_channels_opened = True

        for channel in self._channel_writers:
            if not channel.is_channel_open():
                if channel.invoke_open_channel():
                    channel.get_logger().debug("Channel opened.")
                else:
                    channel.get_logger().debug("Waiting for channel open ...")
                    all_channels_opened = False

        if not all_channels_opened:
            if 0 < self._port_open_timeout_ms:
                if 0 == self._port_open_timeout_start:
                    self._port_open_timeout_start = current_time_millis()
                elif self._port_open_timeout_ms < (
                    current_time_millis() - self._port_open_timeout_start
                ):
                    raise TimeoutError(
                        f"Timeout exceeded {self._port_open_timeout_ms} [ms]"
                    )

            return False

        for channel in self._channel_writers:
            if not channel.is_channel_started():
                channel.start_channel()

        return True

    def _on_port_close(self) -> bool:
        all_channels_closed = True

        for channel in self._channel_writers:
            if channel.is_channel_open() and (channel not in self._port_close_errors):
                try:
                    if channel.can_close() and channel.invoke_close_channel():
                        channel.get_logger().debug("Channel closed.")
                    else:
                        all_channels_closed = False
                        channel.get_logger().debug("Waiting for channel to close...")
                except:  # noqa: E722
                    # We catch the fact that error occurred, but not interrupting the shutdown
                    # process of the other channels
                    self.get_logger().error(traceback.format_exc())
                    self._port_close_errors.add(channel)

        if not all_channels_closed:
            return False

        if 0 < len(self._port_close_errors):
            raise InterruptedError(
                f"Error occurred at port close in the following channels: "
                f"{[channel.get_channel_name() for channel in self._port_close_errors]}"
            )

        if self.is_principal():
            self.get_logger().debug("Closing port gracefully ...")
            time.sleep(5)

        return True

    def send(self, data: Any) -> Any:
        for channel in self._channel_writers:
            channel.invoke_write_records(data)

    def _on_interrupt(self, system_signal: int = None) -> None:
        self.get_logger().warning(f"Interrupted by system signal: {system_signal}")
        self._interrupted = True

    def _on_error(self, source: Any, exception: Exception) -> None:
        if issubclass(
            source, (StateOperationInit, StateOperationRunning, StateOperationShutdown)
        ):
            self._delete_resources = False

        error_channels: set[ChannelWriter] = set()
        for channel in self._channel_writers:
            try:
                if channel.is_channel_open():
                    channel.invoke_sync_send_status_message(ChannelStatus.Error)
            except:  # noqa: E722
                # We catch the fact that error occurred, but not interrupting the error handling
                # process of the other channels
                self.get_logger().error(traceback.format_exc())
                error_channels.add(channel)

        if 0 < len(error_channels):
            raise InterruptedError(
                f"Error occurred at error handling in the following channels: "
                f"{[channel.get_channel_name() for channel in error_channels]}"
            )
