"""Core pyout interface definitions.
"""

import abc
from collections import defaultdict
from collections import OrderedDict
from collections.abc import Mapping
import concurrent.futures as cfut
from concurrent.futures import ThreadPoolExecutor as Pool
from contextlib import contextmanager
from functools import wraps
import inspect
from itertools import chain
from logging import getLogger
import os
import sys
import threading
import time

from pyout.common import ContentWithSummary
from pyout.common import RowNormalizer
from pyout.common import StyleFields
from pyout.common import UnknownColumns
from pyout.field import PlainProcessors

lgr = getLogger(__name__)


class Stream(object, metaclass=abc.ABCMeta):
    """Output stream interface used by Writer.

    Parameters
    ----------
    stream : stream, optional
        Stream to write output to.  Defaults to sys.stdout.
    interactive : boolean, optional
        Whether this stream is interactive.  If not specified, it will be set
        to the return value of `stream.isatty()`.

    Attributes
    ----------
    interactive : bool
    supports_updates : boolean
        If true, the writer supports updating previous lines.
    """

    supports_updates = True

    def __init__(self, stream=None, interactive=None):
        self.stream = stream or sys.stdout
        if interactive is None:
            self.interactive = self.stream.isatty()
        else:
            self.interactive = interactive
        self.supports_updates = self.interactive

    @abc.abstractproperty
    def width(self):
        """Maximum line width.
        """

    @abc.abstractproperty
    def height(self):
        """Maximum number of rows that are visible."""

    @abc.abstractmethod
    def write(self, text):
        """Write `text`.
        """

    @abc.abstractmethod
    def clear_last_lines(self, n):
        """Clear previous N lines.
        """

    @abc.abstractmethod
    def overwrite_line(self, n, text):
        """Go to the Nth previous line and overwrite it with `text`
        """

    @abc.abstractmethod
    def move_to(self, n):
        """Move the Nth previous line.
        """


def skip_if_aborted(method):
    """Decorate Writer `method` to prevent execution if write has been aborted.
    """

    @wraps(method)
    def wrapped(self, *args, **kwds):
        if self._aborted:
            lgr.debug("Write has been aborted; not calling %r", method)
        else:
            return method(self, *args, **kwds)
    return wrapped


