#!/usr/bin/env python3
"""Вычислить байтовое смещение в дампе Википедии по уже залитым статьям.

Нужен один раз при переходе на multistream-дамп: позиция продолжения сменила
смысл — раньше это был номер статьи (и прогон перечитывал дамп с начала),
теперь байтовое смещение блока.

Как считается: берём максимальный page_id среди залитых статей и находим в
индексе дампа блок, которому он принадлежит. Страницы в дампе идут по
возрастанию page_id, поэтому всё, что дальше этого блока, ещё не залито.

Запуск (в контейнере worker-indexer):
    docker compose -f docker-compose.prod.yml exec -T worker-indexer python - \\
        < scripts/ops/wikipedia_resume_offset.py            # показать
    ... python - --apply < scripts/ops/wikipedia_resume_offset.py   # записать
"""

import argparse
import bz2
import os

DEFAULT_INDEX = "/parsers/ruwiki-index.txt.bz2"


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--apply", action="store_true", help="записать смещение в источник")
    ap.add_argument("--index", default=DEFAULT_INDEX, help="путь к индексу дампа")
    args = ap.parse_args()

    from app.db import db_session
    from sqlalchemy import text

    if not os.path.exists(args.index):
        raise SystemExit(f"индекс не найден: {args.index}")

    with db_session() as s:
        row = s.execute(text("""
            SELECT max(split_part(ext_id, ':', 2)::bigint)
            FROM documents WHERE source = 'wikipedia_ru'
        """)).first()
    max_pageid = int(row[0]) if row and row[0] else 0
    print(f"максимальный page_id среди залитых: {max_pageid}")

    if not max_pageid:
        print("залитых статей нет — начинать с нуля")
        return

    # Индекс отсортирован по смещению, page_id внутри растут: ищем блок,
    # в котором лежит наш максимальный, и берём его смещение
    offset = 0
    found_title = ""
    with bz2.open(args.index, "rt", encoding="utf-8") as fh:
        for line in fh:
            parts = line.split(":", 2)
            if len(parts) < 3:
                continue
            off, pid = int(parts[0]), int(parts[1])
            if pid > max_pageid:
                break
            offset, found_title = off, parts[2].strip()

    print(f"блок с этой статьёй начинается на байте {offset} ({found_title[:50]!r})")

    if not args.apply:
        print("\nзапустите с --apply, чтобы записать смещение в источник")
        return

    with db_session() as s:
        s.execute(text("""
            UPDATE parse_sources SET resume_token = :off WHERE source_type = 'wikipedia_ru'
        """), {"off": str(offset)})
        s.commit()
    print(f"смещение {offset} записано в источник")


if __name__ == "__main__":
    main()