RAG-CUT/backend/rag_cut/parsers/text_parser.py

"""Plain-text and markup parsers."""

from __future__ import annotations

import json
import re
from html.parser import HTMLParser
from pathlib import Path

from rag_cut.models import Block, BlockType
from rag_cut.parsers.base import BaseParser


class _HTMLTextExtractor(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self._parts: list[str] = []
        self._heading: tuple[int, str] | None = None
        self._blocks: list[Block] = []
        self._current: list[str] = []
        self._in_heading = False

    def handle_starttag(self, tag: str, attrs) -> None:
        if tag in ("h1", "h2", "h3", "h4", "h5", "h6"):
            self._flush_paragraph()
            self._in_heading = True
            self._heading_level = int(tag[1])

    def handle_endtag(self, tag: str) -> None:
        if tag in ("h1", "h2", "h3", "h4", "h5", "h6"):
            text = "".join(self._current).strip()
            self._current = []
            self._in_heading = False
            if text:
                self._blocks.append(
                    Block(type=BlockType.HEADING, text=text, level=self._heading_level)
                )
        elif tag in ("p", "div", "br", "li"):
            self._flush_paragraph()

    def handle_data(self, data: str) -> None:
        self._current.append(data)

    def _flush_paragraph(self) -> None:
        text = "".join(self._current).strip()
        self._current = []
        if text:
            self._blocks.append(Block(type=BlockType.PARAGRAPH, text=text))

    def get_blocks(self) -> list[Block]:
        self._flush_paragraph()
        return self._blocks


def _parse_markdown(text: str) -> list[Block]:
    blocks: list[Block] = []
    for line in text.splitlines():
        stripped = line.strip()
        if not stripped:
            continue
        m = re.match(r"^(#{1,6})\s+(.+)$", stripped)
        if m:
            blocks.append(
                Block(type=BlockType.HEADING, text=m.group(2).strip(), level=len(m.group(1)))
            )
        else:
            blocks.append(Block(type=BlockType.PARAGRAPH, text=stripped))
    return blocks


def _parse_json(text: str) -> list[Block]:
    data = json.loads(text)
    blocks: list[Block] = []
    if isinstance(data, list):
        for i, item in enumerate(data):
            blocks.append(
                Block(
                    type=BlockType.CODE,
                    text=json.dumps(item, ensure_ascii=False, indent=2),
                    meta={"json_index": i},
                )
            )
    elif isinstance(data, dict):
        for key, value in data.items():
            blocks.append(
                Block(
                    type=BlockType.CODE,
                    text=json.dumps({key: value}, ensure_ascii=False, indent=2),
                    meta={"json_key": key},
                )
            )
    else:
        blocks.append(Block(type=BlockType.PARAGRAPH, text=str(data)))
    return blocks


class TextParser(BaseParser):
    def parse(self, path: Path, assets_dir: Path) -> list[Block]:
        for encoding in ("utf-8-sig", "utf-8", "gbk", "latin-1"):
            try:
                text = path.read_text(encoding=encoding)
                break
            except UnicodeDecodeError:
                continue
        else:
            raise ValueError(f"Cannot decode text file: {path}")

        ext = path.suffix.lower()
        if ext == ".md":
            return _parse_markdown(text)
        if ext in (".html", ".htm"):
            parser = _HTMLTextExtractor()
            parser.feed(text)
            return parser.get_blocks()
        if ext == ".json":
            return _parse_json(text)
        # txt, xml, log — paragraph split on blank lines
        blocks: list[Block] = []
        for para in re.split(r"\n\s*\n", text):
            para = para.strip()
            if para:
                blocks.append(Block(type=BlockType.PARAGRAPH, text=para))
        return blocks