from albedo_eval_service.remote.dataset import EvalSample
from albedo_eval_service.remote.generation import GenerationResult
from albedo_eval_service.shared.observation_format import (
    MAX_CONSECUTIVE_BAD_TURNS,
    truncation_notice,
)
from local_eval.gates import evaluate_side
from local_eval.live_protocol import generate_retrying_bad_turns, is_live_submit


def test_marker_in_first_bash_is_a_submit():
    clause = "echo ALBEDO_TASK_DONE_SUBMIT_NOW && git add -A && git diff --cached"
    text = (
        "<think>done</think>\n\n```bash\necho ALBEDO_TASK_DONE_SUBMIT_NOW && cat patch.txt\n```"
    )
    assert is_live_submit(
        text, command=clause, marker="ALBEDO_TASK_DONE_SUBMIT_NOW"
    )


def test_prose_mention_of_marker_is_not_a_submit():
    text = "I will next run echo ALBEDO_TASK_DONE_SUBMIT_NOW"
    assert not is_live_submit(
        text, command="echo ALBEDO_TASK_DONE_SUBMIT_NOW", marker="ALBEDO_TASK_DONE_SUBMIT_NOW"
    )


def test_gate_accepts_marker_submit_without_exact_tail():
    sample = EvalSample(
        sample_id="s1",
        prompt="p",
        submit_command="echo ALBEDO_TASK_DONE_SUBMIT_NOW && git add -A && git diff --cached",
        submit_marker="ALBEDO_TASK_DONE_SUBMIT_NOW",
    )
    text = "<think>x</think>\n\n```bash\necho ALBEDO_TASK_DONE_SUBMIT_NOW\n```"
    report = evaluate_side(sample, text, [{"role": "assistant", "content": text, "score_target": True}], truncated=False)
    assert report.submitted
    assert "did not issue submit command" not in report.warnings


def test_retrying_bad_turns_uses_the_usable_attempt():
    good = "THOUGHT: look\n\n```bash\nls -la\n```"
    waves = [
        [GenerationResult("s1", truncation_notice(4096), truncated=True)],
        [GenerationResult("s1", "")],
        [GenerationResult("s1", good)],
    ]

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

        def generate(self, samples):
            self.calls.append(samples[0].messages[-1]["content"] if samples[0].messages else "")
            return waves[len(self.calls) - 1]

    gen = _Gen()
    sample = EvalSample(sample_id="s1", prompt="p", messages=[{"role": "user", "content": "task"}])
    out = generate_retrying_bad_turns(gen, [sample])
    assert out[0].text == good
    assert len(gen.calls) == 3
    assert "reached the output token limit" in gen.calls[1]
    assert "Format error" in gen.calls[2]


def test_retrying_bad_turns_gives_up_after_three():
    class _Gen:
        def __init__(self):
            self.calls = 0

        def generate(self, samples):
            self.calls += 1
            return [GenerationResult("s1", "")]

    gen = _Gen()
    sample = EvalSample(sample_id="s1", prompt="p", messages=[{"role": "user", "content": "t"}])
    out = generate_retrying_bad_turns(gen, [sample])
    assert gen.calls == MAX_CONSECUTIVE_BAD_TURNS
    assert out[0].text == ""