#
#  Copyright (C) 2019 Bloomberg Finance LP
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU Lesser General Public
#  License as published by the Free Software Foundation; either
#  version 2 of the License, or (at your option) any later version.
#
#  This library is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.	 See the GNU
#  Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with this library. If not, see <http://www.gnu.org/licenses/>.
#
#  Authors:
#        Angelos Evripiotis <jevripiotis@bloomberg.net>

import os
import datetime
import threading
from contextlib import contextmanager
from typing import Optional, Callable, Iterator, TextIO

from . import _signals
from ._exceptions import BstError
from ._message import Message, MessageType, unconditional_messages
from ._state import State, Task


_RENDER_INTERVAL: datetime.timedelta = datetime.timedelta(seconds=1)


# Time in seconds for which we decide that we want to display subtask information
_DISPLAY_LIMIT: datetime.timedelta = datetime.timedelta(seconds=3)
# If we're in the test suite, we need to ensure that we don't set a limit
if "BST_TEST_SUITE" in os.environ:
    _DISPLAY_LIMIT = datetime.timedelta(seconds=0)


# TimeData class to contain times in an object that can be passed around
# and updated from different places
class _TimeData:
    __slots__ = ["start_time"]

    def __init__(self, start_time: datetime.datetime) -> None:
        self.start_time: datetime.datetime = start_time


class _JobInfo:
    def __init__(self, action_name: str, element_name: str, element_key: str) -> None:
        self.action_name = action_name
        self.element_name = element_name
        self.element_key = element_key


# _MessengerLocal
#
# Thread local storage for the messenger
#
class _MessengerLocal(threading.local):
    def __init__(self) -> None:
        super().__init__()

        # The open file handle for this task
        self.log_handle: Optional[TextIO] = None

        # The filename for this task
        self.log_filename: Optional[str] = None

        # Level of silent messages depth in this task
        self.silence_scope_depth: int = 0

        # Job
        self.job: Optional[_JobInfo] = None


