from __future__ import annotations

import dataclasses as dc
import html
import itertools
from typing import Iterable, Literal, Sequence, cast, get_args

from markdown_it.token import Token

from .src_error import SrcError
from .utils import Freezeable

# FragmentType is used to restrict structural include blocks.
FragmentType = Literal['preface', 'part', 'chapter', 'section', 'appendix']

# the TOC allows every fragment type. it also allows the enclosing book.
# it allows 'example' and 'figure'.
# --experimental-config adds a generic 'content' kind.
TocEntryType = Literal['book', 'preface', 'part', 'chapter', 'section', 'appendix', 'example', 'figure', 'content']

def is_include(token: Token) -> bool:
    return token.type == "fence" and token.info.startswith("{=include=} ")

# toplevel file must contain only the title headings and includes, anything else
# would cause strange rendering.
def _check_book_structure(src: str, tokens: Sequence[Token]) -> None:
    for token in tokens[6:]:
        if not is_include(token):
            raise SrcError(
                src=src,
                description=f"unexpected content; expected structural include",
                token=token,
            )

# much like books, parts may not contain headings other than their title heading.
# this is a limitation of the current renderers and TOC generators that do not handle
# this case well even though it is supported in docbook (and probably supportable
# anywhere else).
def _check_part_structure(src: str,tokens: Sequence[Token]) -> None:
    _check_fragment_structure(src, tokens)
    for token in tokens[3:]:
        if token.type == 'heading_open':
            raise SrcError(
                src=src,
                description="unexpected heading",
                token=token,
            )

# two include blocks must either be adjacent or separated by a heading, otherwise
# we cannot generate a correct TOC (since there'd be nothing to link to between
# the two includes).
def _check_fragment_structure(src: str, tokens: Sequence[Token]) -> None:
    for i, token in enumerate(tokens):
        if is_include(token) \
           and i + 1 < len(tokens) \
           and not (is_include(tokens[i + 1]) or tokens[i + 1].type == 'heading_open'):
            assert token.map
            raise SrcError(
                src=src,
                description="unexpected content; expected heading or structural include",
                token=token,
            )

def check_structure(src: str, kind: TocEntryType, tokens: Sequence[Token]) -> None:
    wanted = { 'h1': 'title' }
    wanted |= { 'h2': 'subtitle' } if kind == 'book' else {}
    for (i, (tag, role)) in enumerate(wanted.items()):
        if len(tokens) < 3 * (i + 1):
            raise RuntimeError(f"missing {role} ({tag}) heading")
        token = tokens[3 * i]
        if token.type != 'heading_open' or token.tag != tag:
            raise SrcError(
                src=src,
                description=f"expected {role} ({tag}) heading",
                token=token,
            )
    for t in tokens[3 * len(wanted):]:
        if t.type != 'heading_open' or not (role := wanted.get(t.tag, '')):
            continue
        raise SrcError(
            src=src,
            description=f"only one {role} heading ({t.markup} [text...]) allowed per "
            f"{kind}, but found a second. "
            "Please remove all such headings except the first or demote the subsequent headings.",
            token=t,
        )

    last_heading_level = 0
    for token in tokens:
        if token.type != 'heading_open':
            continue

        # book subtitle headings do not need an id, only book title headings do.
        # every other headings needs one too. we need this to build a TOC and to
        # provide stable links if the manual changes shape.
        if 'id' not in token.attrs and (kind != 'book' or token.tag != 'h2'):
            raise SrcError(
                src=src,
                description=f"heading does not have an id",
                token=token,
            )

        level = int(token.tag[1:]) # because tag = h1..h6
        if level > last_heading_level + 1:
            raise SrcError(
                src=src,
                description=f"heading skips one or more heading levels, "
                               "which is currently not allowed",
                token=token,
            )
        last_heading_level = level

    if kind == 'book':
        _check_book_structure(src, tokens)
    elif kind == 'part':
        _check_part_structure(src, tokens)
    else:
        _check_fragment_structure(src, tokens)

