from typing import Optional, Any

from docutils.statemachine import StringList
from sphinx.application import Sphinx
from sphinx.ext.autodoc import ClassDocumenter, PropertyDocumenter, MethodDocumenter, bool_option
from sphinx.util import inspect
from sphinx.util.typing import restify


def setup(app: Sphinx) -> None:
  app.setup_extension('sphinx.ext.autodoc')  # Require autodoc extension
  app.add_autodocumenter(CercClassDocumenter)
  app.add_autodocumenter(CercPropertyDocumenter)
  app.add_autodocumenter(CercMethodDocumenter)


class CercMethodDocumenter(MethodDocumenter):
  objtype = 'cercmethod'
  directivetype = 'method'
  member_order = 50
  priority = MethodDocumenter.priority + 10

  @staticmethod
  def _get_type(sig_return):
    types = sig_return.replace('Optional', '').replace('List', '').split('.')
    prefix = ''
    if types[0][0] in '[(':
      prefix = types[0][0]
    return f'{prefix}{types[len(types) -1]}'

  def format_signature(self, **kwargs: Any) -> str:
    sig = super().format_signature(**kwargs)
    if '->' in sig:
      sig_pars = sig.split(' -> ')
      return_type = CercMethodDocumenter._get_type(sig_pars[1])
      return_str = f'{return_type}'
      if 'Optional' in sig_pars[1]:
        return_str = f'Optional{return_str}'
      return f'{sig_pars[0]} -> {return_str}'
    return sig

  def add_content(self, more_content: Optional[StringList], no_docstring: bool = False) -> None:
    source_name = self.get_sourcename()
    docstrings = self.get_doc()
    self.content_indent = '   '
    for i, line in enumerate(self.process_doc(docstrings)):
      if ':return:' in line:
        line = line.replace(':return:', '**Returns**')
      elif ':param ' in line:
        line = line.replace(':param ', '**Parameter** ')
        line = line.replace(':', ' ')
      elif ':alert:' in line:
        line = line.replace(':alert:', '|alert|')
      elif ':cat:' in line:
        line = line.replace(':cat:', '|cat|')
      self.add_line(f'{line}\n', source_name, i)


class CercPropertyDocumenter(PropertyDocumenter):
  objtype = 'cercproperty'
  directivetype = 'property'
  member_order = 60
  priority = PropertyDocumenter.priority + 10

  @staticmethod
  def _get_type(annotations):
    if len(annotations) == 0:
      return annotations
    return_str = str(annotations['return']).replace('typing.', '')
    if 'Union' in return_str:
      return_str = return_str.replace('Union', 'Optional').replace(', None', '').replace('None, ', '')
      return_str = return_str.replace('NoneType, ','').replace(', NoneType', '')
    if 'List' in return_str:
      return_str = return_str.replace('typing.List', '').replace('List', '')
    if "<class '" in return_str:
      return_str = return_str.replace("<class '", '').replace("'>", '')
    if '.' in return_str:
      for i, c in enumerate(reversed(return_str)):
        if c not in '])':
          types_str = return_str.split('.')
          if 'Optional' in return_str:
            return_str = f'{types_str[0][0:8+i]}{types_str[len(types_str)-1]}'
          else:
            return_str = f'{types_str[0][0:i]}{types_str[len(types_str)-1]}'
          break
    return_str = return_str.replace('~', '').replace('Type', '')
    return return_str

  def add_directive_header(self, sig: str) -> None:
    name = self.format_name()
    source_name = self.get_sourcename()
    annotations = self.get_attr(self.object.fget, '__annotations__', None)
    self.retann = CercPropertyDocumenter._get_type(annotations)
    for i, sig_line in enumerate(sig.split("\n")):
      self.add_line(f'.. py:property:: {name}({sig_line}) {(f" -> {self.retann}" if self.retann else "")}', source_name)
    if self.objpath:
      self.add_line(f'   :module: {self.modname}', source_name)
    if inspect.isabstractmethod(self.object):
      self.add_line('      :abstractmethod:', source_name)
    self.add_line(f'      :property:', source_name)

  def add_content(self, more_content: Optional[StringList], no_docstring: bool = False) -> None:
    source_name = self.get_sourcename()
    docstrings = self.get_doc()
    self.content_indent = '   '
    if not no_docstring:
      if not docstrings:
        docstrings.append([])
      for i, line in enumerate(self.process_doc(docstrings)):
        if ':return:' in line:
          line = line.replace(':return:', '**Returns**')
        elif ':alert:' in line:
          line = line.replace(':alert:', '|alert|')
        self.add_line(f'{line}\n', source_name, i)


class CercClassDocumenter(ClassDocumenter):
  objtype = 'cercclass'
  directivetype = 'class'
  priority = 10 + ClassDocumenter.priority
  option_spec = dict(ClassDocumenter.option_spec)
  option_spec['hex'] = bool_option

  def get_bases(self, source_name):
    bases = []
    if hasattr(self.object, '__orig_bases__') and len(self.object.__orig_bases__):
      bases = [restify(cls) for cls in self.object.__orig_bases__]
    elif hasattr(self.object, '__bases__') and len(self.object.__bases__):
      # A normal class
      bases = [restify(cls) for cls in self.object.__bases__]
    if ':class:`object`' in bases:
      bases.remove(':class:`object`')
    cleaned_bases = []
    for base in bases:
      base_tmp = base.replace('`', '').split('.')
      base_tmp = base_tmp[len(base_tmp)-1]
      cleaned_bases.append(base_tmp)
    if len(cleaned_bases) != 0:
      self.add_line('   Inherit: %s' % ', '.join(cleaned_bases), source_name)
      self.add_line(' ', source_name)

  def add_directive_header(self, sig: str) -> None:
    prefix = f'.. py:class:: '
    name = self.format_name()
    source_name = self.get_sourcename()
    self.add_line('%s%s%s' % (prefix, name, sig), source_name)
    if self.analyzer and '.'.join(self.objpath) in self.analyzer.finals:
      self.add_line('   :final:', source_name)
    self.get_bases(source_name)

  def add_content(self, more_content: Optional[StringList], no_docstring: bool = False) -> None:
    source_name = self.get_sourcename()
    docstrings = self.get_doc()
    self.content_indent = '   '
    if not no_docstring:
      if not docstrings:
        docstrings.append([])
      for i, line in enumerate(self.process_doc(docstrings)):
        self.add_line(line, source_name)