"""PDF -> TXT для RAG.

Запуск из корня репо (нужен poppler-utils для pdftotext):
    uv run --with pdfplumber python docs/source/rag/pdf_to_txt.py docs/source/rag/*.pdf --out config/rag/corpus
Сверка результата с PDF:
    python docs/source/rag/check_txt.py docs/source/rag/*.pdf --out config/rag/corpus

- текст вне таблиц, строки склеены в абзацы (в том числе через границу страницы);
- таблицы -> markdown; у таблицы-продолжения на новой странице повторяется шапка,
  строка, разорванная границей страницы, склеивается;
- колонтитулы (повторяющиеся строки вверху/внизу страниц) вырезаются;
- маркеры страниц `--- стр. N ---` для ссылок на источник; абзац, перетёкший на новую
  страницу, склеивается, а место перехода помечается `⟦стр. N⟧`;
- текст с картинок (блок-схемы и т.п.) берётся из figures/<имя pdf>.p<N>.txt
  (ручная транскрипция) и вставляется в конец своей страницы;
- если есть manual/<имя pdf>.txt, берётся он (документы, которые автоматом
  не разбираются, например вёрстка Chromium с выносками).
"""

from __future__ import annotations

import math
import re
import subprocess
import sys
from collections import Counter
from pathlib import Path
from typing import Any

import pdfplumber

FIGURES_DIR = Path(__file__).parent / "figures"
MANUAL_DIR = Path(__file__).parent / "manual"  # вручную выверенные txt: берутся как есть
EDGE_LINES = 4  # сколько строк сверху/снизу страницы проверять на колонтитул
SENTENCE_END = tuple(".:;!?»)")
GRADE_TAIL = re.compile(r"^([ABCАВС]|\(уровень\b.*)$")  # "Уровень убедительности рекомендаций\nA\n(уровень ...)"
BULLET = re.compile(r"^([●•▪◦\-–—*]|\d+[.)]|[а-яa-z][.)]|[IVXА-Я]{1,4}[.)])\s")
HYPHEN = "\x00"  # временная метка переноса "сло-\nво", снимается в resolve_hyphens
SOFT = "\x01"  # временная метка разрыва строки внутри слова в узкой ячейке: "Соверше\nнно"
DEFINITION = re.compile(r"^\S+ [–—] ")  # "рТМС – ..." в списке сокращений
ITEM = re.compile(r"^(\d+\s*[.)\-–]|[А-ЯA-Z]$)")  # "1. СНИЖЕННОЕ НАСТРОЕНИЕ", "И" — пункт, не шапка
OPTION = re.compile(r"^0 □")  # варианты ответа пункта шкалы: "0 □ Отсутствует 1 □ ..."


def norm_edge(line: str) -> str:
    return re.sub(r"\d+", "#", line.strip())


def strip_junk(line: str, junk: set[str]) -> str:
    """Убирает колонтитул: целую строку, удлинённый вариант (на альбомных
    страницах он шире) и прилипший к концу строки ("...на 7Страница 24 из 117")."""
    n = norm_edge(line)
    if n in junk or any(len(j) >= 30 and (n.startswith(j) or j.startswith(n) and len(n) >= 30) for j in junk):
        return ""
    for j in junk:
        if len(j) >= 10 and "#" in j:
            line = re.sub(re.escape(j).replace("\\#", r"\d+") + r"\s*$", "", line)
    return line


def clean_cell(cell: str | None) -> str:
    if not cell:
        return ""
    cell = re.sub(r"(\w)-\s*\n\s*(?=[а-яёa-z])", r"\1" + HYPHEN, cell)  # перенос в узкой ячейке
    cell = re.sub(r"([а-яё])\n(?=[а-яё])", r"\1" + SOFT, cell)
    return re.sub(r"\s*\n\s*", " ", cell).replace("|", "/").strip()


def is_header(row: list[str]) -> bool:
    """Шапка — короткие подписи колонок; у таблиц-опросников её нет и первая
    строка — уже пункт шкалы (его нельзя повторять на страницах-продолжениях)."""
    return not ITEM.match(row[0]) and all(len(c) <= 100 for c in row)


def merge_split_rows(rows: list[list[str]]) -> list[list[str]]:
    """Пункт шкалы, разрезанный на "N. НАЗВАНИЕ" и строку с вариантами "0 □ ...",
    склеивается обратно в одну строку."""
    out: list[list[str]] = []
    for row in rows:
        if out and OPTION.match(row[0]) and "□" not in out[-1][0]:
            out[-1] = [f"{a} {c}".strip() for a, c in zip(out[-1], row)]
        else:
            out.append(row)
    return out


