"""Markdown file parser — frontmatter + wikilink graph + AST tree chunks.

Each chunk carries the **complete heading skeleton** of the document
with its content inlined under the section that owns it; other sections
appear as bare headings so the reader always sees a full document map.

Pipeline: build mistletoe AST → ``MdNode`` tree (sections nest by
heading level) → recursive chunk (try whole subtree; on overflow walk
children — body siblings pack as a run, subsections recurse). Leaf
blocks (table / code / list / paragraph) split on internal boundaries
and each piece is annotated ``[Part X/N]``.
"""

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import frontmatter
from mistletoe.block_token import (
    CodeFence,
    Document,
    Heading,
    List,
    ListItem,
    SetextHeading,
    Table,
    TableRow,
)
from mistletoe.markdown_renderer import BlankLine, MarkdownRenderer

from .base_file_parser import BaseFileParser
from ..component_registry import R
from ...schema import FileChunk, FileEdge, FileFrontMatter, FileNode
from ...utils import hash_text

_PART_RESERVE = 18  # worst-case "[Part NNN/NNN]\n\n" prefix


# -- AST node + helpers ---------------------------------------------------


@dataclass
class MdNode:
    """``root`` (top-level container) / ``section`` (heading + children
    until next equal-or-shallower heading) / ``body`` (one mistletoe
    block, ``block`` keeps the original).

    ``text`` is the rendered subtree (own heading excluded for sections;
    set by build for bodies, by ``_finalize`` for sections/root).
    ``desc_toc`` caches the section-only DFS outline of descendants
    (own heading excluded), used as the TOC suffix when emitting chunks
    inside a section. Line ranges span the full subtree.
    """

    kind: str  # "root" | "section" | "body"
    heading: str | None = None
    level: int = 0
    children: list["MdNode"] = field(default_factory=list)
    block: Any = None
    text: str = ""
    start_line: int = 0
    end_line: int = 0
    desc_toc: str = ""


def _heading_text(node: Any, renderer: MarkdownRenderer) -> str:
    """Heading text without `#` markers (for outline)."""
    rendered = renderer.render(node).rstrip("\n")
    if rendered.startswith("#"):
        return rendered.lstrip("#").strip()
    return rendered.split("\n", 1)[0].strip()


def _dedup_edges(edges: list[FileEdge]) -> list[FileEdge]:
    seen: set[tuple] = set()
    out: list[FileEdge] = []
    for e in edges:
        key = (e.link, e.predicate)
        if key not in seen:
            seen.add(key)
            out.append(e)
    return out


def _finalize(n: MdNode) -> None:
    """Bottom-up post-build pass: propagate line ranges up from children,
    populate ``n.text`` (rendered subtree, own heading excluded for
    sections; bodies already set by build), and ``n.desc_toc`` (DFS
    section outline of descendants — own heading excluded; used as the
    TOC suffix when emitting body chunks inside this section).
    """
    parts: list[str] = []
    desc_lines: list[str] = []
    for c in n.children:
        _finalize(c)
        if c.kind == "section":
            heading = f"{'#' * c.level} {c.heading or ''}"
            parts.append(f"{heading}\n\n{c.text}" if c.text else heading)
            desc_lines.append(f"{heading}\n\n{c.desc_toc}" if c.desc_toc else heading)
        elif c.text:
            parts.append(c.text)
    if n.children:
        first = n.children[0].start_line
        n.start_line = min(n.start_line, first) if n.start_line else first
        n.end_line = max(c.end_line for c in n.children)
    elif n.end_line < n.start_line:
        n.end_line = n.start_line
    if n.kind != "body":
        n.text = "\n\n".join(parts)
    n.desc_toc = "\n\n".join(desc_lines)


def _toc_join(*parts: str) -> str:
    """Concatenate TOC fragments with ``\\n\\n``, skipping empty ones."""
    return "\n\n".join(p for p in parts if p)


def _subtree_toc(n: MdNode) -> str:
    """Section's heading + descendants TOC — its contribution to a parent's
    ``desc_toc``. For root (no own heading) this is just ``desc_toc``."""
    if n.kind != "section" or n.heading is None:
        return n.desc_toc
    heading = f"{'#' * n.level} {n.heading}"
    return f"{heading}\n\n{n.desc_toc}" if n.desc_toc else heading


# -- Parser ---------------------------------------------------------------


