# Copyright 2022 The PyGlove Authors
#
# 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.
"""Utilities for formatting objects."""

import abc
import enum
import io
import sys
from typing import Any, Callable, ContextManager, Dict, List, Optional, Sequence, Set, Tuple
from pyglove.core.object_utils import thread_local


_TLS_STR_FORMAT_KWARGS = '_str_format_kwargs'
_TLS_REPR_FORMAT_KWARGS = '_repr_format_kwargs'


CustomFormatFn = Callable[
    [
        Any,    # Value to format.
        int,    # Root indent.
    ],
    # Returns a string or None. If None, the default `pg.format` will be used.
    Optional[str]
]


def str_format(**kwargs) -> ContextManager[Dict[str, Any]]:
  """Context manager for setting the default format kwargs for __str__."""
  return thread_local.thread_local_arg_scope(_TLS_STR_FORMAT_KWARGS, **kwargs)


def repr_format(**kwargs) -> ContextManager[Dict[str, Any]]:
  """Context manager for setting the default format kwargs for __repr__."""
  return thread_local.thread_local_arg_scope(_TLS_REPR_FORMAT_KWARGS, **kwargs)


class Formattable(metaclass=abc.ABCMeta):
  """Interface for classes whose instances can be pretty-formatted.

  This interface overrides the default ``__repr__`` and ``__str__`` method, thus
  all ``Formattable`` objects can be printed nicely.

  All symbolic types implement this interface.
  """

  # Additional format keyword arguments for `__str__`.
  __str_format_kwargs__ = dict(compact=False, verbose=True)

  # Additional format keyword arguments for `__repr__`.
  __repr_format_kwargs__ = dict(compact=True)

  @abc.abstractmethod
  def format(self,
             compact: bool = False,
             verbose: bool = True,
             root_indent: int = 0,
             **kwargs) -> str:
    """Formats this object into a string representation.

    Args:
      compact: If True, this object will be formatted into a single line.
      verbose: If True, this object will be formatted with verbosity.
        Subclasses should define `verbosity` on their own.
      root_indent: The start indent level for this object if the output is a
        multi-line string.
      **kwargs: Subclass specific keyword arguments.

    Returns:
      A string of formatted object.
    """

  def __str__(self) -> str:
    """Returns the full (maybe multi-line) representation of this object."""
    # NOTE(daiyip): we delegate the formatting logic to `pg.format` instead of
    # `Formattable.format` as `pg.format` could add common functionalities such
    # as `markdown` support.
    return format(self, **self.__str_kwargs__())

  def __str_kwargs__(self) -> Dict[str, Any]:
    """Returns the default format kwargs for __str__."""
    kwargs = dict(self.__str_format_kwargs__)
    kwargs.update(thread_local.thread_local_kwargs(_TLS_STR_FORMAT_KWARGS))
    return kwargs

  def __repr__(self) -> str:
    """Returns a single-line representation of this object."""
    # NOTE(daiyip): we delegate the formatting logic to `pg.format` instead of
    # `Formattable.format` as `pg.format` could add common functionalities such
    # as `markdown` support.
    return format(self, **self.__repr_kwargs__())

  def __repr_kwargs__(self) -> Dict[str, Any]:
    """Returns the default format kwargs for __repr__."""
    kwargs = dict(self.__repr_format_kwargs__)
    kwargs.update(thread_local.thread_local_kwargs(_TLS_REPR_FORMAT_KWARGS))
    return kwargs


class RawText(Formattable):
  """Raw text."""

  def __init__(self, text: str):
    self.text = text

  def format(self, *args, **kwargs):
    del args, kwargs
    return self.text

  def __eq__(self, other: Any) -> bool:
    if isinstance(other, RawText):
      return self.text == other.text
    elif isinstance(other, str):
      return self.text == other
    return False

  def __ne__(self, other: Any) -> bool:
    return not self.__eq__(other)


class BracketType(enum.IntEnum):
  """Bracket types used for complex type formatting."""
  # Round bracket.
  ROUND = 0

  # Square bracket.
  SQUARE = 1

  # Curly bracket.
  CURLY = 2


_BRACKET_CHARS = [
    ('(', ')'),
    ('[', ']'),
    ('{', '}'),
]


def bracket_chars(bracket_type: BracketType) -> Tuple[str, str]:
  """Gets bracket character."""
  return _BRACKET_CHARS[int(bracket_type)]


