"""CPU-only isolated tests for the staged pinned-Qwen parser proposal."""
from __future__ import annotations

import copy
import signal
import time

import pytest

from hessian_recipes import agent_rl_training as proposal
regex = pytest.importorskip("regex")
transformers = pytest.importorskip("transformers")
recursive_parse = pytest.importorskip("transformers.utils.chat_parsing_utils").recursive_parse
qwen3_schema = pytest.importorskip("trl.chat_template_utils").qwen3_schema


@pytest.fixture(scope="module")
def tokenizer():
    import string

    tokenizers = pytest.importorskip("tokenizers")
    specials = ["<|pad|>", "<|im_start|>", "<|im_end|>", "<think>", "</think>",
                "<tool_call>", "</tool_call>"]
    alphabet = list(dict.fromkeys(string.printable + "\x1c\x1d\x1e\x1f\x85\xa0α理由終わり"))
    vocab = {token: i for i, token in enumerate(specials + alphabet + ["ab"])}
    backend = tokenizers.Tokenizer(tokenizers.models.BPE(vocab=vocab, merges=[("a", "b")]))
    backend.decoder = tokenizers.decoders.Fuse()
    tokenizer = transformers.PreTrainedTokenizerFast(
        tokenizer_object=backend, eos_token="<|im_end|>", pad_token="<|pad|>",
        additional_special_tokens=specials[1:],
    )
    return tokenizer


def parse(tokenizer, text):
    ids = tokenizer.encode(text, add_special_tokens=False)
    return proposal._parse_qwen3_bounded(tokenizer, ids), ids


def expected(text):
    return recursive_parse(text, qwen3_schema)


@pytest.mark.parametrize("text", [
    "", "answer", "a\nb\n", "  ", "<|im_end|>", "text<|im_end|>ignored",
    "<think>\nreason\n</think>\n\nanswer<|im_end|>",
    "<think>\n\n</think>\n\nanswer", "<think> </think>answer",
    "<think>理由 α\n</think>\n終わり", "<think>unfinished reasoning",
    "<think>\x1c\x1d</think>\x1eanswer",
    "preface\n<tool_call>\x1f{\"name\":\"x\",\"arguments\":{}}\x85</tool_call>\xa0",
    "<tool_call>{\"name\":\"x\",\"arguments\":{}}</tool_call>",
    "preface\n<tool_call>\n{\"name\":\"x\",\"arguments\":{\"a\":2}}\n</tool_call><|im_end|>",
    "<tool_call>{\"name\":\"x\",\"arguments\":{}}</tool_call>\n"
    "<tool_call>{\"name\":\"y\",\"arguments\":{\"nested\":{\"n\":[1,true,null]}}}</tool_call>",
    "<think>call it</think>\n<tool_call>{\"name\":\"x\",\"arguments\":{}}</tool_call>",
    "<tool_call>x", "ordinary <tool_call> malformed", "a</think>b",
])
def test_native_semantics_and_raw_ids(tokenizer, text):
    actual, ids = parse(tokenizer, text)
    raw = actual.pop("raw_response")
    sampled = actual.pop("raw_token_ids")
    assert actual == expected(tokenizer.decode(ids))
    assert raw == tokenizer.decode(ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)
    assert sampled == ids and sampled is not ids


def test_differential_native_success_corpus(tokenizer):
    # Small cases keep the original stdlib parser safely bounded too.
    for prefix in ("", "text", "<think>r</think>\n", "<think>\n\n</think>\n"):
        for count in range(5):
            calls = '<tool_call>{"name":"x","arguments":{"n":3}}</tool_call>\n' * count
            for suffix in ("", "<|im_end|>", "<tool_call>x"):
                text = prefix + calls + suffix
                actual, ids = parse(tokenizer, text)
                actual.pop("raw_response")
                actual.pop("raw_token_ids")
                assert actual == expected(tokenizer.decode(ids))


def test_pathological_native_input_under_real_token_cap_is_fast_and_exact(tokenizer):
    text = '<tool_call>{"name":"x","arguments":{}}</tool_call>' * 24 + "<tool_call>x"
    ids = tokenizer.encode(text, add_special_tokens=False)
    assert len(ids) < 768
    start = time.monotonic()
    actual = proposal._parse_qwen3_bounded(tokenizer, ids)
    assert time.monotonic() - start < 1.0
    assert actual == {"role": "assistant", "content": text,
                      "raw_response": text, "raw_token_ids": ids}


@pytest.mark.parametrize("text", [
    '<tool_call>not json</tool_call>',
    '<tool_call>{"name": 3, "arguments": {}}</tool_call>',
    '<tool_call>{"name":"x","arguments":3}</tool_call>',
])
def test_native_schema_validation_is_retained(tokenizer, text):
    with pytest.raises((ValueError, TypeError)):
        expected(text)
    with pytest.raises((ValueError, TypeError)):
        parse(tokenizer, text)


