"""Regressions for noncanonical sampled BPE histories and native TRL inputs."""
from __future__ import annotations

from contextlib import contextmanager, nullcontext
import copy
from types import SimpleNamespace

import pytest

from hessian_recipes.agent_rl_training import (
    _audited_trainer_class,
    _install_token_preserving_chat_template,
    audit_policy_tokens,
    configure_policy_tokenizer,
)


class TextTokenizer:
    """A reversible tokenizer where sampled a,b differs from canonical merged ab."""
    padding_side = "left"

    def encode(self, text, add_special_tokens=False):
        assert not add_special_tokens
        ids, i = [], 0
        while i < len(text):
            if text[i:i + 2] == "ab":
                ids.append(300)
                i += 2
            else:
                ids.append(ord(text[i]) + 1)
                i += 1
        return ids

    def decode(self, ids, **kwargs):
        return "".join("ab" if token == 300 else chr(token - 1) for token in ids)

    def pad(self, inputs, *, padding=False, padding_side="left", return_tensors=None,
            max_length=None, return_attention_mask=True):
        assert return_tensors is None
        single = not inputs["input_ids"] or isinstance(inputs["input_ids"][0], int)
        rows = [inputs["input_ids"]] if single else inputs["input_ids"]
        width = max(map(len, rows)) if padding else None
        values, masks = [], []
        for row in rows:
            count = width - len(row) if width else 0
            assert padding_side == "left"
            values.append([0] * count + list(row))
            masks.append([0] * count + [1] * len(row))
        return {"input_ids": values[0] if single else values,
                "attention_mask": masks[0] if single else masks}

    def apply_chat_template(self, conversation, *, tokenize=True, return_dict=True,
                            add_generation_prompt=False, padding=False, **kwargs):
        batch = bool(conversation) and isinstance(conversation[0], (list, tuple))
        conversations = conversation if batch else [conversation]
        texts = []
        for messages in conversations:
            text = ""
            for message in messages:
                text += "<" + message["role"] + ">\n"
                if "raw_response" in message:
                    raw = message["raw_response"]
                    text += message.get("hessian_raw_span_start", "")
                    text += raw
                    text += message.get("hessian_raw_span_end", "")
                    text += "" if raw.endswith("~") else "~"
                else:
                    text += message["content"] + "~"
                text += "\n"
            if add_generation_prompt:
                text += "<assistant>\n"
            texts.append(text)
        if not tokenize:
            return texts if batch else texts[0]
        rows = [self.encode(text) for text in texts]
        result = self.pad({"input_ids": rows if batch else rows[0]}, padding=padding)
        return result if return_dict else result["input_ids"]


def captured(tokenizer, text="ab~"):
    ids = [ord(character) + 1 for character in text]
    return {"role": "assistant", "content": "", "raw_response": text, "raw_token_ids": ids}


def history(tokenizer, text="ab~"):
    return [{"role": "user", "content": "question"}, captured(tokenizer, text),
            {"role": "tool", "content": "observation"}]


def test_noncanonical_sample_is_preserved_in_subsequent_context_and_loss():
    tokenizer = TextTokenizer()
    messages = history(tokenizer)
    original = copy.deepcopy(messages)
    canonical = tokenizer.apply_chat_template(messages, return_dict=False, add_generation_prompt=True)
    initial = tokenizer.apply_chat_template(messages[:1], return_dict=False, add_generation_prompt=True)
    _install_token_preserving_chat_template(tokenizer)
    context = tokenizer.apply_chat_template(messages, return_dict=False, add_generation_prompt=True)
    sample = messages[1]["raw_token_ids"]
    assert canonical != context
    assert tokenizer.decode(canonical) == tokenizer.decode(context)
    assert context[:len(initial)] == initial
    assert context[len(initial):len(initial) + len(sample)] == sample
    second = [ord("z") + 1]
    completion = context[len(initial):] + second
    mask = [1] * len(sample) + [0] * (len(context) - len(initial) - len(sample)) + [1]
    spans = [{"prompt_ids": initial, "completion_ids": sample, "call_index": 0},
             {"prompt_ids": context, "completion_ids": second, "call_index": 1}]
    assert audit_policy_tokens([initial], [completion], [mask], spans)[0]["sampled_token_equality_verified"]
    with pytest.raises(ValueError):
        audit_policy_tokens([initial], [canonical[len(initial):] + second], [mask], spans)
    assert messages == original


def test_original_initial_prompt_and_text_rendering_are_unchanged():
    tokenizer = TextTokenizer()
    initial = [{"role": "user", "content": "ab"}]
    ids = tokenizer.apply_chat_template(initial, return_dict=False, add_generation_prompt=True)
    text = tokenizer.apply_chat_template(history(tokenizer), tokenize=False, add_generation_prompt=True)
    _install_token_preserving_chat_template(tokenizer)
    installed = tokenizer.apply_chat_template
    _install_token_preserving_chat_template(tokenizer)
    assert tokenizer.apply_chat_template is installed
    assert tokenizer.apply_chat_template(initial, return_dict=False, add_generation_prompt=True) == ids
    assert tokenizer.apply_chat_template(history(tokenizer), tokenize=False, add_generation_prompt=True) == text