def kvlist_str(
    kvlist: List[Tuple[str, Any, Any]],
    compact: bool = True,
    verbose: bool = False,
    root_indent: int = 0,
    *,
    label: Optional[str] = None,
    bracket_type: BracketType = BracketType.ROUND,
    custom_format: Optional[CustomFormatFn] = None,
    **kwargs,
) -> str:
  """Formats a list key/value pairs into a comma delimited string.

  Args:
    kvlist: List of tuples in format of
      (key, value, default_value or a tuple of default values)
    compact: If True, format value in kvlist in compact form.
    verbose: If True, format value in kvlist in verbose.
    root_indent: The indent should be applied for values in kvlist if they are
      multi-line.
    label: (Optional) If not None, add label to brace all kv pairs.
    bracket_type: Bracket type used for embracing the kv pairs. Applicable only
      when `name` is not None.
    custom_format: An optional custom format function, which will be applied to
      each value (and child values) in kvlist. If the function returns None, it
      will fall back to the default `pg.format`.
    **kwargs: Keyword arguments that will be passed through unto child
      ``Formattable`` objects.
  Returns:
    A formatted string from a list of key/value pairs delimited by comma.
  """
  s = io.StringIO()
  is_first = True
  bracket_start, bracket_end = bracket_chars(bracket_type)

  child_indent = (root_indent + 1) if label else root_indent
  body = io.StringIO()

  for k, v, d in kvlist:
    if isinstance(d, tuple):
      include_pair = True
      for sd in d:
        if sd == v:
          include_pair = False
          break
    else:
      include_pair = v != d
    if include_pair:
      if not is_first:
        body.write(',')
        body.write(' ' if compact else '\n')
      v = format(
          v, compact=compact,
          verbose=verbose,
          root_indent=child_indent,
          custom_format=custom_format,
          **kwargs
      )
      if not compact:
        body.write(_indent('', child_indent))
      if k:
        body.write(f'{k}={str_ext(v, custom_format, child_indent)}')
      else:
        body.write(str_ext(v, custom_format, child_indent))
      is_first = False

  if label and not is_first and not compact:
    body.write('\n')

  body = body.getvalue()
  if label is None:
    return body
  else:
    s.write(label)
    s.write(bracket_start)
    if body:
      if not compact:
        s.write('\n')
      s.write(body)
      if not compact:
        s.write(_indent('', root_indent))
    s.write(bracket_end)
    return s.getvalue()


def quote_if_str(value: Any) -> Any:
  """Quotes the value if it is a str."""
  if isinstance(value, str):
    return repr(value)
  return value


def comma_delimited_str(value_list: Sequence[Any]) -> str:
  """Gets comma delimited string."""
  return ', '.join(str(quote_if_str(v)) for v in value_list)


def auto_plural(
    number: int, singular: str, plural: Optional[str] = None) -> str:
  """Use singular form if number is 1, otherwise use plural form."""
  if plural is None:
    plural = singular + 's'
  return singular if number == 1 else plural


