El pipeline: un YAML, nueve pasos, siete subcomandos

"""``PipelineConfig``: every data step's config under one YAML (``configs/data/pipeline.yaml``)."""

from __future__ import annotations

from pathlib import Path

from pydantic import Field, model_validator

from rukh import paths
from rukh.config import BaseConfig, load_yaml
from rukh.data.db import DuckDbConfig
from rukh.data.elite import EliteConfig
from rukh.data.elo_bins import EloBinsConfig
from rukh.data.evals import EvalsConfig
from rukh.data.pairs import PairsConfig
from rukh.data.positions import PositionsConfig
from rukh.data.publish import PublishConfig
from rukh.data.puzzles import PuzzlesConfig
from rukh.data.uci import UciConfig

class TokenizeConfig(BaseConfig):
    """Where games come from, where packed tokens and tokenizer artifacts go."""

    uci_dir: str = "data/uci"
    out_dir: str = "data/tokens"
    artifacts_dir: str = "artifacts/tokenizer"
    web_artifacts_dir: str = "artifacts/web"
    fixture_pgn: str = "tests/fixtures/games.pgn"
    train_month: str = "2025-01"
    val_month: str = "2025-02"
    stats_games: str = "data/uci/year=2025/month=01/games.parquet"
    bpe_vocab_size: int = Field(default=4096, ge=100)
    bpe_train_games: int = Field(default=200_000, ge=1)
    stats_n_games: int = Field(default=20_000, ge=1)
    max_len: int = Field(default=200, ge=8)

DUCKDB_STEPS = ("positions", "evals", "puzzles", "elo_bins")


class PipelineConfig(BaseConfig):
    """All step configs; each step reads only its own section.

    ``duckdb:`` is shared: the top-level settings are copied into every step that opens a
    DuckDB connection unless that step declares its own.
    """

    duckdb: DuckDbConfig = Field(default_factory=DuckDbConfig)
    uci: UciConfig = Field(default_factory=UciConfig)
    tokenize: TokenizeConfig = Field(default_factory=TokenizeConfig)
    positions: PositionsConfig = Field(default_factory=PositionsConfig)
    evals: EvalsConfig = Field(default_factory=EvalsConfig)
    puzzles: PuzzlesConfig = Field(default_factory=PuzzlesConfig)
    pairs: PairsConfig = Field(default_factory=PairsConfig)
    elite: EliteConfig = Field(default_factory=EliteConfig)
    elo_bins: EloBinsConfig = Field(default_factory=EloBinsConfig)
    publish: PublishConfig = Field(default_factory=PublishConfig)

    @model_validator(mode="after")
    def _share_duckdb(self) -> PipelineConfig:
        default = DuckDbConfig()
        for name in DUCKDB_STEPS:
            step = getattr(self, name)
            if step.duckdb == default:
                setattr(self, name, step.model_copy(update={"duckdb": self.duckdb}))
        return self

def default_config_path() -> Path:
    return paths.package_root() / "configs" / "data" / "pipeline.yaml"


def load_pipeline(path: Path | None = None) -> PipelineConfig:
    """Load ``pipeline.yaml`` (the shipped one when ``path`` is ``None``)."""
    return load_yaml(path or default_config_path(), PipelineConfig)

# Every step of `rukh data ...` after `fetch`. Paths are relative to the repo root (RUKH_HOME).
duckdb:               # applied to every DuckDB connection of the pipeline
  memory_limit: 32GB
  threads: 0          # 0 = cpu_count()
uci:
  raw_dir: data/raw
  out_dir: data/uci
  min_plies: 20
  max_plies: 300
  workers: 0          # 0 = cpu_count() - 1
tokenize:
  uci_dir: data/uci
  out_dir: data/tokens
  artifacts_dir: artifacts/tokenizer
  web_artifacts_dir: artifacts/web
  fixture_pgn: tests/fixtures/games.pgn
  train_month: "2025-01"
  val_month: "2025-02"
  stats_games: data/uci/year=2025/month=01/games.parquet
  bpe_vocab_size: 4096
  bpe_train_games: 200000
  stats_n_games: 20000
  max_len: 200
positions:
  uci_dir: data/uci
  month: "2025-01"
  out_dir: data/positions
  n_games: 300000
  max_positions: 5000000
  workers: 0
