from sphinx.directives.other import *
from sphinx.directives.other import __


class SpeciPyTocTree(TocTree):
    """
    Extension of the standard Sphinx TocTree that can handle references to documentation pages
    generated by SpeciPy.
    """

    def parse_content(self, toctree: addnodes.toctree) -> None:
        """
        Populate ``toctree['entries']`` and ``toctree['includefiles']`` from content.
        """
        generated_docnames = frozenset(StandardDomain._virtual_doc_names)
        suffixes = self.config.source_suffix
        current_docname = self.env.docname
        glob = toctree['glob']

        # glob target documents
        all_docnames = self.env.found_docs.copy() | generated_docnames
        all_docnames.remove(current_docname)  # remove current document
        frozen_all_docnames = frozenset(all_docnames)

        excluded = Matcher(self.config.exclude_patterns)
        for entry in self.content:
            if not entry:
                continue
            # BEGIN MOD
            entry = self.handle_specipy_entry(entry)
            # END MOD
            # look for explicit titles ("Some Title <document>")
            explicit = explicit_title_re.match(entry)
            url_match = url_re.match(entry) is not None
            if glob and glob_re.match(entry) and not explicit and not url_match:
                pat_name = docname_join(current_docname, entry)
                doc_names = sorted(
                    docname
                    for docname in patfilter(all_docnames, pat_name)
                    # don't include generated documents in globs
                    if docname not in generated_docnames
                )
                if not doc_names:
                    logger.warning(
                        __("toctree glob pattern %r didn't match any documents"),
                        entry,
                        location=toctree,
                    )

                for docname in doc_names:
                    all_docnames.remove(docname)  # don't include it again
                    toctree['entries'].append((None, docname))
                    toctree['includefiles'].append(docname)
                continue

            if explicit:
                ref = explicit.group(2)
                title = explicit.group(1)
                docname = ref
            else:
                ref = docname = entry
                title = None

            # remove suffixes (backwards compatibility)
            for suffix in suffixes:
                if docname.endswith(suffix):
                    docname = docname.removesuffix(suffix)
                    break

            # absolutise filenames
            docname = docname_join(current_docname, docname)
            if url_match or ref == 'self':
                toctree['entries'].append((title, ref))
                continue

            if docname not in frozen_all_docnames:
                if excluded(str(self.env.doc2path(docname, False))):
                    msg = __('toctree contains reference to excluded document %r')
                    subtype = 'excluded'
                else:
                    msg = __('toctree contains reference to nonexisting document %r')
                    subtype = 'not_readable'

                logger.warning(
                    msg, docname, type='toc', subtype=subtype, location=toctree
                )
                self.env.note_reread()
                continue

            if docname in all_docnames:
                all_docnames.remove(docname)
            else:
                logger.warning(
                    __('duplicated entry found in toctree: %s'),
                    docname,
                    location=toctree,
                )

            toctree['entries'].append((title, docname))
            toctree['includefiles'].append(docname)

        # entries contains all entries (self references, external links etc.)
        if 'reversed' in self.options:
            toctree['entries'] = list(reversed(toctree['entries']))
            toctree['includefiles'] = list(reversed(toctree['includefiles']))


    def handle_specipy_entry(self, entry):
        match_obj = AUTOMODEL_PATTERN_FOR_TOC_ENTRY.match(entry)
        if not match_obj:
            return entry
        model_name = match_obj.group('model_name')
        specipy_domain = self.env.domains.get('sp')
        if not specipy_domain:
            # TODO: warn?
            return entry
        model = specipy_domain.model_set.get(model_name)
        if not model:
            # TODO: warn?
            return entry
        return '/' + specipy_domain.doc_name_by_element[model]


BASE_AUTOMODEL_RE = r'<#AUTOMODEL (?P<model_name>[a-zA-Z_][a-zA-Z_0-9]*)#>'

AUTOMODEL_PATTERN_FOR_TOC_ENTRY = re.compile(fr'^{BASE_AUTOMODEL_RE}$')
AUTOMODEL_PATTERN_FOR_DOC = re.compile(fr'^[ \t]*{BASE_AUTOMODEL_RE}[ \t]*$', flags=re.MULTILINE)

def find_model_names_in_toc_trees(raw_doc: str) -> list[str]:
    # TODO: make safer (e.g., check if pattern really appears in a toc tree)
    return [match_obj.group('model_name')
            for match_obj in AUTOMODEL_PATTERN_FOR_DOC.finditer(raw_doc)]