@R.register("md")
class LinkedFileParser(BaseFileParser):
    """Markdown parser: frontmatter + wikilink edges + full-skeleton chunks."""

    def __init__(
            self,
            encoding: str = "utf-8",
            chunk_chars: int = 2000,
            embed_toc: bool = True,
            **kwargs,
    ):
        super().__init__(**kwargs)
        self.encoding = encoding
        self.chunk_chars = max(100, chunk_chars)
        self.embed_toc = embed_toc

    async def parse(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]:
        file_path = Path(path)
        post = frontmatter.loads(file_path.read_text(encoding=self.encoding))
        absolute = str(file_path.absolute())

        chunks: list[FileChunk] = []
        if post.content and post.content.strip():
            with MarkdownRenderer() as renderer:
                tree = self._build_tree(Document(post.content), renderer)
                chunks = self._chunk_node(tree, "", "", absolute, renderer)

        node = FileNode(
            path=absolute,
            st_mtime=file_path.stat().st_mtime,
            chunk_ids=[chunk.id for chunk in chunks],
            edges=_dedup_edges(FileEdge.from_text(post.content)),
            front_matter=FileFrontMatter(**dict(post.metadata)),
        )
        return node, chunks

    def _build_tree(self, doc: Any, renderer: MarkdownRenderer) -> MdNode:
        """Heading-level stack folds mistletoe's flat children into nested
        sections; non-headings attach as ``body`` to the current section
        (or root before the first heading)."""
        root = MdNode(kind="root", start_line=1, end_line=1)
        stack: list[MdNode] = [root]
        for child in doc.children or []:
            if isinstance(child, BlankLine):
                continue
            line = getattr(child, "line_number", None) or stack[-1].start_line
            if isinstance(child, (Heading, SetextHeading)):
                level = max(1, getattr(child, "level", 1))
                while len(stack) > 1 and stack[-1].level >= level:
                    stack.pop()
                sec = MdNode(
                    kind="section",
                    heading=_heading_text(child, renderer),
                    level=level,
                    start_line=line,
                )
                stack[-1].children.append(sec)
                stack.append(sec)
                continue
            rendered = renderer.render(child).rstrip("\n")
            if not rendered:
                continue
            stack[-1].children.append(MdNode(
                kind="body", block=child, text=rendered,
                start_line=line, end_line=line + rendered.count("\n"),
            ))
        _finalize(root)
        return root

    # -- Recursive chunker ------------------------------------------------

    def _chunk_node(
            self, node: MdNode, before: str, after: str,
            path: str, renderer: MarkdownRenderer,
    ) -> list[FileChunk]:
        """Try the whole subtree; on overflow split (leaf) or descend.
        ``before``/``after`` are TOC fragments that bracket each emitted
        chunk's content (chunk text = ``before + content + after``).
        As we descend, the prefix grows with section headings already
        passed and the suffix shrinks correspondingly.
        """
        if not node.text:
            return []
        if node.kind == "section":
            heading_line = f"{'#' * node.level} {node.heading or ''}"
            before_self = _toc_join(before, heading_line)
        else:
            before_self = before
        if len(node.text) <= self.chunk_chars:
            return [self._make_chunk(before_self, node.text, after,
                                     node.start_line, node.end_line, path)]
        if node.kind == "body":
            return self._split_leaf(node, before, after, path, renderer)
        after_inside = _toc_join(node.desc_toc, after)
        sub_tocs = [_subtree_toc(c) for c in node.children if c.kind == "section"]
        chunks: list[FileChunk] = []
        accumulated = before_self
        sec_idx = 0
        run: list[MdNode] = []
        for c in node.children:
            if c.kind == "section":
                if run:
                    chunks.extend(self._chunk_body_run(
                        run, before_self, after_inside, path, renderer))
                    run = []
                remaining = "\n\n".join(sub_tocs[sec_idx + 1:])
                chunks.extend(self._chunk_node(
                    c, accumulated, _toc_join(remaining, after), path, renderer))
                accumulated = _toc_join(accumulated, sub_tocs[sec_idx])
                sec_idx += 1
            else:
                run.append(c)
        if run:
            chunks.extend(self._chunk_body_run(
                run, before_self, after_inside, path, renderer))
        return chunks

    def _chunk_body_run(
            self, run: list[MdNode], before: str, after: str,
            path: str, renderer: MarkdownRenderer,
    ) -> list[FileChunk]:
        """Greedy-pack consecutive body siblings under the same TOC slot.
        No ``[Part X/N]`` markers — distinct blocks, not a leaf split.
        Oversized single body recurses to ``_split_leaf``."""
        composite_size = sum(len(b.text) for b in run) + 2 * max(0, len(run) - 1)
        if composite_size <= self.chunk_chars:
            return [self._make_chunk(before, "\n\n".join(b.text for b in run), after,
                                     run[0].start_line, run[-1].end_line, path)]

        chunks: list[FileChunk] = []
        bucket: list[MdNode] = []
        bucket_chars = 0

        def flush() -> None:
            nonlocal bucket, bucket_chars
            if not bucket:
                return
            chunks.append(self._make_chunk(
                before, "\n\n".join(b.text for b in bucket), after,
                bucket[0].start_line, bucket[-1].end_line, path))
            bucket = []
            bucket_chars = 0

        for body in run:
            if len(body.text) > self.chunk_chars:
                flush()
                chunks.extend(self._split_leaf(body, before, after, path, renderer))
                continue
            sep = 2 if bucket else 0
            if bucket_chars + sep + len(body.text) > self.chunk_chars:
                flush()
                sep = 0
            bucket.append(body)
            bucket_chars += sep + len(body.text)
        flush()
        return chunks

    # -- Leaf splitters: build (text, start, end) units, hand off to packer

    def _split_leaf(
            self, body: MdNode, before: str, after: str,
            path: str, renderer: MarkdownRenderer,
    ) -> list[FileChunk]:
        block = body.block
        if isinstance(block, Table):
            return self._split_table(body, before, after, path)
        if isinstance(block, CodeFence):
            return self._split_code(body, before, after, path)
        if isinstance(block, List):
            return self._split_list(body, before, after, path, renderer)
        return self._split_lines(body, before, after, path)

    def _split_table(
            self, body: MdNode, before: str, after: str, path: str,
    ) -> list[FileChunk]:
        """Repeat header + separator on every chunk."""
        lines = body.text.split("\n")
        header, data = "\n".join(lines[:2]), lines[2:]
        rows = [r for r in (body.block.children or []) if isinstance(r, TableRow)]
        base = body.start_line + 2

        def line_of(i: int) -> int:
            return rows[i].line_number if i < len(rows) and rows[i].line_number else base + i

        units = [(text, line_of(i), line_of(i)) for i, text in enumerate(data)]
        return self._emit_packed(units, before, after, path,
                                 joiner="\n", wrap=f"{header}\n{{inner}}")

    def _split_code(
            self, body: MdNode, before: str, after: str, path: str,
    ) -> list[FileChunk]:
        """Repeat fence opener + closer on every chunk."""
        code = body.block
        indent = " " * (code.indentation or 0)
        fence = f"{indent}{code.delimiter}"
        opener = f"{fence}{code.info_string or ''}"
        raw = (code.children[0].content if code.children else "").rstrip("\n")
        if not raw:
            return []
        start = body.start_line + 1
        units = [(indent + ln, start + i, start + i)
                 for i, ln in enumerate(raw.split("\n"))]
        return self._emit_packed(units, before, after, path, joiner="\n",
                                 wrap=f"{opener}\n{{inner}}\n{fence}",
                                 allow_empty=True)

    def _split_list(
            self, body: MdNode, before: str, after: str,
            path: str, renderer: MarkdownRenderer,
    ) -> list[FileChunk]:
        """Pack list items; oversized items emit alone (overflow accepted)."""
        items = [c for c in (body.block.children or []) if isinstance(c, ListItem)]
        if not items:
            return self._split_lines(body, before, after, path)
        units: list[tuple[str, int, int]] = []
        for it in items:
            text = renderer.render(it).rstrip("\n")
            if not text:
                continue
            line = it.line_number or body.start_line
            units.append((text, line, line + text.count("\n")))
        return self._emit_packed(units, before, after, path,
                                 joiner="\n", wrap="{inner}")

    def _split_lines(
            self, body: MdNode, before: str, after: str, path: str,
    ) -> list[FileChunk]:
        """Last-resort line-greedy split for paragraphs / quotes / html."""
        start = body.start_line
        units = [(line, start + i, start + i)
                 for i, line in enumerate(body.text.split("\n"))]
        return self._emit_packed(units, before, after, path,
                                 joiner="\n", wrap="{inner}")

    def _emit_packed(
            self,
            units: list[tuple[str, int, int]],
            before: str,
            after: str,
            path: str,
            joiner: str,
            wrap: str,
            allow_empty: bool = False,
    ) -> list[FileChunk]:
        """Greedy-pack units into ``wrap`` envelopes; emit each piece.

        Envelope (table header, code fence) counts against ``chunk_chars``;
        TOC (when on) is additive prefix/suffix downstream. Oversized
        units overflow rather than truncate. Multi-piece outputs get
        ``[Part X/N]`` markers; single pieces don't.
        """
        envelope = len(wrap.replace("{inner}", ""))
        budget = max(64, self.chunk_chars - envelope - _PART_RESERVE)
        sep_len = len(joiner)

        parts: list[tuple[str, int, int]] = []
        bucket: list[tuple[str, int, int]] = []
        bucket_chars = 0

        def flush() -> None:
            nonlocal bucket, bucket_chars
            if not bucket:
                return
            inner = joiner.join(t for t, _, _ in bucket)
            parts.append((inner, bucket[0][1], bucket[-1][2]))
            bucket = []
            bucket_chars = 0

        for text, s, e in units:
            if not text and not allow_empty:
                continue
            sep = sep_len if bucket else 0
            if bucket_chars + sep + len(text) > budget:
                flush()
                sep = 0
            bucket.append((text, s, e))
            bucket_chars += sep + len(text)
        flush()

        total = len(parts)
        return [
            self._make_chunk(
                before,
                f"[Part {idx}/{total}]\n\n{wrap.replace('{inner}', inner)}"
                if total > 1 else wrap.replace("{inner}", inner),
                after, s, e, path,
            )
            for idx, (inner, s, e) in enumerate(parts, 1)
        ]

    # -- Emit -------------------------------------------------------------

    def _make_chunk(
            self,
            before: str,
            content: str,
            after: str,
            start_line: int,
            end_line: int,
            path: str,
    ) -> FileChunk:
        """Build one ``FileChunk`` — text is ``before + content + after``
        when ``embed_toc``, otherwise just ``content``."""
        text = _toc_join(before, content, after) if self.embed_toc else content
        return FileChunk(
            id=hash_text(text),
            path=path,
            start_line=start_line,
            end_line=end_line,
            text=text,
        )