def test_batch_left_padding_and_single_output_shapes():
    tokenizer = TextTokenizer()
    _install_token_preserving_chat_template(tokenizer)
    conversations = [history(tokenizer), history(tokenizer, "ababab~")]
    unpadded = tokenizer.apply_chat_template(conversations, return_dict=False)
    encoded = tokenizer.apply_chat_template(conversations, padding=True)
    assert len(encoded["input_ids"][0]) == len(encoded["input_ids"][1])
    for ids, mask, expected in zip(encoded["input_ids"], encoded["attention_mask"], unpadded, strict=True):
        assert [token for token, keep in zip(ids, mask, strict=True) if keep] == expected
        assert mask == sorted(mask)
    assert tokenizer.apply_chat_template(conversations[0], return_dict=False) == unpadded[0]
    assert tokenizer.apply_chat_template(conversations[0])["input_ids"] == unpadded[0]


def test_batched_shared_message_objects_do_not_share_rendering_markers():
    tokenizer = TextTokenizer()
    messages = history(tokenizer)
    _install_token_preserving_chat_template(tokenizer)
    single = tokenizer.apply_chat_template(messages, return_dict=False)
    assert tokenizer.apply_chat_template([messages, messages], return_dict=False) == [single, single]
    assert "hessian_raw_span_start" not in messages[1]


def test_marker_collision_is_avoided_without_altering_tool_text():
    tokenizer = TextTokenizer()
    messages = history(tokenizer)
    messages[2]["content"] = "\x00hessian-sampled-span:0:0:start\x00 ab"
    expected = tokenizer.apply_chat_template(messages, tokenize=False)
    _install_token_preserving_chat_template(tokenizer)
    ids = tokenizer.apply_chat_template(messages, return_dict=False)
    assert tokenizer.decode(ids) == expected


@pytest.mark.parametrize("option", [
    {"truncation": True}, {"continue_final_message": True},
    {"return_assistant_tokens_mask": True}, {"return_offsets_mapping": True},
    {"tokenizer_kwargs": {"add_special_tokens": True}}, {"return_tensors": "tf"},
])
def test_unsupported_options_fail_closed(option):
    tokenizer = TextTokenizer()
    _install_token_preserving_chat_template(tokenizer)
    with pytest.raises(ValueError, match="Unsupported"):
        tokenizer.apply_chat_template(history(tokenizer), **option)


@pytest.mark.parametrize("change", [
    {"raw_token_ids": None}, {"raw_token_ids": [True]}, {"raw_token_ids": [-1]},
    {"raw_token_ids": [123]}, {"hessian_raw_span_start": "injected"},
])
def test_invalid_or_spoofed_metadata_fails_closed(change):
    tokenizer = TextTokenizer()
    messages = history(tokenizer)
    messages[1].update(change)
    _install_token_preserving_chat_template(tokenizer)
    with pytest.raises(ValueError):
        tokenizer.apply_chat_template(messages)


def test_template_that_ignores_preserved_spans_is_rejected():
    tokenizer = TextTokenizer()
    original_apply = tokenizer.apply_chat_template

    def wrong_template(conversation, **kwargs):
        rows = copy.deepcopy(conversation)
        for messages in rows:
            for message in messages:
                message.pop("hessian_raw_span_start", None)
                message.pop("hessian_raw_span_end", None)
        return original_apply(rows, **kwargs)

    tokenizer.apply_chat_template = wrong_template
    _install_token_preserving_chat_template(tokenizer)
    with pytest.raises(ValueError, match="omitted or duplicated"):
        tokenizer.apply_chat_template(history(tokenizer))