evals:
  dataset: Lichess/chess-position-evaluations
  remote_pattern: "hf://datasets/{dataset}/data/data_{index:04d}.parquet"

  n_files: 20
  positions: data/positions/positions.parquet
  out_dir: data/evals
puzzles:
  dataset: Lichess/chess-puzzles
  remote_glob: "hf://datasets/{dataset}/**/*.parquet"
  out_dir: data/puzzles
  max_rating_deviation: 100
  min_plays: 100
  n_test: 2000
  n_train: 50000
  seed: 42
pairs:
  positions_eval: data/evals/positions-eval.parquet
  out_dir: data/pairs
  min_delta_cp: 100
  max_per_phase: 200000
  seed: 42
elite:
  base_url: https://database.nikonoel.fr
  months: ["2025-01", "2025-02"]
  out_dir: data/elite
  min_plies: 20
  max_plies: 300
  workers: 0
elo_bins:
  uci_dir: data/uci
  out_dir: data/elo-bins
  n_per_bin: 50000
  seed: 42
publish:
  owner: chorcat
  publish_dir: data/publish

PipelineOption = Annotated[
    Path | None,
    typer.Option(
        "--config",
        exists=True,
        dir_okay=False,
        readable=True,
        help="Pipeline YAML (default: configs/data/pipeline.yaml).",
    ),
]


def _echo_written(written: list[str]) -> None:
    typer.echo("written:")
    for path in written:
        typer.echo(f"  {path}")

@data_app.command("uci")
def data_uci(
    config: PipelineOption = None,
    workers: Annotated[
        int | None, typer.Option("--workers", help="Override the pool size from the config.")
    ] = None,
    as_json: Annotated[
        bool, typer.Option("--json", help="Print the manifest as JSON only.")
    ] = False,
) -> None:
    """Convert fetched SAN movetext to legal UCI games, one parquet per month."""
    from rukh.data.pipeline import load_pipeline
    from rukh.data.uci import run

    cfg = load_pipeline(config).uci
    if workers is not None:
        cfg = cfg.model_copy(update={"workers": workers})
    manifest = run(cfg)
    if as_json:
        typer.echo(manifest.model_dump_json(indent=2))
        return
    typer.echo(f"months:   {', '.join(manifest.months)}")
    for month, count in manifest.counts.items():
        typer.echo(f"  {month}: {count} games kept")
    typer.echo(f"manifest: {cfg.out_dir}/manifest.json")

def _echo_manifest(manifest: object, as_json: bool, out_dir: str) -> None:
    from rukh.data.manifest import Manifest

    assert isinstance(manifest, Manifest)
    if as_json:
        typer.echo(manifest.model_dump_json(indent=2))
        return
    typer.echo(f"dataset:  {manifest.dataset}")
    if manifest.months:
        typer.echo(f"months:   {', '.join(manifest.months)}")
    typer.echo("counts:")
    for key, value in manifest.counts.items():
        typer.echo(f"  {key}: {value}")
    typer.echo("files:")
    for file in manifest.files:
        typer.echo(f"  {file.path} ({file.bytes} bytes)")
    typer.echo(f"manifest: {out_dir}/manifest.json")


JsonOption = Annotated[bool, typer.Option("--json", help="Print the manifest as JSON.")]


@data_app.command("positions")
def data_positions(config: PipelineOption = None, as_json: JsonOption = False) -> None:
    """Walk sampled games into deduplicated 4-field FENs with phase labels."""
    from rukh.data.pipeline import load_pipeline
    from rukh.data.positions import run

    cfg = load_pipeline(config).positions
    _echo_manifest(run(cfg), as_json, cfg.out_dir)

@data_app.command("evals")
def data_evals(
    config: PipelineOption = None,
    files: Annotated[
        str | None,
        typer.Option("--files", help="Remote file indices, e.g. 0-19, 3 or 0,2,5-7 (default all)."),
    ] = None,
    as_json: JsonOption = False,
) -> None:
    """Semi-join the positions with the remote Stockfish evaluations (resumable per file)."""
    from rukh.data.evals import run
    from rukh.data.pipeline import load_pipeline

    cfg = load_pipeline(config).evals
    _echo_manifest(run(cfg, files=files), as_json, cfg.out_dir)


@data_app.command("puzzles")
def data_puzzles(config: PipelineOption = None, as_json: JsonOption = False) -> None:
    """Filter Lichess puzzles and split them by difficulty band with a fixed seed."""
    from rukh.data.pipeline import load_pipeline
    from rukh.data.puzzles import run

    cfg = load_pipeline(config).puzzles
    _echo_manifest(run(cfg), as_json, cfg.out_dir)