@dc.dataclass(frozen=True)
class XrefTarget:
    id: str
    """link label for `[](#local-references)`"""
    title_html: str
    """toc label"""
    toc_html: str | None
    """text for `<title>` tags and `title="..."` attributes"""
    title: str | None
    """path to file that contains the anchor"""
    path: str
    """whether to drop the `#anchor` from links when expanding xrefs"""
    drop_fragment: bool = False
    """whether to drop the `path.html` from links when expanding xrefs.
       mostly useful for docbook compatibility"""
    drop_target: bool = False

    def href(self) -> str:
        path = "" if self.drop_target else html.escape(self.path, True)
        return path if self.drop_fragment else f"{path}#{html.escape(self.id, True)}"

@dc.dataclass
class TocEntry(Freezeable):
    kind: TocEntryType
    target: XrefTarget
    parent: TocEntry | None = None
    prev: TocEntry | None = None
    next: TocEntry | None = None
    children: list[TocEntry] = dc.field(default_factory=list)
    starts_new_chunk: bool = False
    open_id: str | None = None

    @property
    def root(self) -> TocEntry:
        return self.parent.root if self.parent else self

    @classmethod
    def of(cls, token: Token) -> TocEntry:
        entry = token.meta.get('TocEntry')
        if not isinstance(entry, TocEntry):
            raise RuntimeError('requested toc entry, none found', token)
        return entry

    @classmethod
    def collect_and_link(cls, xrefs: dict[str, XrefTarget], tokens: Sequence[Token]) -> TocEntry:
        entries = cls._collect_entries(xrefs, tokens, 'book')

        def flatten_with_parent(this: TocEntry, parent: TocEntry | None) -> Iterable[TocEntry]:
            this.parent = parent
            return itertools.chain([this], *[ flatten_with_parent(c, this) for c in this.children ])

        flat = list(flatten_with_parent(entries, None))
        prev = flat[0]
        prev.starts_new_chunk = True
        paths_seen = set([prev.target.path])
        for c in flat[1:]:
            if prev.target.path != c.target.path and c.target.path not in paths_seen:
                c.starts_new_chunk = True
                c.prev, prev.next = prev, c
                prev = c
            paths_seen.add(c.target.path)

        for c in flat:
            c.freeze()

        return entries

    @classmethod
    def _collect_entries(cls, xrefs: dict[str, XrefTarget], tokens: Sequence[Token],
                         kind: TocEntryType) -> TocEntry:
        # we assume that check_structure has been run recursively over the entire input.
        # list contains (tag, entry) pairs that will collapse to a single entry for
        # the full sequence.
        entries: list[tuple[str, TocEntry]] = []
        for token in tokens:
            if token.type.startswith('included_') and (included := token.meta.get('included')):
                fragment_type_str = token.type[9:].removesuffix('s')
                assert fragment_type_str in get_args(TocEntryType)
                fragment_type = cast(TocEntryType, fragment_type_str)
                if (nav_label := token.meta.get('nav-label')) is not None:
                    children = [cls._collect_entries(xrefs, fragment, fragment_type)
                                for fragment, _path in included]
                    first = children[0].target
                    group = TocEntry(
                        fragment_type,
                        dc.replace(first, title_html=nav_label, toc_html=nav_label, title=nav_label),
                        children=children,
                        open_id=token.meta.get('nav-id', ''))
                    entries.append(('h1', group))
                    token.meta['TocEntry'] = group
                else:
                    for fragment, _path in included:
                        subentries = cls._collect_entries(xrefs, fragment, fragment_type)
                        entries[-1][1].children.append(subentries)
            elif token.type == 'heading_open' and (id := cast(str, token.attrs.get('id', ''))):
                while len(entries) > 1 and entries[-1][0] >= token.tag:
                    entries[-2][1].children.append(entries.pop()[1])
                entries.append((token.tag,
                                TocEntry(kind if token.tag == 'h1' else 'section', xrefs[id])))
                token.meta['TocEntry'] = entries[-1][1]


        while len(entries) > 1:
            entries[-2][1].children.append(entries.pop()[1])
        return entries[0][1]

_xml_id_translate_table = {
    ord('*'): ord('_'),
    ord('<'): ord('_'),
    ord(' '): ord('_'),
    ord('>'): ord('_'),
    ord('['): ord('_'),
    ord(']'): ord('_'),
    ord(':'): ord('_'),
    ord('"'): ord('_'),
}
# this function is needed to generate option id attributes in the same format as
# the docbook toolchain did to not break existing links. we don't actually use
# xml any more, that's just the legacy we're dealing with and part of our structure
# now.
def make_xml_id(s: str) -> str:
    return s.translate(_xml_id_translate_table)
