"""Датасет для детектора строк из born-digital PDF.

Текстовый слой даёт точные рамки строк бесплатно и в нужном домене — это
разметка, которую иначе пришлось бы рисовать руками. Язык не важен: детектор
ищет строку, а не буквы, поэтому берём все строки подряд, включая русские.

Выпуски, попавшие в оценочную выборку распознавателя, исключаются по имени:
детектор не должен видеть страницы, на которых потом меряется сквозное
качество.
"""
import argparse, json, os, re
from concurrent.futures import ProcessPoolExecutor

import fitz
from PIL import Image

# выпуски, из которых нарезана real_eval (найдены сопоставлением текста),
# плюс весь 2011 год — его выпуск в архиве не опознался однозначно
HELD_OUT = re.compile(r"№ 47 \(2009\)|№ 1 \(2019\)|№ 1 \(2021\)|"
                      r"1 № \(2013\)|№ 8 \(2017\)|№ 16 \(2015\)|\(2011\)")


def page_boxes(page, scale):
    """Рамки строк текстового слоя в пикселях отрендеренной страницы."""
    out = []
    for blk in page.get_text("dict")["blocks"]:
        for ln in blk.get("lines", []):
            if not any(s["text"].strip() for s in ln["spans"]):
                continue
            x0, y0, x1, y1 = ln["bbox"]
            w, h = (x1 - x0) * scale, (y1 - y0) * scale
            # вырожденные и гигантские рамки — мусор вёрстки, не строки
            if w < 6 or h < 4 or h > 400 or w / h > 200:
                continue
            out.append([round(x0 * scale, 1), round(y0 * scale, 1),
                        round(x1 * scale, 1), round(y1 * scale, 1)])
    return out


def one_pdf(args):
    path, out_dir, dpi, max_pages, quality = args
    tag = re.sub(r"[^0-9A-Za-z]+", "_", os.path.splitext(os.path.basename(path))[0])[:48]
    rows, scale = [], dpi / 72.0
    try:
        doc = fitz.open(path)
    except Exception as e:
        return tag, 0, f"не открылся: {e}"
    n = len(doc) if not max_pages else min(len(doc), max_pages)
    for pno in range(n):
        page = doc[pno]
        boxes = page_boxes(page, scale)
        if len(boxes) < 20:          # обложки и полосы без текстового слоя
            continue
        pm = page.get_pixmap(dpi=dpi)
        im = Image.frombytes("RGB", (pm.width, pm.height), pm.samples).convert("L")
        name = f"{tag}_p{pno + 1:03d}.jpg"
        im.save(os.path.join(out_dir, name), quality=quality, optimize=True)
        rows.append({"f": name, "w": im.width, "h": im.height, "b": boxes})
    doc.close()
    if rows:
        with open(os.path.join(out_dir, f"_part_{tag}.jsonl"), "w", encoding="utf-8") as fh:
            for r in rows:
                fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    return tag, len(rows), ""


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True, help="каталог с PDF")
    ap.add_argument("--out", required=True)
    ap.add_argument("--dpi", type=int, default=150)
    ap.add_argument("--pages", type=int, default=0, help="0 = все страницы выпуска")
    ap.add_argument("--quality", type=int, default=85)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--keep-held-out", action="store_true",
                    help="собрать, наоборот, только отложенные выпуски")
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)

    files = sorted(f for f in os.listdir(a.src) if f.lower().endswith(".pdf"))
    picked = [os.path.join(a.src, f) for f in files
              if bool(HELD_OUT.search(f)) == a.keep_held_out]
    print(f"PDF всего {len(files)}, берём {len(picked)}", flush=True)

    tasks = [(p, a.out, a.dpi, a.pages, a.quality) for p in picked]
    done = pages = 0
    with ProcessPoolExecutor(max_workers=a.workers) as ex:
        for tag, k, err in ex.map(one_pdf, tasks):
            done += 1
            pages += k
            if err:
                print(f"  [{done}/{len(tasks)}] {tag}: {err}", flush=True)
            elif done % 10 == 0:
                print(f"  [{done}/{len(tasks)}] страниц собрано {pages}", flush=True)

    # склеиваем части: каждый воркер пишет свой файл, общий индекс собираем в конце
    with open(os.path.join(a.out, "index.jsonl"), "w", encoding="utf-8") as fh:
        for name in sorted(os.listdir(a.out)):
            if not name.startswith("_part_"):
                continue
            p = os.path.join(a.out, name)
            with open(p, encoding="utf-8") as src:
                fh.write(src.read())
            os.remove(p)
    with open(os.path.join(a.out, "index.jsonl"), encoding="utf-8") as fh:
        rows = sum(1 for _ in fh)
    print(f"готово: страниц {rows} в {a.out}")


if __name__ == "__main__":
    main()