"""TDD tests for compiler/stages/split.py."""

from __future__ import annotations

import json
from pathlib import Path

import pytest
from compiler.stages.split import (
    HOLDOUT_RATE,
    _parse_paragraphs,
    assert_holdout_disjoint,
    split_corpus,
)

FIXTURES_HTTP = Path(__file__).parent.parent / "fixtures" / "http"

SAMPLE_MD = """\
---
source_id: test-src
title: Test
rights: public-domain
fetch_type: gutenberg
---

# Test Title

This is the first paragraph of sufficient length to be included in the corpus.

This is the second paragraph and it too has enough characters to meet the minimum.

Third paragraph here with enough words to pass the minimum character threshold check.

Fourth paragraph here with enough words to pass the minimum character threshold check.

Fifth paragraph here with enough words to pass the minimum character threshold check.

Sixth paragraph here with enough words to pass the minimum character threshold check.

Seventh paragraph here with enough words to pass the minimum character threshold check.

Eighth paragraph here with enough words to pass the minimum character threshold check.

Ninth paragraph here with enough words to pass the minimum character threshold check.

Tenth paragraph here with enough words to pass the minimum character threshold check.
"""


# ── _parse_paragraphs ─────────────────────────────────────────────────────────


def test_parse_paragraphs_strips_frontmatter():
    paras = _parse_paragraphs(SAMPLE_MD)
    for p in paras:
        assert "source_id" not in p
        assert "---" not in p


def test_parse_paragraphs_filters_short():
    long_para = (
        "This is a long enough paragraph to pass the minimum character threshold."
        " More words added here to make it longer."
    )
    md = f"---\nfoo: bar\n---\n\n# Title\n\nShort.\n\n{long_para}"
    paras = _parse_paragraphs(md)
    assert not any(p == "Short." for p in paras)
    assert any("long enough" in p for p in paras)


def test_parse_paragraphs_returns_nonempty_for_real_fixture():
    raw = (FIXTURES_HTTP / "gutenberg_2680.txt").read_text()
    # Simulate what normalize produces
    from compiler.stages.normalize import _strip_gutenberg
    body = _strip_gutenberg(raw)
    pseudo_md = f"---\nsource_id: x\n---\n\n# Title\n\n{body}"
    paras = _parse_paragraphs(pseudo_md)
    assert len(paras) > 0


# ── split_corpus ─────────────────────────────────────────────────────────────


def test_split_corpus_produces_holdout_file(tmp_path):
    corpus_path = tmp_path / "corpus" / "test-src.md"
    corpus_path.parent.mkdir(parents=True)
    corpus_path.write_text(SAMPLE_MD)
    style_dir = tmp_path / "style"

    split_corpus({"test-src": corpus_path}, style_dir)

    holdout = style_dir / "holdout.jsonl"
    assert holdout.exists()
    records = [json.loads(ln) for ln in holdout.read_text().splitlines() if ln.strip()]
    assert len(records) > 0


def test_split_corpus_holdout_rate_approximately_correct(tmp_path):
    corpus_path = tmp_path / "corpus" / "test-src.md"
    corpus_path.parent.mkdir(parents=True)
    corpus_path.write_text(SAMPLE_MD)
    style_dir = tmp_path / "style"

    indexable = split_corpus({"test-src": corpus_path}, style_dir)

    holdout_path = style_dir / "holdout.jsonl"
    held = [json.loads(ln) for ln in holdout_path.read_text().splitlines() if ln.strip()]
    total = len(indexable["test-src"]) + len(held)
    actual_rate = len(held) / total if total else 0
    # Allow 2x the target rate as a generous tolerance for small corpora
    assert actual_rate <= HOLDOUT_RATE * 2 + 0.05


def test_holdout_disjoint_passes(tmp_path):
    corpus_path = tmp_path / "corpus" / "test-src.md"
    corpus_path.parent.mkdir(parents=True)
    corpus_path.write_text(SAMPLE_MD)
    style_dir = tmp_path / "style"

    indexable = split_corpus({"test-src": corpus_path}, style_dir)
    assert_holdout_disjoint(indexable, style_dir)  # must not raise


def test_holdout_disjoint_fails_on_overlap(tmp_path):
    style_dir = tmp_path / "style"
    style_dir.mkdir()

    leaking_para = "This paragraph is in both indexed and holdout."
    (style_dir / "holdout.jsonl").write_text(
        json.dumps({"source_id": "s", "text": leaking_para}) + "\n"
    )
    indexable = {"s": [leaking_para, "Another paragraph that is not in holdout."]}

    with pytest.raises(AssertionError, match="holdout"):
        assert_holdout_disjoint(indexable, style_dir)


def test_split_corpus_multiple_sources(tmp_path):
    corpus_dir = tmp_path / "corpus"
    corpus_dir.mkdir()
    style_dir = tmp_path / "style"

    for i in range(3):
        (corpus_dir / f"src-{i}.md").write_text(SAMPLE_MD)

    corpus_paths = {f"src-{i}": corpus_dir / f"src-{i}.md" for i in range(3)}
    indexable = split_corpus(corpus_paths, style_dir)
    assert_holdout_disjoint(indexable, style_dir)

    holdout = style_dir / "holdout.jsonl"
    records = [json.loads(ln) for ln in holdout.read_text().splitlines() if ln.strip()]
    source_ids = {r["source_id"] for r in records}
    assert len(source_ids) == 3