from local_train.think import thought_text, wrap_completion


def test_wrap_converts_thought_prefix_and_keeps_one_bash_block():
    gold = "THOUGHT: Look at retention.go\n\n```bash\ncat /testbed/retention.go\n```\n```bash\nls\n```"
    out = wrap_completion(gold)
    assert out is not None
    assert out.startswith("<think>\nLook at retention.go")
    assert "</think>" in out
    assert out.count("```bash") == 1
    assert "cat /testbed/retention.go" in out
    assert "ls\n```" not in out


def test_wrap_override_rewrites_submit_command():
    gold = "THOUGHT: done\n```bash\necho COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\n```"
    out = wrap_completion(gold, "echo ALBEDO_TASK_DONE_SUBMIT_NOW")
    assert out is not None
    assert "echo ALBEDO_TASK_DONE_SUBMIT_NOW" in out
    assert "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" not in out.split("</think>", 1)[-1]


def test_wrap_rejects_missing_bash():
    assert wrap_completion("THOUGHT: just thinking") is None


def test_existing_think_block_is_reused():
    gold = "<think>plan</think>\n\n```bash\nsed -i 's/a/b/' foo.py\n```"
    assert thought_text(gold) == "plan"
    out = wrap_completion(gold)
    assert out is not None
    assert "<think>\nplan\n</think>" in out