#!/usr/bin/env python3
"""Local tests for the accepted-baseline parser path."""

import json

from script import build_prompt, count_expected_answers, parse_model_output


BASELINE_SYSTEM = (
    "You solve International Linguistics Olympiad problems using only the data in the problem. "
    "Infer the pattern from the examples. Think silently. Return valid JSON only: "
    '{"answers": ["..."]}'
)


def check_parse(raw: str, expected: list[str]) -> None:
    actual = parse_model_output(raw, None)
    assert actual == expected, (raw, actual, expected)

    encoded = json.dumps(actual, ensure_ascii=False)
    assert json.loads(encoded) == expected


def main() -> None:
    assert count_expected_answers("Translate:\n1. mi\n2. yu") is None
    assert count_expected_answers("") is None

    messages = build_prompt(
        {
            "task_type": "match_letter",
            "eval_type": "single",
            "context": "1. context item should not affect answer count",
            "query": "Translate this one item.",
        },
        None,
    )
    assert messages[0]["content"] == BASELINE_SYSTEM
    assert "Task type: match_letter" in messages[1]["content"]
    assert "Return one answer per numbered item." in messages[1]["content"]
    assert "Use the context examples" not in messages[1]["content"]

    check_parse("1. abc\n2. def", ["abc", "def"])
    check_parse("1. abc\n2. def\n3. ghi\n4. jkl\n5. mno", ["abc", "def", "ghi", "jkl", "mno"])
    check_parse("(a) ká  [ŋa]\n- čʼa\n2) O'Neil", ["ká  [ŋa]", "čʼa", "O'Neil"])
    check_parse("'áŋa' [x] /t͡sʼ/\n(kept)", ["'áŋa' [x] /t͡sʼ/", "(kept)"])
    check_parse("  keep  internal   spaces  ", ["keep  internal   spaces"])
    check_parse("", [])
    check_parse('{"answers": ["mi go", "yu sleep"]}', ['{"answers": ["mi go", "yu sleep"]}'])
    check_parse('["ká [ŋa]", "čʼa"]', ['["ká [ŋa]", "čʼa"]'])


if __name__ == "__main__":
    main()