from __future__ import annotations

import json
from pathlib import Path

import pytest
from scripts.prepare_answer_batches import (
    build_batches,
    load_tuning_questions,
    prepare,
    validate_answer_output,
)


def test_batch_builder_excludes_reserved_and_caps_batch_size(tmp_path: Path) -> None:
    path = tmp_path / "questions.json"
    path.write_text(
        json.dumps(
            [
                {"id": "Q01", "partition": "tuning", "question": "First?"},
                {"id": "Q99", "partition": "reserved", "question": "Secret?"},
                {"id": "Q02", "partition": "tuning", "question": "Second?"},
            ]
        ),
        encoding="utf-8",
    )
    questions = load_tuning_questions(path)
    batches = build_batches(questions, 2)
    assert [[row["id"] for row in batch] for batch in batches] == [["Q01", "Q02"]]
    with pytest.raises(ValueError, match="between 1 and 4"):
        build_batches(questions, 5)


def test_prepared_batches_include_limits_but_no_reserved_content(tmp_path: Path) -> None:
    source = tmp_path / "questions.json"
    source.write_text(
        json.dumps(
            [
                {"id": f"Q{i:02d}", "partition": "tuning", "question": f"Question {i}?"}
                for i in range(1, 7)
            ]
            + [{"id": "Q99", "partition": "reserved", "question": "Reserved answer clue"}]
        ),
        encoding="utf-8",
    )
    output = tmp_path / "out"
    result = prepare(
        type("Args", (), {"questions": source, "output": output, "batch_size": 4})()
    )
    assert result["question_count"] == 6
    assert result["batch_count"] == 2
    prompts = "\n".join(path.read_text(encoding="utf-8") for path in output.glob("*-prompt.md"))
    assert "Never repeat an identical search" in prompts
    assert "40000" in prompts
    assert "Reserved answer clue" not in prompts
    assert all((path.stat().st_mode & 0o077) == 0 for path in output.iterdir())


def test_answer_validator_requires_exact_ids_and_complete_citations(tmp_path: Path) -> None:
    expected = tmp_path / "expected.json"
    answers = tmp_path / "answers.json"
    expected.write_text(json.dumps([{"id": "Q01"}, {"id": "Q02"}]), encoding="utf-8")
    answers.write_text(
        json.dumps(
            {
                "answers": [
                    {"id": "Q01", "answer": "A", "citations": [], "evidence_gap": False},
                    {
                        "id": "Q02",
                        "answer": "B",
                        "citations": [
                            {
                                "node_id": "n",
                                "claim_id": "c",
                                "source_revision_id": "r",
                                "extraction_id": "e",
                                "locator": {"kind": "text_span"},
                                "quote": "exact quote",
                            }
                        ],
                        "evidence_gap": False,
                    },
                ]
            }
        ),
        encoding="utf-8",
    )
    report = validate_answer_output(expected, answers)
    assert report["status"] == "valid"

    malformed = json.loads(answers.read_text(encoding="utf-8"))
    malformed["answers"].pop()
    answers.write_text(json.dumps(malformed), encoding="utf-8")
    report = validate_answer_output(expected, answers)
    assert report["status"] == "invalid"
    assert any("exactly match" in issue for issue in report["issues"])