"""Native Markdown IR builder: ``extract_markdown`` output → :class:`IRDoc`.

Input contract: the block dicts and side tables produced by
:func:`lightrag.parser.markdown.extract.extract_markdown`, threaded through
``NativeMarkdownParser.extract`` as ``(blocks, _, metadata)``. ``metadata``
carries the marker→payload side tables under the ``md_*`` keys.

Each block's ``content`` holds self-closing markers (``<mdtable ref=…/>`` /
``<mdequation ref=…/>`` / ``<mddrawing ref=…/>``). This builder rewrites them
into IR placeholder tokens (``{{TBL:k}}`` / ``{{EQ:k}}`` / ``{{IMG:k}}``) and
builds the matching :class:`IRTable` / :class:`IREquation` / :class:`IRDrawing`
from the side tables.

Image bytes are carried in ``metadata["md_assets"]`` (keyed by a stable
identity), so assets are declared as ``AssetSpec(source=bytes)`` and the writer
materializes + deduplicates them. External-link images carry no bytes and are
rendered verbatim through ``IRDrawing.path_override``.
"""

from __future__ import annotations

import itertools
import re
from collections.abc import Callable
from pathlib import Path
from typing import Any

from lightrag.parser._html_table import (
    extract_html_table_info,
    extract_thead_html,
    html_table_inner_body,
    unwrap_html_table,
)
from lightrag.parser.markdown.extract import (
    DRAWING_MARKER_RE,
    EQUATION_MARKER_RE,
    PREFACE_HEADING,
    TABLE_MARKER_RE,
)
from lightrag.sidecar.ir import (
    AssetSpec,
    IRBlock,
    IRDoc,
    IRDrawing,
    IREquation,
    IRPosition,
    IRTable,
)


def _placeholder_keyspace() -> Callable[[str], str]:
    counter = itertools.count(1)
    return lambda prefix: f"{prefix}{next(counter)}"


class NativeMarkdownIRBuilder:
    """Translate ``extract_markdown`` output into an :class:`IRDoc`.

    Stateless — instantiate per call. ``asset_dir_name`` is the relative name
    of ``<base>.blocks.assets/``; the writer applies it as the on-disk prefix
    via ``block_drawing_path_style="with_prefix"``.
    """

    def normalize(
        self,
        blocks: list[dict[str, Any]],
        *,
        document_name: str,
        asset_dir_name: str,
        parse_metadata: dict[str, Any] | None = None,
    ) -> IRDoc:
        meta = parse_metadata or {}
        tables_meta: dict[str, dict] = meta.get("md_tables") or {}
        equations_meta: dict[str, str] = meta.get("md_equations") or {}
        drawings_meta: dict[str, dict] = meta.get("md_drawings") or {}
        assets_meta: dict[str, dict] = meta.get("md_assets") or {}

        next_key = _placeholder_keyspace()
        ir_blocks: list[IRBlock] = []
        assets: list[AssetSpec] = []
        seen_asset_refs: set[str] = set()
        doc_title = ""

        for block in blocks:
            content = block.get("content") or ""
            heading = block.get("heading") or ""
            level = int(block.get("level", 0) or 0)
            parent_headings = list(block.get("parent_headings") or [])

            tables: list[IRTable] = []
            equations: list[IREquation] = []
            drawings: list[IRDrawing] = []

            def _replace_table(match: "re.Match[str]") -> str:
                ref = match.group(1)
                spec = tables_meta.get(ref) or {}
                placeholder = next_key("tb")
                tables.append(_build_table(placeholder, spec))
                return f"{{{{TBL:{placeholder}}}}}"

            def _replace_equation(match: "re.Match[str]") -> str:
                ref = match.group(1)
                latex = equations_meta.get(ref, "")
                placeholder = next_key("eq")
                equations.append(
                    IREquation(placeholder_key=placeholder, latex=latex, is_block=True)
                )
                return f"{{{{EQ:{placeholder}}}}}"

            def _replace_drawing(match: "re.Match[str]") -> str:
                ref = match.group(1)
                spec = drawings_meta.get(ref) or {}
                placeholder = next_key("im")
                drawings.append(
                    _build_drawing(
                        placeholder, spec, assets_meta, assets, seen_asset_refs
                    )
                )
                return f"{{{{IMG:{placeholder}}}}}"

            content = TABLE_MARKER_RE.sub(_replace_table, content)
            content = EQUATION_MARKER_RE.sub(_replace_equation, content)
            content = DRAWING_MARKER_RE.sub(_replace_drawing, content)

            anchor = heading if heading and heading != PREFACE_HEADING else None
            positions = [IRPosition(type="heading", anchor=anchor)]

            ir_blocks.append(
                IRBlock(
                    content_template=content,
                    heading=heading,
                    level=level,
                    parent_headings=parent_headings,
                    positions=positions,
                    tables=tables,
                    drawings=drawings,
                    equations=equations,
                )
            )

            if not doc_title and level == 1 and heading and heading != PREFACE_HEADING:
                doc_title = heading

        if not doc_title:
            doc_title = Path(document_name).stem or document_name

        return IRDoc(
            document_name=document_name,
            document_format=Path(document_name).suffix.lower().lstrip("."),
            doc_title=doc_title,
            split_option={},
            blocks=ir_blocks,
            assets=assets,
            bbox_attributes=None,
        )


def _build_table(placeholder: str, spec: dict) -> IRTable:
    if spec.get("kind") == "html":
        raw = str(spec.get("html") or "")
        html = unwrap_html_table(raw) or None
        body_override = html_table_inner_body(html) or None if html else None
        info = extract_html_table_info(html or "")
        return IRTable(
            placeholder_key=placeholder,
            rows=None,
            html=html,
            num_rows=info.num_rows,
            num_cols=info.num_cols,
            table_header=extract_thead_html(html or ""),
            body_override=body_override,
        )
    # pipe table
    rows = [[str(c) for c in row] for row in (spec.get("rows") or [])]
    header = spec.get("header")
    num_rows = len(rows)
    num_cols = max((len(r) for r in rows), default=0)
    if header:
        num_cols = max(num_cols, max((len(h) for h in header), default=0))
    return IRTable(
        placeholder_key=placeholder,
        rows=rows,
        html=None,
        num_rows=num_rows,
        num_cols=num_cols,
        table_header=header if header else None,
    )


def _build_drawing(
    placeholder: str,
    spec: dict,
    assets_meta: dict[str, dict],
    assets: list[AssetSpec],
    seen_asset_refs: set[str],
) -> IRDrawing:
    fmt = str(spec.get("fmt") or "")
    src = str(spec.get("src") or "")
    if spec.get("kind") == "external":
        return IRDrawing(
            placeholder_key=placeholder,
            asset_ref="",
            fmt=fmt,
            src=src,
            path_override=str(spec.get("url") or src),
        )
    asset_ref = str(spec.get("asset_ref") or "")
    if asset_ref and asset_ref not in seen_asset_refs:
        asset = assets_meta.get(asset_ref) or {}
        assets.append(
            AssetSpec(
                ref=asset_ref,
                suggested_name=str(asset.get("suggested_name") or asset_ref),
                source=asset.get("data"),
            )
        )
        seen_asset_refs.add(asset_ref)
    return IRDrawing(
        placeholder_key=placeholder,
        asset_ref=asset_ref,
        fmt=fmt,
        src=src,
        path_override=None,
    )


__all__ = ["NativeMarkdownIRBuilder"]