#!/usr/bin/env python3
"""Полигон, пакет-7 (A3): нарезка книг стенда в общий JSONL-субстрат для майнинга.

Выход — одна строка на чанк: {"chapter": int, "chunk_idx": int, "source": str}. Это ровно тот
контракт, который `internal/miner.Chunk` ждёт от вызывающего (см. pipeline/mining.go:55-58 —
Chapter/ChunkIdx/NSource), поэтому боевой Go-майнер и мои собственные скрипты по en/ja едят
ОДИН И ТОТ ЖЕ субстрат и их выходы сравнимы.

ЧЕСТНАЯ ОГОВОРКА (пре-регистрирована до прогонов): это НЕ боевой чанкер. Боевой режет по
output-бюджету модели (`internal/chunk`), здесь же — по абзацам до ~target символов внутри главы.
Для майнера различие затрагивает только границы n-грамм и пер-чанковый счёт вхождений, поэтому
эффект замерян отдельно (--target-свип), а не объявлен пренебрежимым.

$0: ни одного сетевого вызова, ни одной модели.
"""
from __future__ import annotations

import argparse
import json
import re
import sys
import unicodedata
import zipfile
from html.parser import HTMLParser
from pathlib import Path

# --- маркеры глав по книгам стенда -------------------------------------------------------------
# Каждое правило — (regex, группа с номером, парсер номера). Правила ДАННЫЕ, а не ветвление по книге
# в коде: добавить книгу = добавить строку.
CJK_NUM = {"〇": 0, "零": 0, "一": 1, "二": 2, "三": 3, "四": 4, "五": 5,
           "六": 6, "七": 7, "八": 8, "九": 9}
JP_FULLWIDTH = {chr(0xFF10 + i): str(i) for i in range(10)}


def cjk_int(s: str) -> int:
    """Китайская/японская запись числа → int (十/百/千 позиционно). Только то, что реально
    встречается в заголовках глав: до 9999."""
    s = s.strip()
    if s.isdigit():
        return int(s)
    if all(c in JP_FULLWIDTH for c in s):
        return int("".join(JP_FULLWIDTH[c] for c in s))
    total, section, digit = 0, 0, 0
    for ch in s:
        if ch in CJK_NUM:
            digit = CJK_NUM[ch]
        elif ch == "十":
            section += (digit or 1) * 10
            digit = 0
        elif ch == "百":
            section += (digit or 1) * 100
            digit = 0
        elif ch == "千":
            section += (digit or 1) * 1000
            digit = 0
        elif ch == "万":
            total += (section + digit) * 10000
            section = digit = 0
        else:
            return -1
    return total + section + digit


EN_ORD = ("one two three four five six seven eight nine ten eleven twelve thirteen fourteen "
          "fifteen sixteen seventeen eighteen nineteen twenty").split()


def en_int(s: str) -> int:
    s = s.strip().lower().replace("-", " ")
    parts = s.split()
    if len(parts) == 1 and parts[0] in EN_ORD:
        return EN_ORD.index(parts[0]) + 1
    if len(parts) == 2 and parts[0] == "twenty" and parts[1] in EN_ORD:
        return 20 + EN_ORD.index(parts[1]) + 1
    if s.isdigit():
        return int(s)
    return -1


RULES = {
    # 蛊真人: `第1节 …` / `第一节 …` — тот же маркер/юнит, что читает боевой langpack (heading.txt: 第 + 章节節回)
    "guzhenren": (re.compile(r"^\s*第\s*([0-9]+|[〇零一二三四五六七八九十百千万]+)\s*[节節章]"), cjk_int),
    # 金瓶梅 (ctext-срез): `===== 第4回 =====`
    "jinpingmei": (re.compile(r"^=+\s*第\s*([0-9]+)\s*回\s*=+"), cjk_int),
    # 異世界魔術師 (syosetu-дамп): `001. 第１話`
    "isekai": (re.compile(r"^\s*[0-9]{3}\.\s*第\s*([0-9０-９]+)\s*話"), cjk_int),
    # Fifty Shades (txt): `CHAPTER ONE` в теле (в оглавлении — `    Chapter One`, отсекается регистром)
    "fifty": (re.compile(r"^CHAPTER\s+([A-Z\-]+)\s*$"), en_int),
    # Empire of the Dawn (epub): заголовки внутри xhtml не структурны — режем по файлам сплита
    "epub-files": (None, None),
}


class _Text(HTMLParser):
    """Минимальный html→text без внешних зависимостей (bs4/ebooklib в venv полигона нет)."""

    def __init__(self) -> None:
        super().__init__(convert_charrefs=True)
        self.parts: list[str] = []
        self._skip = 0

    def handle_starttag(self, tag, attrs):
        if tag in ("script", "style"):
            self._skip += 1
        elif tag in ("p", "div", "br", "h1", "h2", "h3", "h4", "li"):
            self.parts.append("\n")

    def handle_endtag(self, tag):
        if tag in ("script", "style") and self._skip:
            self._skip -= 1
        elif tag in ("p", "div", "h1", "h2", "h3", "h4", "li"):
            self.parts.append("\n")

    def handle_data(self, data):
        if not self._skip:
            self.parts.append(data)

    def text(self) -> str:
        return re.sub(r"\n{3,}", "\n\n", "".join(self.parts))