@data_app.command("pairs")
def data_pairs(config: PipelineOption = None, as_json: JsonOption = False) -> None:
    """Build phase-balanced DPO pairs from the multi-PV evaluations."""
    from rukh.data.pairs import run
    from rukh.data.pipeline import load_pipeline

    cfg = load_pipeline(config).pairs
    _echo_manifest(run(cfg), as_json, cfg.out_dir)


@data_app.command("elite")
def data_elite(config: PipelineOption = None, as_json: JsonOption = False) -> None:
    """Download the Lichess Elite Database months and convert them to UCI."""
    from rukh.data.elite import run
    from rukh.data.pipeline import load_pipeline

    cfg = load_pipeline(config).elite
    _echo_manifest(run(cfg), as_json, cfg.out_dir)


@data_app.command("elo-bins")
def data_elo_bins(config: PipelineOption = None, as_json: JsonOption = False) -> None:
    """Sample up to n_per_bin games per 100-Elo bin of the average rating."""
    from rukh.data.elo_bins import run
    from rukh.data.pipeline import load_pipeline

    cfg = load_pipeline(config).elo_bins
    _echo_manifest(run(cfg), as_json, cfg.out_dir)

"""Lichess Elite Database (2500+ vs 2300+, no bullet): download, extract, convert to UCI."""

from __future__ import annotations

import re
import zipfile
from collections.abc import Iterator
from pathlib import Path

import pyarrow.parquet as pq
from pydantic import Field

from rukh.config import BaseConfig
from rukh.data.manifest import FileHash, Manifest
from rukh.data.parallel import run_batches
from rukh.data.uci import (
    OUT_SCHEMA,
    RAW_COLUMNS,
    convert_batch,
    resolve_workers,
    sha256_file,
    write_results,
)
from rukh.paths import resolve

GAMES_FILE = "games.parquet"
BATCH_GAMES = 20_000
_HEADER_RE = re.compile(r'^\[(\w+)\s+"(.*)"\]\s*$')


class EliteConfig(BaseConfig):
    """Which monthly zips to fetch and the ply bounds of the conversion."""

    base_url: str = "https://database.nikonoel.fr"
    months: list[str] = Field(default=["2025-01", "2025-02"], min_length=1)
    out_dir: str = "data/elite"
    min_plies: int = Field(default=20, ge=0)
    max_plies: int = Field(default=300, ge=1)
    workers: int = Field(default=0, ge=0)

def zip_name(month: str) -> str:
    return f"lichess_elite_{month}.zip"


def zip_url(cfg: EliteConfig, month: str) -> str:
    return f"{cfg.base_url.rstrip('/')}/{zip_name(month)}"


def _download(url: str, dest: Path) -> None:
    """Stream ``url`` to ``dest`` (tests monkeypatch this)."""
    import requests

    dest.parent.mkdir(parents=True, exist_ok=True)
    tmp = dest.with_suffix(dest.suffix + ".part")
    with requests.get(url, stream=True, timeout=120) as response:
        response.raise_for_status()
        with tmp.open("wb") as fh:
            for chunk in response.iter_content(chunk_size=1 << 20):
                fh.write(chunk)
    tmp.replace(dest)


def extract_pgn(zip_path: Path, dest_dir: Path) -> Path:
    """Extract the (single) ``.pgn`` member of the zip; returns its path."""
    with zipfile.ZipFile(zip_path) as archive:
        members = [m for m in archive.namelist() if m.lower().endswith(".pgn")]
        if len(members) != 1:
            raise ValueError(f"{zip_path.name}: expected one .pgn member, found {members}")
        archive.extract(members[0], dest_dir)
    return dest_dir / members[0]