def table_to_md(header: list[str], body: list[list[str]]) -> str:
    out = ["| " + " | ".join(header) + " |", "|" + "---|" * len(header)]
    out += ["| " + " | ".join(r) + " |" for r in body]
    return "\n".join(out)


def reflow(lines: list[str]) -> str:
    """Склеивает строки, разорванные вёрсткой, обратно в абзацы."""
    paras: list[str] = []
    for raw in lines:
        line = raw.strip()
        if not line:
            continue
        prev = paras[-1] if paras else ""
        if prev and GRADE_TAIL.match(line):
            paras[-1] = prev + " " + line
            continue
        joinable = (
            prev
            and line[0].islower()
            and not BULLET.match(line)
            and not DEFINITION.match(line)
        )
        if joinable and prev.endswith("-") and not prev.endswith(" -"):
            paras[-1] = prev[:-1] + HYPHEN + line  # перенос или составное слово
        elif joinable:
            paras[-1] = prev + " " + line
        else:
            paras.append(line)
    return "\n".join(paras)


def region_text(pdf_path: Path, num: int, page: Any, top: float, bottom: float) -> str:
    """Текст горизонтальной полосы страницы. pdftotext надёжнее pdfplumber на
    буквах с диакритикой (у pdfplumber выходит "K o z a r i ć")."""
    y = math.floor(top)
    h = math.ceil(bottom) - y
    if h <= 0:
        return ""
    cmd = ["pdftotext", "-nopgbrk", "-f", str(num), "-l", str(num), "-x", "0", "-y", str(y),
           "-W", str(math.ceil(page.width)), "-H", str(h), str(pdf_path), "-"]
    return subprocess.run(cmd, capture_output=True, text=True, check=True).stdout


def word_bag(text: str) -> Counter[str]:
    return Counter(w.lower() for w in re.findall(r"\w+", text.replace(HYPHEN, "")))


def page_blocks(pdf_path: Path, num: int, page: Any) -> list[tuple[str, Any]]:
    """Текстовые полосы и таблицы страницы сверху вниз."""
    # у таблицы, разорванной страницей, нет нижней (или верхней) рамки — без неё
    # pdfplumber теряет крайнюю строку; дорисовываем линии по концам вертикалей
    ends = {round(y) for e in page.vertical_edges if e["height"] > 5 for y in (e["top"], e["bottom"])}
    settings = {"explicit_horizontal_lines": sorted(ends)}
    tables = []
    for t in sorted(page.find_tables(settings), key=lambda t: t.bbox[1]):
        rows = [[clean_cell(c) for c in row] for row in t.extract()]
        # рамка, внутри которой текст лежит не в ячейках, — не таблица: иначе текст пропадёт
        band = word_bag(region_text(pdf_path, num, page, t.bbox[1], t.bbox[3]))
        if band - word_bag(" ".join(" ".join(r) for r in rows)):
            continue
        tables.append((t.bbox[1], t.bbox[3], rows))
    blocks: list[tuple[str, Any]] = []
    top = 0.0
    for t_top, t_bottom, rows in tables:
        if t_top < top:  # вложенная/перекрывающаяся таблица
            continue
        blocks.append(("text", region_text(pdf_path, num, page, top, t_top)))
        blocks.append(("table", rows))
        top = t_bottom
    blocks.append(("text", region_text(pdf_path, num, page, top, page.height)))
    return blocks


def resolve_hyphens(text: str) -> str:
    """"клинико-\nпсихопатологические" — составное слово, "антихо-\nлинерги-\nческие" —
    перенос. Склеиваем без дефиса, если склейка встречается в документе целиком или
    хотя бы одна половинка — не слово (обрывок слога); иначе оставляем дефис."""
    vocab = {w.lower() for w in re.findall(r"\w+", re.sub(rf"\w*[{HYPHEN}{SOFT}]\w*", " ", text))}

    def fix(m: re.Match[str]) -> str:
        a, b = m.group(1), m.group(2)
        whole = (a + b).lower() in vocab
        compound = a.lower() in vocab and b.lower() in vocab
        return a + b if whole or not compound else f"{a}-{b}"

    pattern = re.compile(rf"(\w+){HYPHEN}(\w+)")
    while HYPHEN in text:
        new = pattern.sub(fix, text, count=0)
        if new == text:
            break
        text = new
    text = text.replace(HYPHEN, "-")

    def fix_soft(m: re.Match[str]) -> str:
        a, b = m.group(1), m.group(2)
        broken = (a + b).lower() in vocab and not (a.lower() in vocab and b.lower() in vocab)
        return a + b if broken else f"{a} {b}"

    return re.sub(rf"(\w+){SOFT}(\w+)", fix_soft, text).replace(SOFT, " ")