def test_unknown_schema_and_zero_length_iterator_mutation_fail_closed(tokenizer):
    ids = tokenizer.encode("x")
    wrong = copy.deepcopy(qwen3_schema)
    wrong["properties"]["tool_calls"]["x-regex-iterator"] = "()"
    with pytest.raises(ValueError, match="exact pinned"):
        proposal._parse_qwen3_bounded(tokenizer, ids, schema=wrong)


def test_limits_reject_before_decode(tokenizer):
    with pytest.raises(ValueError, match="shape or size"):
        proposal._parse_qwen3_bounded(tokenizer, [1] * (proposal._QWEN_PARSE_MAX_IDS + 1))
    with pytest.raises(ValueError, match="shape or size"):
        proposal._parse_qwen3_bounded(tokenizer, [True])


def test_character_ceiling_and_budget_cover_all_parser_phases(tokenizer, monkeypatch):
    monkeypatch.setattr(proposal, "_QWEN_PARSE_MAX_CHARS", 1)
    with pytest.raises(ValueError, match="character"):
        parse(tokenizer, "long")
    monkeypatch.setattr(proposal, "_QWEN_PARSE_MAX_CHARS", 131072)
    monkeypatch.setattr(proposal, "_QWEN_PARSE_SECONDS", -1)
    with pytest.raises(TimeoutError, match="deadline"):
        parse(tokenizer, "x")


def test_recursive_validation_and_iterator_are_inside_shared_deadline(tokenizer, monkeypatch):
    monkeypatch.setattr(proposal, "_QWEN_PARSE_SECONDS", .01)
    original = regex.finditer

    def slow_iterator(*args, **kwargs):
        time.sleep(.03)
        return original(*args, **kwargs)

    monkeypatch.setattr(regex, "finditer", slow_iterator)
    with pytest.raises(TimeoutError):
        parse(tokenizer, '<tool_call>{"name":"x","arguments":{}}</tool_call>')


def test_own_regex_timeout_propagates_without_becoming_fallback(tokenizer, monkeypatch):
    error = TimeoutError("regex timed out")

    def fail(*args, **kwargs):
        raise error

    monkeypatch.setattr(regex, "search", fail)
    with pytest.raises(TimeoutError) as caught:
        parse(tokenizer, "x")
    assert caught.value is error


def test_foreign_sigalrm_propagates_even_with_same_exception_text(tokenizer, monkeypatch):
    # This test process owns this temporary alarm. The proposed parser must not
    # set/cancel it or catch the foreign exception, including identical messages.
    original_handler = signal.getsignal(signal.SIGALRM)
    original_timer = signal.getitimer(signal.ITIMER_REAL)
    foreign = TimeoutError("regex timed out")

    def guard(signum, frame):
        raise foreign

    def interrupted(*args, **kwargs):
        time.sleep(.2)

    monkeypatch.setattr(regex, "search", interrupted)
    signal.signal(signal.SIGALRM, guard)
    try:
        signal.setitimer(signal.ITIMER_REAL, .02)
        with pytest.raises(TimeoutError) as caught:
            parse(tokenizer, "x")
        assert caught.value is foreign
    finally:
        signal.setitimer(signal.ITIMER_REAL, 0)
        signal.signal(signal.SIGALRM, original_handler)
        signal.setitimer(signal.ITIMER_REAL, *original_timer)


def test_success_does_not_steal_or_renew_existing_alarm(tokenizer):
    original_handler = signal.getsignal(signal.SIGALRM)
    original_timer = signal.getitimer(signal.ITIMER_REAL)
    try:
        signal.setitimer(signal.ITIMER_REAL, 2)
        parse(tokenizer, "ordinary")
        remaining, interval = signal.getitimer(signal.ITIMER_REAL)
        assert 0 < remaining < 2 and interval == 0
        assert signal.getsignal(signal.SIGALRM) == original_handler
    finally:
        signal.setitimer(signal.ITIMER_REAL, 0)
        signal.setitimer(signal.ITIMER_REAL, *original_timer)


def test_recursive_schema_work_is_inside_shared_deadline(tokenizer, monkeypatch):
    import transformers.utils.chat_parsing_utils as parsing

    def slow_parse(*args, **kwargs):
        time.sleep(.03)
        return {"role": "assistant", "content": "x"}

    monkeypatch.setattr(proposal, "_QWEN_PARSE_SECONDS", .01)
    monkeypatch.setattr(parsing, "recursive_parse", slow_parse)
    with pytest.raises(TimeoutError):
        parse(tokenizer, "x")