def iter_pgn_rows(pgn_path: Path, month: str) -> Iterator[dict[str, object]]:
    """Yield raw rows (``RAW_COLUMNS``) from a PGN file with a light line parser.

    A new game starts when a header follows movetext or repeats a tag of the current block, so
    a game with headers but no moves is dropped instead of polluting the next one. Without a
    ``Site`` tag the game id falls back to ``month:index`` (every game still gets its own id);
    a missing or unparsable Elo is recorded as ``None``, never as 0.
    """
    headers: dict[str, str] = {}
    movetext: list[str] = []
    in_moves = False
    index = 0

    def flush(index: int) -> dict[str, object] | None:
        if not headers or not movetext:
            return None
        return {
            "Site": headers.get("Site") or f"{month}:{index}",
            "movetext": " ".join(movetext),
            "WhiteElo": _int(headers.get("WhiteElo")),
            "BlackElo": _int(headers.get("BlackElo")),
            "Result": headers.get("Result", "*"),
            "TimeControl": headers.get("TimeControl"),
            "UTCDate": headers.get("UTCDate", headers.get("Date")),
            "ECO": headers.get("ECO"),
            "month": month,
        }

    with pgn_path.open(encoding="utf-8", errors="replace") as fh:
        for line in fh:
            line = line.strip()
            if line.startswith("["):
                match = _HEADER_RE.match(line)
                if match is None:
                    continue
                if in_moves or match.group(1) in headers:
                    row = flush(index)
                    if row is not None:
                        yield row
                    index += 1
                    headers, movetext, in_moves = {}, [], False
                headers[match.group(1)] = match.group(2)
            elif line:
                in_moves = True
                movetext.append(line)
    row = flush(index)
    if row is not None:
        yield row

def _int(value: str | None) -> int | None:
    """``None`` when the header is missing or not a number (never a fabricated 0)."""
    try:
        return int(value) if value is not None else None
    except ValueError:
        return None


def _batches(rows: Iterator[dict[str, object]]) -> Iterator[list[dict[str, object]]]:
    batch: list[dict[str, object]] = []
    for row in rows:
        for name in RAW_COLUMNS:
            row.setdefault(name, None)
        batch.append(row)
        if len(batch) >= BATCH_GAMES:
            yield batch
            batch = []
    if batch:
        yield batch

def convert_pgns(pgns: list[tuple[str, Path]], dst: Path, cfg: EliteConfig) -> dict[str, int]:
    """Convert every ``(month, pgn)`` into one UCI parquet with the ``uci.py`` schema."""
    totals = {"rows": 0, "kept": 0, "illegal": 0, "short": 0, "long": 0}
    workers = resolve_workers(cfg.workers)

    def jobs() -> Iterator[tuple[list[dict[str, object]], int, int]]:
        for month, pgn in pgns:
            for batch in _batches(iter_pgn_rows(pgn, month)):
                yield batch, cfg.min_plies, cfg.max_plies

    dst.parent.mkdir(parents=True, exist_ok=True)
    with pq.ParquetWriter(dst, OUT_SCHEMA, compression="zstd") as writer:
        write_results(writer, run_batches(convert_batch, jobs(), workers), totals)
    return totals