class Writer(object):
    """Base class implementing the core handling logic of pyout output.

    To define a writer, a subclass should inherit Writer and define __init__ to
     call Writer.__init__ and then the _init method.
    """
    def __init__(self, columns=None, style=None, stream=None,
                 interactive=None, mode=None, continue_on_failure=True,
                 wait_for_top=3, max_workers=None):
        if columns and not isinstance(columns, (list, OrderedDict)):
            self._columns = list(columns)
        else:
            self._columns = columns or None

        self._ids = None

        self._last_content_len = 0
        self._last_summary = None
        self._normalizer = None

        self._pool = None
        if max_workers is None and sys.version_info < (3, 8):
            # ThreadPoolExecutor's max_workers didn't get a default until
            # Python 3.5, and that default was changed in 3.8.  Use Python
            # 3.8's default for consistent behavior.
            max_workers = min(32, (os.cpu_count() or 1) + 4)
        self._max_workers = max_workers
        self._lock = None
        self._aborted = False
        self._futures = defaultdict(list)
        self._continue_on_failure = continue_on_failure

        self._wait_for_top = wait_for_top
        self._mode = mode
        self._write_fn = None

        self._stream = None
        self._width_from_stream = False
        self._content = None

    def _init(self, style, streamer, processors=None):
        """Do writer-specific setup.

        Parameters
        ----------
        style : dict
            Style, as passed to __init__.
        streamer : interface.Stream
            A stream interface that takes __init__'s `stream` and `interactive`
            arguments into account.
        processors : field.StyleProcessors, optional
            A writer-specific processors instance.  Defaults to
            field.PlainProcessors().
        """
        self._stream = streamer
        self._init_mode(streamer)

        style = style or {}
        if style.get("width_") is None:
            lgr.debug("Setting width to stream width: %s",
                      self._stream.width)
            style["width_"] = self._stream.width
            self._width_from_stream = True
        self._content = ContentWithSummary(
            StyleFields(style, processors or PlainProcessors()))

    def _init_prewrite(self, table_width=None):
        self._content.init_columns(self._columns, self.ids,
                                   table_width=table_width)
        self._normalizer = RowNormalizer(self._columns,
                                         self._content.fields.style)

    def _init_mode(self, streamer):
        value = self._mode
        lgr.debug("Initializing mode with given value of %s", value)
        if value is None:
            if streamer.interactive:
                if streamer.supports_updates:
                    value = "update"
                else:
                    value = "incremental"
            else:
                value = "final"
        valid = {"update", "incremental", "final"}
        if value not in valid:
            raise ValueError("{!r} is not a valid mode: {!r}"
                             .format(value, valid))

        lgr.debug("Setting write mode to %r", value)
        self._mode = value
        if value == "incremental":
            self._write_fn = self._write_incremental
        elif value == "final":
            self._write_fn = self._write_final
        else:
            if self._stream.supports_updates and self._stream.interactive:
                self._write_fn = self._write_update
            else:
                raise ValueError("Stream {} does not support updates"
                                 .format(self._stream))

    def __enter__(self):
        return self

    def __exit__(self, _exc_type, exc_value, _tb):
        failed = None
        if exc_value is not None:
            self._abort(msg="\n{!r} raised\n".format(exc_value))
        else:
            try:
                failed = self.wait()
            except KeyboardInterrupt:
                lgr.debug("Caught KeyboardInterrupt "
                          "while waiting for asynchronous workers")
                self._abort(msg="\nKeyboard interrupt registered\n")
                # Raise so that caller can decide how to handle.
                raise

        if self._mode == "final":
            self._stream.write(str(self._content))
        if self._mode != "update" and self._last_summary is not None:
            self._stream.write(str(self._last_summary))

        if failed:
            self._print_async_exceptions(failed)

    @property
    def ids(self):
        """A list of unique IDs used to identify a row.

        If not explicitly set, it defaults to the first column name.
        """
        if self._ids is None:
            if self._columns:
                if isinstance(self._columns, OrderedDict):
                    return [list(self._columns.keys())[0]]
                return [self._columns[0]]
        else:
            return self._ids

    @ids.setter
    def ids(self, columns):
        self._ids = columns

    def _process_futures(self):
        """Process each future as it completes.

        If _continue_on_failure is false, raise the exception of the first
        failed future encountered.  Otherwise return a list of futures that had
        an exception.
        """
        failed = []
        lgr.debug("Waiting for asynchronous calls")
        continue_on_failure = self._continue_on_failure
        for id_key, futures in self._futures.items():
            for future in cfut.as_completed(futures):
                lgr.debug("Processing future %s", future)
                if not future.cancelled() and future.exception():
                    if continue_on_failure:
                        failed.append((id_key, future))
                    else:
                        future.result()  # Raise exception.
        return failed

    def _print_async_exceptions(self, failed_futures):
        import traceback

        # Prevent any remaining callbacks from writing to stream.
        with self._write_lock():
            self._aborted = True

        n_failed = len(failed_futures)
        stream = self._stream
        with self._write_lock():
            stream.write("\n\n")
            stream.write("ERROR: {} asynchronous worker{} failed\n\n"
                         .format(n_failed, "" if n_failed == 1 else "s"))
            for id_key, future in failed_futures:
                try:
                    future.result()
                except Exception:
                    stream.write(
                        "Producing value for row {} failed:\n{}\n"
                        .format(id_key, traceback.format_exc()))

    @skip_if_aborted
    def _abort(self, cause=None, msg=None):
        if self._pool is None:
            # No asynchronous calls; there's nothing to abort.
            return

        with self._write_lock():
            self._aborted = cause or True
            stream = self._stream
            if msg:
                stream.write(msg)
            futures = list(chain(*self._futures.values()))
            for f in futures:
                lgr.debug("Calling .cancel() with for %s", f)
                f.cancel()
        n_running = len([f for f in futures if f.running()])
        stream.write("Canceled pending asynchronous workers. "
                     "{} worker{} already running\n"
                     .format(n_running, "" if n_running == 1 else "s"))
        # Note: We can't call shutdown() with wait=True here.  That will
        # trigger a RuntimeError in underlying <thread>.join() call.
        self._pool.shutdown(wait=False)

    def wait(self):
        """Wait for asynchronous calls to return.

        Returns
        -------
        A list of futures for asynchronous calls had an exception.
        """
        lgr.debug("Waiting for asynchronous calls")
        if self._pool is None:
            return
        aborted = self._aborted
        if aborted:
            if isinstance(aborted, cfut.Future):
                aborted.result()  # Raise exception.
        else:
            failed = self._process_futures()
            self._pool.shutdown(wait=True)
            lgr.debug("Pool shut down")
            return failed

    @contextmanager
    def _write_lock(self):
        """Acquire and release the lock around output calls.

        This should allow multiple threads or processes to write output
        reliably.  Code that modifies the `_content` attribute should also do
        so within this context.
        """
        if self._lock:
            lgr.debug("Acquiring write lock")
            self._lock.acquire()
        try:
            yield
        finally:
            if self._lock:
                lgr.debug("Releasing write lock")
                self._lock.release()

    def _write(self, row, style=None):
        with self._write_lock():
            if self._width_from_stream and self._mode != "final":
                width_current = self._content.fields.style["width_"]
                width_stream = self._stream.width
                if width_stream is not None and width_current != width_stream:
                    lgr.debug("Current stream width (%d) different "
                              "than last recorded (%d). Updating",
                              width_stream, width_current)
                    self._init_prewrite(table_width=width_stream)
            try:
                self._write_fn(row, style)
            except UnknownColumns as exc:
                self._columns.extend(exc.unknown_columns)
                self._init_prewrite()
                self._write_fn(row, style)

    def _get_last_summary_length(self):
        last_summary = self._last_summary
        return len(last_summary.splitlines()) if last_summary else 0

    def _write_update(self, row, style=None):
        last_summary_len = self._get_last_summary_length()

        content, status, summary = self._content.update(row, style)

        if last_summary_len > 0:
            # Clear the summary because 1) it has very likely changed, 2)
            # it makes the counting for row updates simpler, 3) and it is
            # possible for the summary lines to shrink.
            lgr.debug("Clearing summary of %d line(s)", last_summary_len)
            self._stream.clear_last_lines(last_summary_len)


        single_row_updated = False
        if isinstance(status, int):
            height = self._stream.height
            n_visible = min(
                height - last_summary_len - 1,  # -1 for current line.
                self._last_content_len)

            n_back = self._last_content_len - status
            if n_back > n_visible:
                lgr.debug("Cannot move back %d rows for update; "
                          "only %d visible rows",
                          n_back, n_visible)
                status = "repaint"
                content = str(self._content)
            else:
                lgr.debug("Moving up %d line(s) to overwrite line %d with %r",
                          n_back, status, row)
                self._stream.overwrite_line(n_back, content)
                single_row_updated = True

        if not single_row_updated:
            if status == "repaint":
                lgr.debug("Moving up %d line(s) to repaint the whole thing. "
                          "Blame row %r",
                          self._last_content_len, row)
                self._stream.move_to(self._last_content_len)
            self._stream.write(content)

        if summary is not None:
            self._stream.write(summary)
            lgr.debug("Wrote summary")
        self._last_content_len = len(self._content)
        self._last_summary = summary

    def _write_incremental(self, row, style=None):
        content, status, summary = self._content.update(row, style)
        if isinstance(status, int):
            lgr.debug("Duplicating line %d with %r", status, row)
        elif status == "repaint":
            lgr.debug("Duplicating the whole thing.  Blame row %r", row)
        self._stream.write(content)
        self._last_summary = summary

    def _write_final(self, row, style=None):
        _, _, summary = self._content.update(row, style)
        self._last_summary = summary

    @skip_if_aborted
    def _write_async_result(self, id_vals, cols, result):
        lgr.debug("Received result for %s: %s",
                  cols, result)
        if isinstance(result, Mapping):
            lgr.debug("Processing result as mapping")
        elif isinstance(result, tuple):
            lgr.debug("Processing result as tuple")
            result = dict(zip(cols, result))
        elif len(cols) == 1:
            lgr.debug("Processing result as atom")
            result = {cols[0]: result}
        else:
            raise ValueError(
                "Expected tuple or mapping for columns {!r}, got {!r}"
                .format(cols, result))
        result.update(id_vals)
        self._write(result)

    @skip_if_aborted
    def _start_callables(self, row, callables):
        """Start running `callables` asynchronously.
        """
        id_key = tuple(row[c] for c in self.ids)
        id_vals = {c: row[c] for c in self.ids}

        if self._pool is None:
            lgr.debug("Initializing pool with max workers=%s",
                      self._max_workers)
            self._pool = Pool(max_workers=self._max_workers)
        if self._lock is None:
            lgr.debug("Initializing lock")
            self._lock = threading.Lock()

        for cols, fn in callables:
            gen = None
            if inspect.isgeneratorfunction(fn):
                gen = fn()
            elif inspect.isgenerator(fn):
                gen = fn

            def check_result(future):
                if future.cancelled():
                    ok = False
                elif future.exception():
                    ok = False
                    if not self._continue_on_failure:
                        self._abort(cause=future)
                else:
                    ok = True
                return ok

            if gen:
                lgr.debug("Wrapping generator for cols %r of row %r",
                          cols, id_vals)

                def async_fn():
                    for i in gen:
                        self._write_async_result(id_vals, cols, i)

                callback = check_result
            else:
                async_fn = fn

                def callback(future):
                    if check_result(future):
                        self._write_async_result(
                            id_vals, cols, future.result())

            try:
                future = self._pool.submit(async_fn)
            except RuntimeError as exc:
                # We can get here if, between entering this method call and
                # calling .submit(), _aborted was set by a callback.
                if self._aborted:
                    lgr.debug(
                        "Submitting callable for %s failed "
                        "because pool is already shutdown: %s",
                        id_key, exc)
                else:
                    raise
            else:
                future.add_done_callback(callback)
                lgr.debug("Registering future %s for %s", future, id_key)
                self._futures[id_key].append(future)

    def top_nrows_done(self, n, height=None):
        """Check if the top N rows' asynchronous workers are done.

        Parameters
        ----------
        n : int
            Consider this many of the top rows (e.g., 1 would consider just the
            first row).
        height : int, optional
            Take this as the terminal height instead of querying the stream.

        Returns
        -------
        True if the asynchronous workers for the top N rows have finished, and
        False if they have not.  None is returned if Tabular is not operating
        in "update" mode.
        """
        if self._mode != "update" or not self._content:
            return None
        #          0|..                                  <|
        #          1|..                                   |
        #          2|..                                   |
        # top_idx  3|oo <|       <|          <|           |
        #          4|oo  |--- n=3 |           |           |
        #          5|oo <|        |           |           |--- content
        #          6|oo           |           |           |    length
        #          7|oo           |--- n_free |           | (including header)
        #          8|oo           |           |    stream |
        #          9|oo           |           |--- height |
        #         10|oo          <|           |          <|
        #           |oo <|                    |
        #           |oo  |--- summary         |
        #           |oo <|                    |
        #           |oo <------ cursor       <|
        last_summary_len = self._get_last_summary_length()
        n_free = (height or self._stream.height) - last_summary_len - 1
        top_idx = self._last_content_len - n_free

        if top_idx < 0:
            # The content lines haven't yet filled the screen.
            return True

        idxs = (top_idx + i for i in range(min(n, n_free)))
        id_keys = (self._content.get_idkey(i) for i in idxs
                   if i is not None)

        futures = self._futures
        top_futures = list(chain(*(futures[k] for k in id_keys)))
        if not top_futures:
            # These rows have no registered producers.
            return True
        return all(f.done() for f in top_futures)

    def _maybe_wait_on_top_rows(self):
        n = self._wait_for_top
        if n:
            waited = 0
            secs = 0.5
            height = self._stream.height
            while self.top_nrows_done(n, height) is False:
                time.sleep(secs)
                waited += 1
            if waited:
                lgr.debug("Waited for %s cycles of sleeping %s seconds",
                          waited, secs)
                # Wait a bit longer so that the caller has a chance to see the
                # last updated row if it about to go off screen.
                time.sleep(secs)

    @skip_if_aborted
    def __call__(self, row, style=None):
        """Write styled `row`.

        Parameters
        ----------
        row : mapping, sequence, or other
            If a mapping is given, the keys are the column names and values are
            the data to write.  After the initial set of columns is defined
            (via the constructor's `columns` argument or based on inferring the
            names from the first row passed), rows can still include new keys,
            in which case the list of known columns will be expanded.

            For a sequence, the items represent the values and are taken to be
            in the same order as the constructor's `columns` argument.  Any
            other object type should have an attribute for each column
            specified via `columns`.

            Instead of a plain value, a column's value can be a tuple of the
            form (initial_value, producer).  If a producer is is a generator
            function or a generator object, each item produced replaces
            `initial_value`.  Otherwise, a producer should be a function that
            will be called with no arguments and that returns the value with
            which to replace `initial_value`.  For both generators and normal
            functions, the execution will happen asynchronously.

            Directly supplying a producer as the value rather than
            (initial_value, producer) is shorthand for ("", producer).

            The producer can return an update for multiple columns.  To do so,
            the keys of `row` should include a tuple with the column names and
            the produced value should be a tuple with the same order as the key
            or a mapping from column name to the updated value.  A mapping's
            keys may include unknown columns; these will be added to the set of
            known columns.

            Using the (initial_value, producer) form requires some additional
            steps.  The `ids` property should be set unless the first column
            happens to be a suitable id.  Also, to instruct the program to wait
            for the updated values, the instance calls should be followed by a
            call to the `wait` method or the instance should be used as a
            context manager.
        style : dict, optional
            Each top-level key should be a column name and the value should be
            a style dict that overrides the class instance style.
        """
        self._maybe_wait_on_top_rows()
        if self._columns is None:
            self._columns = self._infer_columns(row)
            lgr.debug("Inferred columns: %r", self._columns)
        if self._normalizer is None:
            self._init_prewrite()

        callables, row = self._normalizer(row)
        self._write(row, style)
        if callables:
            lgr.debug("Starting callables for row %r", row)
            self._start_callables(row, callables)

    @staticmethod
    def _infer_columns(row):
        try:
            columns = list(row.keys())
        except AttributeError:
            raise ValueError("Can't infer columns from data")
        # Make sure we don't have any multi-column keys.
        flat = []
        for column in columns:
            if isinstance(column, tuple):
                flat.extend(column)
            else:
                flat.append(column)
        return flat

    def __getitem__(self, key):
        """Get the (normalized) row for `key`.

        This interface is focused on _writing_ output, and the caller usually
        knows the values.  However, this method can be useful for retrieving
        values that were produced asynchronously (see __call__).

        Parameters
        ----------
        key : tuple
            Unique ID for a row, as specified by the `ids` property.

        Returns
        -------
        A dictionary with the row's current value.
        """
        try:
            return self._content[key]
        except KeyError as exc:
            # Suppress context.
            raise KeyError(exc) from None

    @contextmanager
    def outside_write(self, clear=False):
        """Stop writing rows and yield control to the caller.

        This context manager allows callers to interrupt the table output while
        writing their own output.  On exit, the entire table is rewritten.

        Parameters
        ----------
        clear : bool
            Before yielding to the caller, clear the visible rows.  This only
            has an effect in "update" mode.
        """
        update = self._mode == "update"
        with self._write_lock():
            if update and clear:
                n_lines = min(
                    # -1 for the last empty line of screen.
                    self._stream.height - 1,
                    self._last_content_len + self._get_last_summary_length())
                self._stream.clear_last_lines(n_lines)
            yield
            if update:
                self._stream.write(str(self._content))
                last_summary = self._last_summary
                if last_summary:
                    self._stream.write(last_summary)