def tiny_qwen_tokenizer():
    tokenizers = pytest.importorskip("tokenizers")
    transformers = pytest.importorskip("transformers")
    trl = pytest.importorskip("trl.chat_template_utils")
    import string
    specials = ["<|pad|>", "<|im_start|>", "<|im_end|>", "<think>", "</think>",
                "<tool_call>", "</tool_call>"]
    vocab = {token: i for i, token in enumerate(specials + list(string.printable) + ["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:],
    )
    tokenizer.chat_template = trl.qwen3_training_chat_template
    configure_policy_tokenizer(tokenizer)
    return tokenizer


@pytest.mark.parametrize(("clip", "mixed"), [(False, False), (True, False), (False, True)])
def test_actual_pinned_trl_tool_loop_passes_exact_ids_to_model(monkeypatch, tmp_path, clip, mixed):
    torch = pytest.importorskip("torch")
    pytest.importorskip("jmespath")
    grpo = pytest.importorskip("trl.trainer.grpo_trainer")
    from transformers import Trainer
    tokenizer = tiny_qwen_tokenizer()
    aid, bid = tokenizer.convert_tokens_to_ids(["a", "b"])
    merged = tokenizer.convert_tokens_to_ids("ab")

    def sampled(text):
        result = []
        for token in tokenizer.encode(text, add_special_tokens=False):
            result.extend([aid, bid] if token == merged else [token])
        return result

    first = [sampled('<think>ab</think>\n<tool_call>\n{"name":"lookup","arguments":{}}\n</tool_call><|im_end|>'),
             sampled('<think>abab</think>\n<tool_call>\n{"name":"lookup","arguments":{}}\n</tool_call><|im_end|>')]
    second = [sampled('ab<|im_end|>'), first[1] if mixed else sampled('ab<|im_end|>')]
    third = [sampled('ab<|im_end|>')]

    class Model:
        training = True
        config = SimpleNamespace(max_position_embeddings=4096)

        def __init__(self):
            self.calls = []

        def generate(self, input_ids, attention_mask, **kwargs):
            self.calls.append((input_ids.clone(), attention_mask.clone()))
            values = [first, second, third][len(self.calls) - 1]
            width = max(map(len, values))
            suffix = torch.tensor([row + [tokenizer.pad_token_id] * (width - len(row)) for row in values])
            return torch.cat((input_ids, suffix), dim=1)

    @contextmanager
    def unwrap(*args, **kwargs):
        yield model

    monkeypatch.setattr(grpo, "unwrap_model_for_generation", unwrap)
    monkeypatch.setattr(grpo, "profiling_context", lambda *args, **kwargs: nullcontext())
    monkeypatch.setattr(Trainer, "_prepare_inputs", lambda self, inputs: inputs)
    model = Model()
    trainer = object.__new__(_audited_trainer_class(grpo.GRPOTrainer))
    trainer.model = trainer.model_wrapped = model
    trainer.accelerator = SimpleNamespace(device=torch.device("cpu"))
    trainer.args = SimpleNamespace(ds3_gather_for_generation=False, output_dir=str(tmp_path))
    trainer.is_fsdp_enabled = trainer.use_vllm = trainer.use_transformers_paged = False
    trainer.processing_class = tokenizer
    trainer.chat_template = tokenizer.chat_template
    trainer.chat_template_kwargs = {"enable_thinking": True}
    trainer.generation_kwargs = {}
    trainer.generation_config = None
    trainer.eos_token_id = tokenizer.eos_token_id
    trainer.tools = [{"type": "function", "function": {"name": "lookup", "description": "Read data.",
                     "parameters": {"type": "object", "properties": {}}}}]
    trainer._sync_tool_dicts = [{"lookup": lambda: "observation"}] * 2
    trainer._async_tool_dicts = [{}, {}]
    trainer.max_tool_calling_iterations = 2 if mixed else 1
    trainer.max_completion_length = 2048
    trainer._sampled_calls = []
    prompts = [[{"role": "user", "content": "short"}],
               [{"role": "user", "content": "a longer question"}]]
    prompt_ids, completion_ids, logprobs, _ = trainer._generate_single_turn(prompts)
    completions = [[tokenizer.parse_response(ids)] for ids in completion_ids]
    assert [row[0]["raw_token_ids"] for row in completions] == first
    if clip:
        contexts = [messages + [completion[0], {"role": "tool", "name": "lookup", "content": "observation"}]
                    for messages, completion in zip(prompts, completions, strict=True)]
        next_ids = tokenizer.apply_chat_template(contexts, tools=trainer.tools,
                    add_generation_prompt=True, return_dict=False, enable_thinking=True)
        trainer.max_completion_length = max(len(context) - len(prompt)
                    for context, prompt in zip(next_ids, prompt_ids, strict=True)) + 2
    masks, _, retained, _, _, _ = trainer._tool_call_loop(prompts, prompt_ids, completion_ids,
                                                        completions, logprobs)
    audits = audit_policy_tokens(prompt_ids, retained, masks, trainer._sampled_calls)
    assert len(model.calls) == (3 if mixed else 2)
    for index, audit in enumerate(audits):
        input_ids, attention = model.calls[1]
        actual_context = input_ids[index][attention[index].bool()].tolist()
        first_span, next_span, *later_spans = audit["sampled_spans"]
        assert actual_context == next_span["prompt_ids"]
        if mixed and index == 1:
            assert len(later_spans) == 1
            last_ids, last_attention = model.calls[2]
            assert last_ids.shape[0] == 1  # the other rollout already finished
            assert last_ids[0][last_attention[0].bool()].tolist() == later_spans[0]["prompt_ids"]
        else:
            assert not later_spans
        offset = len(prompt_ids[index])
        assert actual_context[offset:offset + len(first[index])] == first[index]
        assert sum(masks[index]) == sum(span["retained_tokens"] for span in audit["sampled_spans"])
        assert audit["external_tokens"] > 0
        assert first_span["sampled_completion_ids"] == first[index]
        assert audit["sampled_token_equality_verified"]
    assert model.calls[1][0].shape[0] == 2
    assert model.calls[1][1][0, 0] == 0  # shorter history is left padded
    if clip:
        assert any(audit["sampled_spans"][-1]["retained_tokens"] < len(second[index])
                   for index, audit in enumerate(audits))