def read_epub(path: Path) -> list[tuple[int, str]]:
    """EPUB → [(порядковый номер документа, текст)] в порядке spine из content.opf."""
    with zipfile.ZipFile(path) as z:
        opf_name = next((n for n in z.namelist() if n.endswith(".opf")), None)
        order: list[str] = []
        if opf_name:
            opf = z.read(opf_name).decode("utf-8", "replace")
            base = str(Path(opf_name).parent)
            ids = dict(re.findall(r'<item\s[^>]*id="([^"]+)"[^>]*href="([^"]+)"', opf))
            ids.update({i: h for h, i in re.findall(
                r'<item\s[^>]*href="([^"]+)"[^>]*id="([^"]+)"', opf)})
            for idref in re.findall(r'<itemref[^>]*idref="([^"]+)"', opf):
                href = ids.get(idref)
                if not href:
                    continue
                name = href if base in ("", ".") else f"{base}/{href}"
                if name in z.namelist():
                    order.append(name)
        if not order:
            order = sorted(n for n in z.namelist()
                           if n.lower().endswith((".xhtml", ".html", ".htm")))
        raw = []
        for i, name in enumerate(order, 1):
            p = _Text()
            p.feed(z.read(name).decode("utf-8", "replace"))
            t = p.text().strip()
            if t:
                raw.append((i, t))
        # Колонтитулы epub: строка, повторяющаяся в >30% документов сплита, — это шапка/подвал
        # (серия, автор, название), а не текст. Первый прогон без этого фильтра вывел «dawn» и
        # «Vampire Empire» в верхушку кандидатов с рассеянием 125/126 глав — чистый артефакт ингеста.
        from collections import Counter as _C
        docfreq: _C = _C()
        for _, t in raw:
            for ln in {x.strip() for x in t.split("\n") if x.strip()}:
                docfreq[ln] += 1
        boiler = {ln for ln, c in docfreq.items() if c > 0.3 * len(raw)}
        out = []
        for i, t in raw:
            kept = "\n".join(ln for ln in t.split("\n") if ln.strip() not in boiler)
            if kept.strip():
                out.append((i, kept))
        if boiler:
            print(f"  epub: отброшено {len(boiler)} колонтитульных строк "
                  f"(порог >30% документов): {sorted(boiler)[:5]}", file=sys.stderr)
        return out


def paragraphs(body: str) -> list[str]:
    """Абзацы: непустые строки. Мягкая нормализация пробелов, содержимое не трогаем."""
    return [ln.strip() for ln in body.replace("\r\n", "\n").split("\n") if ln.strip()]


def pack(paras: list[str], target: int) -> list[str]:
    """Склейка абзацев в чанки ~target символов (детерминированно, жадно)."""
    chunks, cur, n = [], [], 0
    for p in paras:
        if cur and n + len(p) > target:
            chunks.append("\n".join(cur))
            cur, n = [], 0
        cur.append(p)
        n += len(p)
    if cur:
        chunks.append("\n".join(cur))
    return chunks


def split_text(text: str, rule_name: str, target: int, max_ch: int) -> list[dict]:
    rx, num = RULES[rule_name]
    lines = text.replace("\r\n", "\n").split("\n")
    chapters: list[tuple[int, list[str]]] = []
    cur_no = 0
    cur: list[str] = []
    for ln in lines:
        m = rx.match(ln) if rx else None
        if m:
            n = num(m.group(1))
            if n > 0:
                if cur:
                    chapters.append((cur_no, cur))
                cur_no, cur = n, []
                continue
        cur.append(ln)
    if cur:
        chapters.append((cur_no, cur))
    # Пролог/оглавление (глава 0) отбрасываем: у 蛊真人 это аннотация, у Fifty Shades — оглавление,
    # и в обоих случаях это не тело книги. Решение принято ДО замеров и одинаково для всех книг.
    chapters = [(n, b) for n, b in chapters if n > 0]
    if max_ch:
        chapters = [c for c in chapters if c[0] <= max_ch]
    out = []
    for no, body in chapters:
        for idx, ch in enumerate(pack(paragraphs("\n".join(body)), target)):
            out.append({"chapter": no, "chunk_idx": idx, "source": ch})
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--book", required=True, type=Path)
    ap.add_argument("--rule", required=True, choices=sorted(RULES))
    ap.add_argument("--target", type=int, default=2000, help="целевой размер чанка в символах")
    ap.add_argument("--max-chapters", type=int, default=0)
    ap.add_argument("--out", required=True, type=Path)
    a = ap.parse_args()

    if a.rule == "epub-files":
        docs = read_epub(a.book)
        if a.max_chapters:
            docs = docs[: a.max_chapters]
        rows = []
        for no, body in docs:
            for idx, ch in enumerate(pack(paragraphs(body), a.target)):
                rows.append({"chapter": no, "chunk_idx": idx, "source": ch})
    else:
        raw = a.book.read_text(encoding="utf-8", errors="replace")
        rows = split_text(raw, a.rule, a.target, a.max_chapters)

    a.out.parent.mkdir(parents=True, exist_ok=True)
    with a.out.open("w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    chars = sum(len(r["source"]) for r in rows)
    chaps = len({r["chapter"] for r in rows})
    print(f"{a.out}: chunks={len(rows)} chapters={chaps} chars={chars}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())