def run(cfg: EliteConfig) -> Manifest:
    """Download each month's zip (once), extract, convert and write the manifest."""
    out_dir = resolve(cfg.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    files: list[FileHash] = []
    sources: dict[str, str] = {}
    pgns: list[tuple[str, Path]] = []
    for month in cfg.months:
        zip_path = out_dir / zip_name(month)
        url = zip_url(cfg, month)
        if not zip_path.is_file():
            _download(url, zip_path)
        sources[month] = url
        files.append(
            FileHash(
                path=zip_path.name, sha256=sha256_file(zip_path), bytes=zip_path.stat().st_size
            )
        )
        pgns.append((month, extract_pgn(zip_path, out_dir)))
    target = out_dir / GAMES_FILE
    totals = convert_pgns(pgns, target, cfg)
    files.append(FileHash(path=GAMES_FILE, sha256=sha256_file(target), bytes=target.stat().st_size))
    manifest = Manifest(
        dataset="Lichess Elite Database (database.nikonoel.fr)",
        months=list(cfg.months),
        filters={
            "source_urls": sources,
            "min_plies": cfg.min_plies,
            "max_plies": cfg.max_plies,
            "legal_only": True,
            "conversion": totals,
        },
        counts={"games": totals["kept"]},
        files=files,
    )
    (out_dir / "manifest.json").write_text(manifest.model_dump_json(indent=2) + "\n", "utf-8")
    return manifest

"""Tests for rukh.data.elo_bins: bin formula, per-bin cap and determinism."""

from __future__ import annotations

import shutil
from pathlib import Path

import polars as pl
import pytest

from rukh.data.elo_bins import EloBinsConfig, run

pytestmark = pytest.mark.unit


def _stage(rukh_home: Path, repo_root: Path) -> None:
    src = repo_root / "tests" / "fixtures" / "games.parquet"
    for month in ("01", "02"):
        target = rukh_home / "data" / "uci" / "year=2025" / f"month={month}" / "games.parquet"
        target.parent.mkdir(parents=True)
        shutil.copy(src, target)


def test_run_caps_each_bin(rukh_home: Path, repo_root: Path) -> None:
    _stage(rukh_home, repo_root)
    manifest = run(EloBinsConfig(n_per_bin=2, seed=1))
    frame = pl.read_parquet((rukh_home / "data" / "elo-bins" / "games.parquet").as_posix())
    assert "bin" in frame.columns
    per_bin = frame.group_by("bin").len()
    assert per_bin["len"].max() <= 2
    expected = ((pl.col("white_elo") + pl.col("black_elo")) // 2) // 100 * 100
    assert frame.filter(pl.col("bin") != expected).height == 0
    assert manifest.months == ["2025-01", "2025-02"]
    assert all(v <= 2 for v in manifest.counts.values())
    assert manifest.filters["n_per_bin"] == 2

def test_run_handles_ids_at_the_edge_of_int64(rukh_home: Path, repo_root: Path) -> None:
    """``game_id`` is a signed 64-bit hash: adding the seed to it would overflow."""
    src = pl.read_parquet((repo_root / "tests" / "fixtures" / "games.parquet").as_posix())
    edges = [2**63 - 1, -(2**63), 2**63 - 2, -(2**63) + 1]
    ids = (edges * ((src.height // len(edges)) + 1))[: src.height]
    target = rukh_home / "data" / "uci" / "year=2025" / "month=01" / "games.parquet"
    target.parent.mkdir(parents=True)
    src.with_columns(pl.Series("game_id", ids, dtype=pl.Int64)).write_parquet(target.as_posix())
    manifest = run(EloBinsConfig(n_per_bin=1, seed=42))
    assert sum(manifest.counts.values()) > 0


def test_run_is_deterministic(rukh_home: Path, repo_root: Path) -> None:
    _stage(rukh_home, repo_root)
    run(EloBinsConfig(n_per_bin=1, seed=3))
    first = pl.read_parquet((rukh_home / "data" / "elo-bins" / "games.parquet").as_posix())
    run(EloBinsConfig(n_per_bin=1, seed=3))
    second = pl.read_parquet((rukh_home / "data" / "elo-bins" / "games.parquet").as_posix())
    assert first.equals(second)
    assert first.height == first["bin"].n_unique()

"""Tests for rukh.data.elite: monkeypatched download, PGN parsing and UCI conversion."""

from __future__ import annotations

import zipfile
from pathlib import Path

import polars as pl
import pytest

from rukh.data import elite as elite_module
from rukh.data.elite import EliteConfig, extract_pgn, iter_pgn_rows, run, zip_url
from rukh.data.uci import convert_rows, game_id

pytestmark = pytest.mark.unit

PGN = """[Event "Rated Blitz game"]
[Site "https://lichess.org/abc1"]
[Date "2025.01.03"]
[White "A"]
[Black "B"]
[Result "1-0"]
[UTCDate "2025.01.03"]
[WhiteElo "2612"]
[BlackElo "2380"]
[ECO "C50"]
[TimeControl "180+2"]

1. e4 e5 2. Nf3 Nc6 3. Bc4 Bc5 4. c3 Nf6 5. d4 exd4 6. cxd4 Bb4+ 7. Nc3 Nxe4 8. O-O Bxc3
9. d5 Bf6 10. Re1 Ne7 11. Rxe4 d6 1-0

[Event "Rated Blitz game"]
[Site "https://lichess.org/abc2"]
[Result "0-1"]
[UTCDate "2025.01.04"]
[WhiteElo "2540"]
[BlackElo "2610"]
[ECO "A00"]
[TimeControl "300+0"]

1. f3 e5 2. g4 Qh4# 0-1

[Event "Rated Blitz game"]
[Site "https://lichess.org/abc3"]
[Result "1/2-1/2"]
[UTCDate "2025.01.04"]
[WhiteElo "2500"]
[BlackElo "2400"]
[TimeControl "300+0"]

1. e4 e5 2. Nf3 Nc6 3. Bb5 a6 4. Ba4 Nf6 5. O-O Be7 6. Re1 b5 7. Bb3 d6 8. c3 O-O 9. h3 Nb8
10. d4 Nbd7 11. Nbd2 Bb7 1/2-1/2

"""


def test_zip_url() -> None:
    cfg = EliteConfig()
    assert zip_url(cfg, "2025-01") == "https://database.nikonoel.fr/lichess_elite_2025-01.zip"


def test_iter_pgn_rows(tmp_path: Path) -> None:
    pgn = tmp_path / "x.pgn"
    pgn.write_text(PGN, encoding="utf-8")
    rows = list(iter_pgn_rows(pgn, "2025-01"))
    assert [r["Site"] for r in rows] == [f"https://lichess.org/abc{i}" for i in (1, 2, 3)]
    assert rows[0]["WhiteElo"] == 2612 and rows[0]["Result"] == "1-0"
    assert rows[0]["movetext"].startswith("1. e4 e5") and rows[0]["movetext"].endswith("d6 1-0")
    assert rows[1]["ECO"] == "A00" and rows[2]["ECO"] is None
    assert rows[2]["month"] == "2025-01"


NO_SITE_PGN = """[Event "Rated Blitz game"]
[Result "1-0"]
[WhiteElo "2600"]

1. e4 e5 2. Nf3 Nc6 3. Bc4 Bc5 4. c3 Nf6 5. d4 exd4 6. cxd4 Bb4+ 7. Nc3 Nxe4 8. O-O Bxc3
9. d5 Bf6 10. Re1 Ne7 11. Rxe4 d6 1-0

[Event "Rated Blitz game"]
[Result "0-1"]

1. d4 d5 2. Nf3 Nf6 3. Nc3 Nc6 4. Bf4 Bf5 5. Qd2 Qd7 6. O-O-O O-O-O 7. e3 e6 8. Bd3 Bd6
9. Ne5 Ne4 10. Bxe4 Bxe5 0-1

"""

HEADERS_ONLY_PGN = """[Event "Rated Blitz game"]
[Site "https://lichess.org/empty"]
[Result "*"]
[WhiteElo "2700"]
[BlackElo "2650"]

[Event "Rated Blitz game"]
[Site "https://lichess.org/real"]
[Result "1-0"]
[ECO "C50"]

1. e4 e5 2. Nf3 Nc6 3. Bc4 Bc5 4. c3 Nf6 5. d4 exd4 6. cxd4 Bb4+ 7. Nc3 Nxe4 8. O-O Bxc3
9. d5 Bf6 10. Re1 Ne7 11. Rxe4 d6 1-0

"""


def test_iter_pgn_rows_without_site_gives_each_game_its_own_id(tmp_path: Path) -> None:
    pgn = tmp_path / "no-site.pgn"
    pgn.write_text(NO_SITE_PGN, encoding="utf-8")
    rows = list(iter_pgn_rows(pgn, "2025-01"))
    assert [r["Site"] for r in rows] == ["2025-01:0", "2025-01:1"]
    ids = convert_rows(rows, min_plies=0, max_plies=300)["columns"]["game_id"]  # type: ignore[index]
    assert ids == [game_id("2025-01:0"), game_id("2025-01:1")]
    assert len(set(ids)) == 2


def test_iter_pgn_rows_missing_elo_is_none(tmp_path: Path) -> None:
    pgn = tmp_path / "no-site.pgn"
    pgn.write_text(NO_SITE_PGN, encoding="utf-8")
    rows = list(iter_pgn_rows(pgn, "2025-01"))
    assert rows[0]["WhiteElo"] == 2600 and rows[0]["BlackElo"] is None
    columns = convert_rows(rows, min_plies=0, max_plies=300)["columns"]  # type: ignore[index]
    assert columns["black_elo"] == [None, None]


def test_iter_pgn_rows_drops_a_game_without_movetext(tmp_path: Path) -> None:
    pgn = tmp_path / "headers-only.pgn"
    pgn.write_text(HEADERS_ONLY_PGN, encoding="utf-8")
    rows = list(iter_pgn_rows(pgn, "2025-01"))
    assert len(rows) == 1
    assert rows[0]["Site"] == "https://lichess.org/real"
    # The dropped game's headers must not leak into the next one.
    assert rows[0]["WhiteElo"] is None and rows[0]["ECO"] == "C50"


def _fake_download(pgn_text: str, calls: list[str]):  # type: ignore[no-untyped-def]
    def download(url: str, dest: Path) -> None:
        calls.append(url)
        dest.parent.mkdir(parents=True, exist_ok=True)
        month = url.rsplit("_", 1)[1].removesuffix(".zip")
        with zipfile.ZipFile(dest, "w") as archive:
            archive.writestr(f"lichess_elite_{month}.pgn", pgn_text)

    return download


def test_run_downloads_once_and_converts(rukh_home: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    calls: list[str] = []
    monkeypatch.setattr(elite_module, "_download", _fake_download(PGN, calls))
    cfg = EliteConfig(months=["2025-01", "2025-02"], workers=1)
    manifest = run(cfg)
    assert calls == [zip_url(cfg, "2025-01"), zip_url(cfg, "2025-02")]
    out = rukh_home / "data" / "elite"
    frame = pl.read_parquet((out / "games.parquet").as_posix())
    # The 4-ply game is too short; the other two survive, once per month.
    assert frame.height == 4
    assert frame["month"].to_list() == ["2025-01", "2025-01", "2025-02", "2025-02"]
    assert frame["uci"][0].startswith("e2e4 e7e5 g1f3 b8c6 f1c4 f8c5")
    assert frame.schema["white_elo"] == pl.Int16
    assert str(frame["utc_date"][0]) == "2025-01-03"
    assert manifest.counts == {"games": 4}
    assert manifest.filters["conversion"]["short"] == 2
    assert [f.path for f in manifest.files] == [
        "lichess_elite_2025-01.zip",
        "lichess_elite_2025-02.zip",
        "games.parquet",
    ]
    assert manifest.filters["source_urls"]["2025-01"].endswith("lichess_elite_2025-01.zip")
    # A second run reuses the zips on disk.
    run(cfg)
    assert len(calls) == 2
    assert extract_pgn(out / "lichess_elite_2025-01.zip", out).name == "lichess_elite_2025-01.pgn"

"""Tests for rukh.data.db: every connection gets the spill directory, memory limit and threads."""

from __future__ import annotations

import os
from pathlib import Path

import pytest

from rukh.data.db import TMP_DIRNAME, DuckDbConfig, connect, temp_directory
from rukh.data.pipeline import DUCKDB_STEPS, PipelineConfig

pytestmark = pytest.mark.unit


def _settings(con: object) -> dict[str, object]:
    row = con.execute(  # type: ignore[attr-defined]
        "SELECT current_setting('temp_directory'), current_setting('memory_limit'), "
        "current_setting('threads')"
    ).fetchone()
    return {"temp_directory": row[0], "memory_limit": row[1], "threads": row[2]}


def test_connect_applies_the_settings(rukh_home: Path) -> None:
    con = connect(DuckDbConfig(memory_limit="2GB", threads=3))
    try:
        settings = _settings(con)
    finally:
        con.close()
    assert settings["temp_directory"] == (rukh_home / "data" / TMP_DIRNAME).as_posix()
    assert settings["memory_limit"] == "1.8 GiB"  # DuckDB reports GiB
    assert settings["threads"] == 3
    assert (rukh_home / "data" / TMP_DIRNAME).is_dir()


def test_connect_defaults_to_32gb_and_every_core(rukh_home: Path) -> None:
    con = connect()
    try:
        settings = _settings(con)
    finally:
        con.close()
    assert settings["memory_limit"] == "29.8 GiB"  # 32 GB
    assert settings["threads"] == (os.cpu_count() or 4)
    assert temp_directory() == rukh_home / "data" / TMP_DIRNAME

def test_pipeline_shares_the_duckdb_section() -> None:
    cfg = PipelineConfig.model_validate({"duckdb": {"memory_limit": "8GB", "threads": 2}})
    for name in DUCKDB_STEPS:
        step = getattr(cfg, name)
        assert step.duckdb.memory_limit == "8GB" and step.duckdb.threads == 2
    # A step may still override it.
    cfg = PipelineConfig.model_validate(
        {"duckdb": {"memory_limit": "8GB"}, "evals": {"duckdb": {"memory_limit": "1GB"}}}
    )
    assert cfg.evals.duckdb.memory_limit == "1GB"
    assert cfg.positions.duckdb.memory_limit == "8GB"


def test_shipped_pipeline_yaml_declares_duckdb() -> None:
    from rukh.data.pipeline import load_pipeline

    cfg = load_pipeline()
    assert cfg.duckdb.memory_limit == "32GB"
    assert cfg.puzzles.duckdb == cfg.duckdb