def convert(pdf_path: Path) -> str:
    manual = MANUAL_DIR / f"{pdf_path.stem}.txt"
    if manual.exists():
        return manual.read_text()
    with pdfplumber.open(pdf_path) as pdf:
        pages = [page_blocks(pdf_path, n, p) for n, p in enumerate(pdf.pages, 1)]

    # колонтитулы: строки у края страницы, повторяющиеся на >= 40% страниц
    edge: Counter[str] = Counter()
    for page in pages:
        lines = [l for kind, b in page if kind == "text" for l in b.splitlines() if l.strip()]
        edge.update({norm_edge(l) for l in lines[:EDGE_LINES] + lines[-EDGE_LINES:]})
    junk = {k for k, n in edge.items() if n >= max(3, 0.4 * len(pages))}

    # blocks: ["page", N] | ["text", str] | ["table", (header, body)] | ["figure", str]
    blocks: list[list[Any]] = []
    for num, page in enumerate(pages, 1):
        blocks.append(["page", num])
        first_on_page = True
        for kind, b in page:
            prev = next((x for x in reversed(blocks) if x[0] != "page"), None)
            if kind == "text":
                text = reflow([strip_junk(l, junk) for l in b.splitlines()])
                if not text:
                    continue
                if first_on_page and prev and prev[0] == "text" and text[0].islower():
                    # абзац перетёк на новую страницу: дописываем его хвост к началу, а место
                    # перехода помечаем ⟦стр. N⟧ — чанкер по метке ставит верную страницу в цитату
                    head, _, rest = text.partition("\n")
                    prev[1] = prev[1] + f" ⟦стр. {num}⟧ " + head
                    if rest:
                        blocks.append(["text", rest])
                else:
                    blocks.append(["text", text])
            else:
                rows = [r for r in b if any(r)]
                if not rows:
                    continue
                continues = (
                    first_on_page
                    and prev is not None
                    and prev[0] == "table"
                    and len(prev[1][0]) == len(rows[0])
                    and rows[0] != prev[1][0]
                )
                if continues:
                    header, body = prev[1]
                    tail = body if body else [header]
                    if not rows[0][0] or rows[0][0][:1].islower() or (
                        OPTION.match(rows[0][0]) and "□" not in tail[-1][0]
                    ):
                        # строка таблицы разорвана границей страницы
                        tail[-1] = [f"{a} {c}".strip() for a, c in zip(tail[-1], rows[0])]
                        if not body:
                            prev[1] = (tail[0], body)
                        rows = rows[1:]
                    if rows and is_header(header):
                        blocks.append(["table", (header, rows)])
                    elif rows:  # шапки нет — не дублируем первый пункт
                        blocks.append(["table", (rows[0], rows[1:])])
                else:
                    blocks.append(["table", (rows[0], rows[1:])])
            first_on_page = False
        fig = FIGURES_DIR / f"{pdf_path.stem}.p{num}.txt"
        if fig.exists():
            blocks.append(["figure", fig.read_text().strip()])

    out = [f"# {pdf_path.stem}", f"Источник: {pdf_path.name}, страниц: {len(pages)}"]
    for kind, b in blocks:
        if kind == "page":
            out.append(f"--- стр. {b} ---")
        elif kind == "text":
            out.append(b)
        elif kind == "table":
            header, *body = merge_split_rows([b[0], *b[1]])
            out.append(table_to_md(header, body))
        else:
            out.append("[Изображение: текст перенесён вручную]\n" + b)
    return resolve_hyphens("\n\n".join(out)) + "\n"


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("pdfs", nargs="+", type=Path)
    parser.add_argument("--out", type=Path, help="папка для txt (по умолчанию рядом с pdf)")
    args = parser.parse_args()
    for src in args.pdfs:
        dst = (args.out / src.name if args.out else src).with_suffix(".txt")
        dst.write_text(convert(src))
        print(f"{src.name} -> {dst}")