def format(   # pylint: disable=redefined-builtin
    value: Any,
    compact: bool = False,
    verbose: bool = True,
    root_indent: int = 0,
    list_wrap_threshold: int = 80,
    strip_object_id: bool = False,
    include_keys: Optional[Set[str]] = None,
    exclude_keys: Optional[Set[str]] = None,
    markdown: bool = False,
    max_str_len: Optional[int] = None,
    max_bytes_len: Optional[int] = None,
    *,
    custom_format: Optional[CustomFormatFn] = None,
    **kwargs,
) -> str:
  """Formats a (maybe) hierarchical value with flags.

  Args:
    value: The value to format.
    compact: If True, this object will be formatted into a single line.
    verbose: If True, this object will be formatted with verbosity. Subclasses
      should define `verbosity` on their own.
    root_indent: The start indent level for this object if the output is a
      multi-line string.
    list_wrap_threshold: A threshold in number of characters for wrapping a list
      value in a single line.
    strip_object_id: If True, format object as '<class-name>(...)' other than
      'object at <address>'.
    include_keys: A set of keys to include from the top-level dict or object.
    exclude_keys: A set of keys to exclude from the top-level dict or object.
      Applicable only when `include_keys` is set to None.
    markdown: If True, use markdown notion to quote the formatted object.
    max_str_len: The max length of the string to be formatted. If the string is
      longer than this length, it will be truncated.
    max_bytes_len: The max length of the bytes to be formatted. If the bytes is
      longer than this length, it will be truncated.
    custom_format: An optional custom format function, which will be applied to
      each value (and child values) in kvlist. If the function returns None, it
      will fall back to the default `pg.format`.
    **kwargs: Keyword arguments that will be passed through unto child
      ``Formattable`` objects.

  Returns:
    A string representation for `value`.
  """
  # We allow custom_format to intercept the entire value if it's present.
  if custom_format is not None:
    result = custom_format(value, root_indent)
    if result is not None:
      return maybe_markdown_quote(result, markdown)

  exclude_keys = exclude_keys or set()

  def _should_include_key(key: str) -> bool:
    if include_keys:
      return key in include_keys
    return key not in exclude_keys

  def _format_child(v):
    return format(
        v,
        compact=compact,
        verbose=verbose,
        root_indent=root_indent + 1,
        list_wrap_threshold=list_wrap_threshold,
        strip_object_id=strip_object_id,
        max_str_len=max_str_len,
        max_bytes_len=max_bytes_len,
        custom_format=custom_format,
        **kwargs
    )

  # `markdown` is only applied at the outter most level,
  # so we disable markdown format in the sub-tree for `str` and `repr`.
  with str_format(markdown=False), repr_format(markdown=False):
    if isinstance(value, Formattable):
      s = value.format(
          compact=compact,
          verbose=verbose,
          root_indent=root_indent,
          list_wrap_threshold=list_wrap_threshold,
          strip_object_id=strip_object_id,
          include_keys=include_keys,
          exclude_keys=exclude_keys,
          max_str_len=max_str_len,
          max_bytes_len=max_bytes_len,
          custom_format=custom_format,
          **kwargs
      )
    elif isinstance(value, (list, tuple)):
      # Always try compact representation if length is not too long.
      open_bracket, close_bracket = bracket_chars(
          BracketType.SQUARE if isinstance(value, list) else BracketType.ROUND)
      s = [open_bracket]
      s.append(', '.join([_format_child(elem) for elem in value]))
      s.append(close_bracket)
      s = [''.join(s)]
      if not compact and len(s[-1]) > list_wrap_threshold:
        s = [f'{open_bracket}\n']
        s.append(',\n'.join([
            _indent(_format_child(elem), root_indent + 1)
            for elem in value
        ]))
        s.append('\n')
        s.append(_indent(close_bracket, root_indent))
    elif isinstance(value, dict):
      if compact or not value:
        s = ['{']
        s.append(', '.join([
            f'{k!r}: {_format_child(v)}'
            for k, v in value.items() if _should_include_key(k)
        ]))
        s.append('}')
      else:
        s = ['{\n']
        s.append(',\n'.join([
            _indent(f'{k!r}: {_format_child(v)}', root_indent + 1)
            for k, v in value.items() if _should_include_key(k)
        ]))
        s.append('\n')
        s.append(_indent('}', root_indent))
    else:
      if isinstance(value, str):
        if max_str_len is not None and len(value) > max_str_len:
          value = value[:max_str_len] + '...'
        s = [repr_ext(value, custom_format, root_indent)]
      elif isinstance(value, bytes):
        if max_bytes_len is not None and len(value) > max_bytes_len:
          value = value[:max_bytes_len] + b'...'
        s = [repr_ext(value, custom_format, root_indent)]
      else:
        s = [repr_ext(value, custom_format, root_indent)
             if compact else str_ext(value, custom_format, root_indent)]
        if strip_object_id and 'object at 0x' in s[-1]:
          s = [f'{value.__class__.__name__}(...)']
  return maybe_markdown_quote(''.join(s), markdown)


def _maybe_custom_format(
    v: Any,
    default_fn: Callable[[Any], str],
    custom_format: Optional[CustomFormatFn],
    root_indent: int,
) -> str:
  if custom_format is None:
    return default_fn(v)
  x = custom_format(v, root_indent)
  if x is None:
    return default_fn(v)
  return x


def str_ext(
    v: Any,
    custom_format: Optional[CustomFormatFn] = None,
    root_indent: int = 0,
) -> str:
  """"str operator with special format support."""
  return _maybe_custom_format(v, str, custom_format, root_indent)


def repr_ext(
    v: Any,
    custom_format: Optional[CustomFormatFn] = None,
    root_indent: int = 0,
) -> str:
  """repr operator with special format support."""
  return _maybe_custom_format(v, repr, custom_format, root_indent)


def maybe_markdown_quote(s: str, markdown: bool = True) -> str:
  """Maybe quote the formatted string with markdown."""
  if not markdown:
    return s
  if '\n' not in s:
    return f'`{s}`'
  else:
    return f'```\n{s}\n```'


def printv(v: Any, **kwargs):
  """Prints formatted value."""
  fs = kwargs.pop('file', sys.stdout)
  print(format(v, **kwargs), file=fs)


def _indent(text: str, indent: int) -> str:
  return ' ' * 2 * indent + text
