"""Markdown documents (model cards, docs served as .md, READMEs): YAML front matter, headings, tables, links, plain text."""
from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Any

import yaml

_FM = re.compile(r"\A---[ \t]*\r?\n(.*?)\r?\n---[ \t]*\r?\n", re.DOTALL)
_HEADING = re.compile(r"^(#{1,6})\s+(.*?)\s*#*\s*$", re.MULTILINE)
_LINK = re.compile(r"\[([^\]]*)\]\((https?://[^)\s]+)\)")
_TABLE_ROW = re.compile(r"^\s*\|(.+)\|\s*$")
_SEP_ROW = re.compile(r"^\s*\|?\s*:?-{2,}:?\s*(\|\s*:?-{2,}:?\s*)*\|?\s*$")


@dataclass
class MarkdownDoc:
    front_matter: dict[str, Any] = field(default_factory=dict)
    body: str = ""
    headings: list[tuple[int, str]] = field(default_factory=list)
    tables: list[dict[str, Any]] = field(default_factory=list)
    links: list[tuple[str, str]] = field(default_factory=list)
    text: str = ""

    def section(self, title_pattern: str) -> str:
        """Body text of the first heading matching `title_pattern` (case-insensitive) up to the next heading of same/higher level."""
        rx = re.compile(title_pattern, re.IGNORECASE)
        lines = self.body.split("\n")
        start = None
        level = 0
        for i, line in enumerate(lines):
            m = _HEADING.match(line)
            if m and start is None and rx.search(m.group(2)):
                start = i + 1
                level = len(m.group(1))
                continue
            if start is not None and m and len(m.group(1)) <= level:
                return "\n".join(lines[start:i]).strip()
        return "\n".join(lines[start:]).strip() if start is not None else ""


def parse_front_matter(text: str) -> tuple[dict[str, Any], str]:
    m = _FM.match(text)
    if not m:
        return {}, text
    try:
        data = yaml.safe_load(m.group(1)) or {}
        if not isinstance(data, dict):
            data = {"_value": data}
    except yaml.YAMLError:
        data = {}
    return data, text[m.end():]


def parse_markdown(text: str) -> MarkdownDoc:
    fm, body = parse_front_matter(text)
    doc = MarkdownDoc(front_matter=fm, body=body)
    doc.headings = [(len(m.group(1)), m.group(2).strip()) for m in _HEADING.finditer(body)]
    doc.links = [(m.group(2), m.group(1)) for m in _LINK.finditer(body)]
    doc.tables = _tables(body)
    doc.text = _to_text(body)
    return doc


def _tables(body: str) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    lines = body.split("\n")
    i = 0
    while i < len(lines):
        if _TABLE_ROW.match(lines[i]) and i + 1 < len(lines) and _SEP_ROW.match(lines[i + 1]):
            headers = _cells(lines[i])
            rows: list[list[str]] = []
            i += 2
            while i < len(lines) and _TABLE_ROW.match(lines[i]):
                rows.append(_cells(lines[i]))
                i += 1
            out.append({"caption": None, "headers": headers, "rows": rows})
        else:
            i += 1
    return out


def _cells(line: str) -> list[str]:
    inner = line.strip()
    inner = inner.removeprefix("|")
    inner = inner.removesuffix("|")
    return [re.sub(r"\*\*|`", "", c).strip() for c in inner.split("|")]


def _to_text(body: str) -> str:
    s = re.sub(r"```.*?```", " ", body, flags=re.DOTALL)
    s = re.sub(r"!\[[^\]]*\]\([^)]*\)", " ", s)
    s = re.sub(r"\[([^\]]*)\]\([^)]*\)", r"\1", s)
    s = re.sub(r"<[^>]+>", " ", s)
    s = re.sub(r"^[#>*\-|\s]+", "", s, flags=re.MULTILINE)
    s = re.sub(r"[ \t]+", " ", s)
    return re.sub(r"\n{3,}", "\n\n", s).strip()


__all__ = ["MarkdownDoc", "parse_front_matter", "parse_markdown"]