# Messenger()
#
# The messenger object.
#
# This is used to propagate messages either from the main context or
# from task contexts in such a way that messages are propagated to
# the frontend and also optionally recorded to a task log file when
# the message is issued from a task context.
#
class Messenger:
    def __init__(self) -> None:
        self._state: Optional[State] = None  # The State object

        #
        # State related to simple tasks, these drive the status bar
        # when ongoing activities occur outside of an active scheduler
        #
        self._active_simple_tasks: int = 0  # Number of active simple tasks
        self._next_render: Optional[datetime.datetime] = None  # The time of the next render
        self._render_status_cb: Optional[Callable[[], None]] = None  # The render callback

        # Thread local storage
        self._locals: _MessengerLocal = _MessengerLocal()

        # The callback to call when propagating messages
        #
        # FIXME: The message handler is currently not strongly typed,
        #        as it uses a kwarg, we cannot declare it with Callable.
        #        We can use `Protocol` to strongly type this with python >= 3.8
        self._message_handler = None

    def setup_new_action_context(self, action_name: str, element_name: str, element_key: str) -> None:
        self._locals.silence_scope_depth = 0
        self._locals.job = _JobInfo(action_name, element_name, element_key)

    # set_message_handler()
    #
    # Sets the handler for any status messages propagated through
    # the messenger.
    #
    def set_message_handler(self, handler) -> None:
        self._message_handler = handler

    # set_state()
    #
    # Sets the State object within the Messenger
    #
    # Args:
    #    state: The state to set
    #
    def set_state(self, state: State) -> None:
        self._state = state

    # set_render_status_cb()
    #
    # Sets the callback to use to render status
    #
    # Args:
    #    callback: The Callback to be notified
    #
    def set_render_status_cb(self, callback: Callable[[], None]) -> None:
        self._render_status_cb = callback

    # message():
    #
    # Proxies a message back to the caller, this is the central
    # point through which all messages pass.
    #
    # Args:
    #    message: A Message object
    #
    def message(self, message: Message) -> None:
        # If we are recording messages, dump a copy into the open log file.
        self._record_message(message)

        # Always add the log filename automatically
        message.logfile = self._locals.log_filename

        is_silenced = self._silent_messages()
        job = self._locals.job

        if job is not None:
            # Automatically add message information from the job context
            message.action_name = job.action_name
            message.task_element_name = job.element_name
            message.task_element_key = job.element_key

            # Don't forward LOG messages from jobs
            if message.message_type == MessageType.LOG:
                return

            # Don't forward JOB messages if they are currently silent
            if is_silenced and (message.message_type not in unconditional_messages):
                return

        # Send it off to the log handler (can be the frontend,
        # or it can be the child task which will propagate
        # to the frontend)
        assert self._message_handler
        self._message_handler(message, is_silenced=is_silenced)

    # status():
    #
    # A core facing convenience method for issuing STATUS messages
    #
    # Args:
    #    brief: The brief status message
    #    detail: An optional detailed message
    #    kwargs: Additional Message constructor keyword arguments
    #
    def status(self, brief: str, *, detail: Optional[str] = None, **kwargs) -> None:
        message = Message(MessageType.STATUS, brief, detail=detail, **kwargs)
        self.message(message)

    # info():
    #
    # A core facing convenience method for issuing INFO messages
    #
    # Args:
    #    brief: The brief info message
    #    detail: An optional detailed message
    #    kwargs: Additional Message constructor keyword arguments
    #
    def info(self, brief: str, *, detail: Optional[str] = None, **kwargs) -> None:
        message = Message(MessageType.INFO, brief, detail=detail, **kwargs)
        self.message(message)

    # warn():
    #
    # A core facing convenience method for issuing WARN messages
    #
    # Args:
    #    brief: The brief warning message
    #    detail: An optional detailed message
    #    kwargs: Additional Message constructor keyword arguments
    #
    def warn(self, brief: str, *, detail: Optional[str] = None, **kwargs) -> None:
        message = Message(MessageType.WARN, brief, detail=detail, **kwargs)
        self.message(message)

    # error():
    #
    # A core facing convenience method for issuing ERROR messages
    #
    # Args:
    #    brief: The brief error message
    #    detail: An optional detailed message
    #    kwargs: Additional Message constructor keyword arguments
    #
    def error(self, brief: str, *, detail: Optional[str] = None, **kwargs) -> None:
        message = Message(MessageType.ERROR, brief, detail=detail, **kwargs)
        self.message(message)

    # bug():
    #
    # A core facing convenience method for issuing BUG messages
    #
    # Args:
    #    brief: The brief bug message
    #    detail: An optional detailed message
    #    kwargs: Additional Message constructor keyword arguments
    #
    def bug(self, brief: str, *, detail: Optional[str] = None, **kwargs) -> None:
        message = Message(MessageType.BUG, brief, detail=detail, **kwargs)
        self.message(message)

    # silence()
    #
    # A context manager to silence messages, this behaves in
    # the same way as the `silent_nested` argument of the
    # timed_activity() context manager: all but
    # _message.unconditional_messages will be silenced.
    #
    # Args:
    #    actually_silence: Whether to actually do the silencing, if
    #                      False then this context manager does not
    #                      affect anything.
    #
    @contextmanager
    def silence(self, *, actually_silence: bool = True) -> Iterator[None]:
        if not actually_silence:
            yield None
            return

        self._locals.silence_scope_depth += 1
        try:
            yield None
        finally:
            assert self._locals.silence_scope_depth > 0
            self._locals.silence_scope_depth -= 1

    # timed_activity()
    #
    # Context manager for performing timed activities and logging those
    #
    # Args:
    #    activity_name: The name of the activity
    #    detail: An optional detailed message, can be multiline output
    #    silent_nested: If True, all nested messages are silenced except for unconditionaly ones
    #    kwargs: Remaining Message() constructor keyword arguments.
    #
    @contextmanager
    def timed_activity(
        self, activity_name: str, *, detail: str = None, silent_nested: bool = False, **kwargs
    ) -> Iterator[None]:
        with self.timed_suspendable() as timedata:
            try:
                # Push activity depth for status messages
                message = Message(MessageType.START, activity_name, detail=detail, **kwargs)
                self.message(message)
                with self.silence(actually_silence=silent_nested):
                    yield None

            except BstError:
                # Note the failure in status messages and reraise, the scheduler
                # expects an error when there is an error.
                elapsed = datetime.datetime.now() - timedata.start_time
                message = Message(MessageType.FAIL, activity_name, elapsed=elapsed, **kwargs)
                self.message(message)
                raise

            elapsed = datetime.datetime.now() - timedata.start_time
            message = Message(MessageType.SUCCESS, activity_name, elapsed=elapsed, **kwargs)
            self.message(message)

    # simple_task()
    #
    # Context manager for creating a task to report progress to.
    #
    # Args:
    #    activity_name: The name of the activity
    #    task_name: Optionally, the task name for the frontend during this task
    #    detail: An optional detailed message, can be multiline output
    #    silent_nested: If True, all nested messages are silenced except for unconditionaly ones
    #    kwargs: Remaining Message() constructor keyword arguments.
    #
    # Yields:
    #    Task: A Task object that represents this activity, principally used to report progress
    #
    @contextmanager
    def simple_task(
        self, activity_name: str, *, task_name: str = None, detail: str = None, silent_nested: bool = False, **kwargs
    ) -> Iterator[Optional[Task]]:
        # Bypass use of State when none exists (e.g. tests)
        if not self._state:
            with self.timed_activity(activity_name, detail=detail, silent_nested=silent_nested, **kwargs):
                yield None
            return

        if not task_name:
            task_name = activity_name

        with self.timed_suspendable() as timedata:
            try:
                message = Message(MessageType.START, activity_name, detail=detail, **kwargs)
                self.message(message)

                task = self._state.add_task(task_name, activity_name, task_name)
                task.set_task_changed_callback(self._render_status)
                self._active_simple_tasks += 1
                if not self._next_render:
                    self._next_render = datetime.datetime.now() + _RENDER_INTERVAL

                with self.silence(actually_silence=silent_nested):
                    yield task

            except BstError:
                elapsed = datetime.datetime.now() - timedata.start_time
                message = Message(MessageType.FAIL, activity_name, elapsed=elapsed, **kwargs)
                self.message(message)
                raise
            finally:
                self._state.remove_task(task_name)
                self._active_simple_tasks -= 1
                if self._active_simple_tasks == 0:
                    self._next_render = None

            elapsed = datetime.datetime.now() - timedata.start_time
            detail = None

            if task.current_progress is not None and elapsed > _DISPLAY_LIMIT:
                if task.maximum_progress is not None:
                    detail = "{} of {} subtasks processed".format(task.current_progress, task.maximum_progress)
                else:
                    detail = "{} subtasks processed".format(task.current_progress)
            message = Message(MessageType.SUCCESS, activity_name, elapsed=elapsed, detail=detail, **kwargs)
            self.message(message)

    # recorded_messages()
    #
    # Records all messages in a log file while the context manager
    # is active.
    #
    # In addition to automatically writing all messages to the
    # specified logging file, an open file handle for process stdout
    # and stderr will be available via the Messenger.get_log_handle() API,
    # and the full logfile path will be available via the
    # Messenger.get_log_filename() API.
    #
    # Args:
    #    filename: A logging directory relative filename,
    #              the pid and .log extension will be automatically
    #              appended
    #
    #    logdir: The path to the log file directory.
    #
    # Yields:
    #    The fully qualified log filename
    #
    @contextmanager
    def recorded_messages(self, filename: str, logdir: str) -> Iterator[str]:
        # We dont allow recursing in this context manager, and
        # we also do not allow it in the main process.
        assert not hasattr(self._locals, "log_handle") or self._locals.log_handle is None
        assert not hasattr(self._locals, "log_filename") or self._locals.log_filename is None

        # Create the fully qualified logfile in the log directory,
        # appending the pid and .log extension at the end.
        self._locals.log_filename = os.path.join(logdir, "{}.{}.log".format(filename, os.getpid()))

        # Ensure the directory exists first
        directory = os.path.dirname(self._locals.log_filename)
        os.makedirs(directory, exist_ok=True)

        with open(self._locals.log_filename, "a") as logfile:

            # Write one last line to the log and flush it to disk
            def flush_log():

                # If the process currently had something happening in the I/O stack
                # then trying to reenter the I/O stack will fire a runtime error.
                #
                # So just try to flush as well as we can at SIGTERM time
                try:
                    logfile.write("\n\nForcefully terminated\n")
                    logfile.flush()
                except RuntimeError:
                    os.fsync(logfile.fileno())

            self._locals.log_handle = logfile
            with _signals.terminator(flush_log):
                yield self._locals.log_filename

            self._locals.log_handle = None
            self._locals.log_filename = None

    # get_log_handle()
    #
    # Fetches the active log handle, this will return the active
    # log file handle when the Messenger.recorded_messages() context
    # manager is active
    #
    # Returns:
    #    The active logging file handle, or None
    #
    def get_log_handle(self) -> Optional[TextIO]:
        return self._locals.log_handle

    # get_log_filename()
    #
    # Fetches the active log filename, this will return the active
    # log filename when the Messenger.recorded_messages() context
    # manager is active
    #
    # Returns:
    #    The active logging filename, or None
    #
    def get_log_filename(self) -> Optional[str]:
        return self._locals.log_filename

    # timed_suspendable()
    #
    # A contextmanager that allows an activity to be suspended and can
    # adjust for clock drift caused by suspending
    #
    # Yields:
    #    An object that contains the time the activity started
    #
    @contextmanager
    def timed_suspendable(self) -> Iterator[_TimeData]:
        # Note: timedata needs to be in a namedtuple so that values can be
        # yielded that will change
        timedata = _TimeData(start_time=datetime.datetime.now())
        stopped_time = None

        def stop_time():
            nonlocal stopped_time
            stopped_time = datetime.datetime.now()

        def resume_time():
            sleep_time = datetime.datetime.now() - stopped_time
            timedata.start_time += sleep_time

        with _signals.suspendable(stop_time, resume_time):
            yield timedata

    # _silent_messages():
    #
    # Returns:
    #    (bool): Whether messages are currently being silenced
    #
    def _silent_messages(self) -> bool:
        return self._locals.silence_scope_depth > 0

    # _record_message()
    #
    # Records the message if recording is enabled
    #
    # Args:
    #    message: The message to record
    #
    def _record_message(self, message: Message) -> None:

        if self._locals.log_handle is None:
            return

        INDENT = "    "
        EMPTYTIME = "--:--:--"
        template = "[{timecode: <8}] {type: <7}"

        # If this message is associated with an element or source plugin, print the
        # full element name and key for the instance.
        element_key = ""
        if message.element_key:
            template += " [{element_key}]"
            element_key = message.element_key.brief

        element_name = ""
        if message.element_name:
            template += " {element_name}"
            element_name = message.element_name

        template += ": {message}"

        detail = ""
        if message.detail is not None:
            template += "\n\n{detail}"
            detail = message.detail.rstrip("\n")
            detail = INDENT + INDENT.join(detail.splitlines(True))

        timecode = EMPTYTIME
        if message.message_type in (MessageType.SUCCESS, MessageType.FAIL):
            hours, remainder = divmod(int(message.elapsed.total_seconds()), 60 ** 2)
            minutes, seconds = divmod(remainder, 60)
            timecode = "{0:02d}:{1:02d}:{2:02d}".format(hours, minutes, seconds)

        text = template.format(
            timecode=timecode,
            element_key=element_key,
            element_name=element_name,
            type=message.message_type.upper(),
            message=message.message,
            detail=detail,
        )

        # Write to the open log file
        self._locals.log_handle.write("{}\n".format(text))
        self._locals.log_handle.flush()

    # _render_status()
    #
    # Calls the render status callback set in the messenger, but only if a
    # second has passed since it last rendered.
    #
    def _render_status(self) -> None:
        assert self._next_render

        # self._render_status_cb()
        now = datetime.datetime.now()
        if self._render_status_cb and now >= self._next_render:
            self._render_status_cb()
            self._next_render = now + _RENDER_INTERVAL