LLM-UI/vwts_scraper/run.py

"""Основной цикл: ThreadPoolExecutor + диапазоны страниц → мерж."""

import sys, logging, shutil, json
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from urllib.parse import urlparse
from .config import OUTPUT_DIR, WORKERS, BAN_TEST, DELAY
from .fetch import get, make_session
from .paginate import count, verify_last
from .parse import parse_topics, parse_posts, parse_tags
from .state import load, save, merge_states
from .output import save_topic
from .throttle import Throttle, BanDetector

log = logging.getLogger(__name__)


# ── Глобальный throttle (устанавливается в run()) ─────────────────────
_throttle = None
import fetch as _fetch


def _worker(section_url: str, section_slug: str, section_name: str,
            pages: range, worker_id: int):
    """Один поток: обходит свой диапазон страниц, пишет в worker_N/."""
    session = make_session()
    st = {"topics_done": {}, "pages_done": [], "file_index": 0, "topic_count": 0}
    base = f"{urlparse(section_url).scheme}://{urlparse(section_url).netloc}"
    consecutive_404 = 0

    # Папка воркера
    worker_dir = OUTPUT_DIR / section_slug / f"worker_{worker_id}"
    worker_dir.mkdir(parents=True, exist_ok=True)

    for pg in pages:
        url = section_url if pg == 1 else f"{section_url}page-{pg}"
        log.info(f"[W{worker_id}] 📄 стр.{pg}")

        page_soup = get(url, session)
        if not page_soup:
            consecutive_404 += 1
            if consecutive_404 >= 3:
                log.info(f"[W{worker_id}] ⏹️  3 пустых — завершаем")
                break
            continue
        consecutive_404 = 0

        topics = parse_topics(page_soup)
        log.info(f"[W{worker_id}]    Тем: {len(topics)}")

        for i, tp in enumerate(topics, 1):
            tid = tp["id"]
            if str(tid) in st["topics_done"]:
                continue

            log.info(f"[W{worker_id}]    [{i}/{len(topics)}] #{tid}: {tp['title'][:80]}")
            topic_soup = get(tp["url"], session)
            if not topic_soup:
                st["topics_done"][str(tid)] = tp["title"]
                continue

            topic_pages = count(topic_soup)
            posts, tags = [], parse_tags(topic_soup)

            for tpp in range(1, topic_pages + 1):
                if tpp == 1:
                    posts += parse_posts(topic_soup)
                else:
                    tp2 = get(f"{base}/forum/topic/{tid}/page-{tpp}", session)
                    if tp2:
                        posts += parse_posts(tp2)

            log.info(f"[W{worker_id}]      {len(posts)} постов ({topic_pages} стр.), теги: {tags}")
            save_topic(section_slug, section_name, tp, posts, tags, st,
                       worker_dir=worker_dir)
            st["topics_done"][str(tid)] = tp["title"]

        st["pages_done"].append(pg)

    # Сохраняем состояние воркера
    _state_file = worker_dir / "state.json"
    _state_file.write_text(json.dumps(st, ensure_ascii=False, indent=2),
                           encoding="utf-8")
    session.close()
    return worker_id, st["topic_count"]


def _merge_workers(section_slug: str):
    """Собрать part_*.jsonl из всех worker_N/ в корень раздела с единой нумерацией."""
    root = OUTPUT_DIR / section_slug

    # Собираем все part_*.jsonl из worker_N/
    all_parts = []
    for wdir in sorted(root.glob("worker_*")):
        if wdir.is_dir():
            for f in sorted(wdir.glob("part_*.jsonl")):
                all_parts.append(f)

    log.info(f"🔗 Мерж {len(all_parts)} файлов из {len(list(root.glob('worker_*')))} воркеров")

    # Переименовываем с единой нумерацией
    for idx, src in enumerate(all_parts):
        dst = root / f"part_{idx:04d}.jsonl"
        shutil.move(str(src), str(dst))

    # Удаляем пустые worker-папки
    for wdir in root.glob("worker_*"):
        if wdir.is_dir():
            try:
                wdir.rmdir()
            except OSError:
                pass

    # Пишем объединённый state.json
    merge_states(section_slug)

    log.info(f"✅ Мерж: {len(all_parts)} файлов → {root}")


def run(section_url: str, workers: int = WORKERS,
        no_delay: bool = False, ban_test: bool = True):
    """Главный цикл с многопоточностью.

    Args:
        section_url: полный URL раздела
        workers: количество потоков
        no_delay: True = без координации задержек
        ban_test: True = проверить бан перед стартом
    """
    section_slug = section_url.rstrip("/").split("/")[-1]

    # ── Автоопределение бана ──────────────────────────────────────────
    throttle_enabled = not no_delay
    if ban_test and throttle_enabled and workers > 1:
        if not BanDetector.test(section_url, BAN_TEST):
            throttle_enabled = False

    global _throttle
    _throttle = Throttle(DELAY, enabled=throttle_enabled)
    _fetch.throttle = _throttle

    log.info(f"⚙️  Потоков: {workers}, координация: {'вкл' if throttle_enabled else 'выкл'}")

    # ── Страница 1: название + количество страниц ─────────────────────
    soup = get(section_url, make_session())
    if not soup:
        log.error("❌ Сайт недоступен")
        sys.exit(1)

    section_name = soup.select_one("h1")
    section_name = section_name.text.strip() if section_name else section_slug

    total_pages = count(soup)
    total_pages = verify_last(section_url, total_pages, make_session())
    log.info(f"📋 {section_name} | страниц: {total_pages}")

    # ── Диапазоны страниц ─────────────────────────────────────────────
    if workers > total_pages:
        workers = total_pages

    base_size = total_pages // workers
    remainder = total_pages % workers
    ranges = []
    start = 1
    for w in range(workers):
        size = base_size + (1 if w < remainder else 0)
        ranges.append(range(start, start + size))
        start += size

    log.info(f"📦 Диапазоны: {[f'{r.start}-{r[-1]}' for r in ranges]}")

    # ── Запуск потоков ────────────────────────────────────────────────
    total_topics = 0
    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = [
            pool.submit(_worker, section_url, section_slug, section_name,
                        rng, i)
            for i, rng in enumerate(ranges)
        ]
        for fut in as_completed(futures):
            wid, cnt = fut.result()
            total_topics += cnt
            log.info(f"[W{wid}] ✅ завершён: {cnt} тем")

    # ── Мерж ──────────────────────────────────────────────────────────
    _merge_workers(section_slug)

    # ── Итог ──────────────────────────────────────────────────────────
    log.info("=" * 60)
    log.info(f"✅ '{section_name}' — {total_topics} тем "
             f"в {len(list((OUTPUT_DIR/section_slug).glob('part_*.jsonl')))} частях")
    for f in sorted((OUTPUT_DIR / section_slug).glob("part_*.jsonl")):
        log.info(f"   📄 {f.name} ({f.stat().st_size/1024/1024:.1f} MB)")
    log.info("=" * 60)