def test_noncanonical_sampled_ids_survive_configured_native_entrypoint(tokenizer, monkeypatch):
    from trl.chat_template_utils import qwen3_training_chat_template

    tokenizer.chat_template = qwen3_training_chat_template
    metadata = proposal.configure_policy_tokenizer(tokenizer)
    assert metadata["response_parser"]["version"] == "qwen3-bounded-single-span-v1"
    aid, bid = tokenizer.convert_tokens_to_ids(["a", "b"])
    sampled = [aid, bid]
    assert tokenizer.encode("ab", add_special_tokens=False) != sampled
    actual = tokenizer.parse_response(sampled)
    assert actual["raw_token_ids"] == sampled and actual["raw_response"] == "ab"
    assert actual["content"] == "ab"

    # Native TRL catches ValueError/TypeError for malformed output, but an
    # actual outer wall guard must remain observable through the full entrypoint.
    from trl.chat_template_utils import parse_response as native_entrypoint

    foreign = TimeoutError("original experiment wall budget")

    def interrupt(*args, **kwargs):
        raise foreign

    monkeypatch.setattr(regex, "search", interrupt)
    with pytest.raises(TimeoutError) as caught:
        native_entrypoint(tokenizer, sampled)
    assert caught.value is foreign


def test_actual_regex_engine_preserves_foreign_alarm_exception():
    # A capped synthetic decoder isolates the regex C-engine from tokenization.
    # Missing </think> makes this root branch expensive; no private output used.
    decoder = type("Decoder", (), {
        "decode": lambda self, ids, **kwargs: "<think>" + "a " * 20000})()
    original_handler = signal.getsignal(signal.SIGALRM)
    original_timer = signal.getitimer(signal.ITIMER_REAL)
    foreign = TimeoutError("foreign original eight-hour wall guard")

    def guard(signum, frame):
        raise foreign

    signal.signal(signal.SIGALRM, guard)
    try:
        signal.setitimer(signal.ITIMER_REAL, .02)
        with pytest.raises(TimeoutError) as caught:
            proposal._parse_qwen3_bounded(decoder, [1])
        assert caught.value is foreign
    finally:
        signal.setitimer(signal.ITIMER_REAL, 0)
        signal.signal(signal.SIGALRM, original_handler)
        signal.setitimer(signal.ITIMER_REAL, *original_timer)


@pytest.mark.parametrize("argument", [
    "</tool_call>", "<tool_call>", "<tool_call></tool_call>",
    "ab</tool_call>", "<|im_end|>", "</tool_call></tool_call>",
])
def test_delimiter_text_in_json_retains_native_success_or_error(tokenizer, argument):
    import json

    body = json.dumps({"name": "x", "arguments": {"text": argument}})
    text = "<tool_call>" + body + "</tool_call>"
    ids = tokenizer.encode(text, add_special_tokens=False)
    try:
        native = expected(tokenizer.decode(ids))
    except (ValueError, TypeError) as error:
        with pytest.raises(type(error)):
            proposal._parse_qwen3_bounded(tokenizer, ids)
    else:
        actual = proposal._parse_qwen3_bounded(tokenizer, ids)
        actual.pop("raw_response")
        actual.pop("raw_token_ids")
        assert actual == native


def test_embedded_close_matches_full_native_trl_fallback_with_sampled_ids(tokenizer):
    import json
    from types import SimpleNamespace

    from trl.chat_template_utils import parse_response as native_entrypoint
    from trl.chat_template_utils import qwen3_training_chat_template

    text = "<tool_call>" + json.dumps(
        {"name": "x", "arguments": {"text": "ab</tool_call>"}}) + "</tool_call>"
    canonical = tokenizer.encode(text, add_special_tokens=False)
    aid, bid, merged = tokenizer.convert_tokens_to_ids(["a", "b", "ab"])
    sampled = [piece for token in canonical for piece in ([aid, bid] if token == merged else [token])]
    assert sampled != canonical
    before = list(sampled)
    baseline = SimpleNamespace(
        parse_response=lambda ids: expected(tokenizer.decode(ids)),
        decode=tokenizer.decode, eos_token=tokenizer.eos_token)
    with pytest.raises(ValueError):
        baseline.parse_response(sampled)
    tokenizer.chat_template = qwen3_training_chat_template
    proposal.configure_policy_tokenizer(tokenizer)
    with pytest.raises(ValueError):
        tokenizer.parse_response(sampled)
    expected_fallback = native_entrypoint(baseline, sampled)
    actual_fallback = native_entrypoint(tokenizer, sampled)
    assert actual_fallback == expected_fallback == {
        "role": "assistant", "content": tokenizer.decode(sampled, skip_special_tokens=True)}
    assert "tool_calls" not in actual_fallback
    assert sampled == before


def test_escaped_closing_delimiter_keeps_valid_json_tool_semantics(tokenizer):
    import json

    body = json.dumps({"name": "x", "arguments": {"text": "</tool_call>"}})
    text = "<tool_call>" + body.replace("</tool_call>", r"<\/tool_call>") + "</tool_call>"
    actual, ids = parse(tokenizer, text)
    native = expected(tokenizer.decode(ids))
    assert actual["tool_calls"] == native["tool_calls"]
    assert actual["tool_calls"][0]["function"]["arguments"] == {"text": "</tool_call>"}
    assert actual["raw_token_ids"] == ids