"""Adversarial invariants for the frozen synthetic reward, without model inference."""
from __future__ import annotations

import copy
import json

import pytest

from hessian_recipes.agent_rl_tasks import GermanAgentEnv, generate_tasks, score_episode
from hessian_recipes.agent_rl_training import milestone_training_reward


def _call(name, arguments):
    return {"role": "assistant", "content": "", "tool_calls": [{
        "type": "function", "function": {"name": name, "arguments": arguments},
    }]}


def _completed_tools():
    row = generate_tasks("train", 1, seed=8765)[0]
    env = GermanAgentEnv()
    env.reset(**row)
    observed = json.loads(env.lookup(row["task"]["record_id"]))
    expression = f"{observed['quantity']}*{observed['unit_price']}"
    result = json.loads(env.calculator(expression))["result"]
    messages = [_call("lookup", {"record_id": row["task"]["record_id"]}),
                {"role": "tool", "name": "lookup", "content": json.dumps(observed)},
                _call("calculator", {"expression": expression}),
                {"role": "tool", "name": "calculator", "content": json.dumps({"result": result})}]
    return row, env, messages, result


@pytest.mark.parametrize("kind", ["unknown_name", "unexpected_keyword"])
def test_native_python_rejections_fail_strict_success_but_are_outside_trace_penalty(kind):
    row, env, messages, answer = _completed_tools()
    if kind == "unknown_name":
        name, arguments = "unknown_tool", {}
        with pytest.raises(AttributeError):
            getattr(env, name)(**arguments)
    else:
        name, arguments = "lookup", {"record_id": row["task"]["record_id"], "unexpected": "value"}
        with pytest.raises(TypeError):
            getattr(env, name)(**arguments)
    messages += [_call(name, arguments),
                 {"role": "tool", "name": name, "content": "{'error': 'upstream Python dispatch rejection'}"},
                 {"role": "assistant", "content": f"<answer>{answer}</answer>"}]
    score = score_episode(env, messages)
    reward = milestone_training_reward(score)
    assert score["success"] is False and score["reward"] == 0
    assert len(score["trace"]) == 2 and all(call["ok"] for call in score["trace"])
    assert reward["training_reward"] == 0.4
    assert reward["trusted_tool_errors"] == 0  # This is NOT a count of native dispatch errors.


def test_unexecuted_call_at_turn_cap_is_not_a_native_error_observation():
    _, env, messages, _ = _completed_tools()
    messages.append(_call("unknown_tool", {}))
    assert sum(message["role"] == "tool" for message in messages) == 2
    assert len(env.trace) == 2
    assert score_episode(env, messages)["reward"] == 0
    assert milestone_training_reward(score_episode(env, messages))["training_reward"] == 0.4


@pytest.mark.parametrize("expression", ["{answer}", "1*{answer}", "({q}*{p})+0", "{q}*{p}*1"])
def test_arithmetically_equivalent_shortcuts_cannot_replace_required_retrieved_operands(expression):
    row = generate_tasks("train", 1, seed=991)[0]
    env = GermanAgentEnv()
    env.reset(**row)
    facts = json.loads(env.lookup(row["task"]["record_id"]))
    answer = facts["quantity"] * facts["unit_price"]
    observation = json.loads(env.calculator(expression.format(answer=answer, q=facts["quantity"], p=facts["unit_price"])))
    assert observation["error"] == "must_multiply_retrieved_values"
    assert score_episode(env, f"<answer>{answer}</answer>")["reward"] == 0


def test_large_generated_splits_keep_identities_and_prompt_fact_boundary():
    splits = {split: generate_tasks(split, 2000, seed=20261020) for split in ("train", "dev", "test")}
    identity_sets = [{row["task"]["task_id"] for row in rows} for rows in splits.values()]
    record_sets = [{row["task"]["record_id"] for row in rows} for rows in splits.values()]
    assert len(set.union(*identity_sets)) == 6000
    assert len(set.union(*record_sets)) == 6000
    for rows in splits.values():
        for row in rows:
            assert row["task"]["record_id"] in row["prompt"]
            assert '"quantity":' not in row["prompt"] and '"unit_price":' not in row["prompt"]
            assert "<tool_response>" not in row["prompt"]
    pairs = {(row["task"]["private_facts"]["quantity"], row["task"]["private_facts"]["unit_price"])
             for row in splits["train"]}
    assert len(pairs) <= 80  # A narrow task family, regardless of the number of opaque records.


def test_unsupported_record_requests_do_not_return_hidden_facts():
    row = generate_tasks("train", 1)[0]
    for record in ("", "beleg-incorrect", None, 17, True, [], {}, row["task"]["record_id"].upper()):
        env = GermanAgentEnv()
        env.reset(**row)
        result = json.loads(env.lookup(record))
        assert result == {"error": "unknown_record"}
        assert not env._lookup_success


def test_wrong_terminal_answer_never_becomes_success_after_repeated_legal_calls():
    _, env, messages, answer = _completed_tools()
    for final in ("<answer>-1</answer>", "<answer>0</answer>", "<answer>999999</answer>"):
        transcript = copy.deepcopy(messages) + [{"role": "assistant", "content": final}]
        assert score_episode(env, transcript)["reward"] == 0
        assert milestone_training_reward(score_episode(env, transcript))["training_reward"] == 0.4
    valid = messages + [{"role": "assistant", "content": f"<answer>{answer}</answer>"}]
    assert score_episode(env, valid)["reward"] == 1