import pytest

from tiny_computer_agent import (
    Action, BrainError, DecisionMenu, Element, GoalUnreachable, ModelBrain,
    Observation, extract_json_object, extract_json_objects, parse_decision,
)

FRAME = Observation(elements=(
    Element("Files", x=10, y=5, width=40, height=20),
    Element("Save", x=200, y=260, width=60, height=20),
    Element("Cancel", x=300, y=260, width=60, height=20),
), text="untitled document")


def test_menu_is_numbered_in_reading_order_and_capped():
    observation = Observation(elements=tuple(
        Element(f"b{i}", x=(i % 2) * 100, y=(i // 2) * 50) for i in range(6)
    ))
    menu = DecisionMenu(observation, max_items=3)
    assert [item.index for item in menu] == [1, 2, 3]
    assert len(menu) == 3 and bool(menu)
    assert "1." in menu.render()


def test_menu_lookup_tolerates_what_models_send():
    menu = DecisionMenu(FRAME)
    assert menu.item(1).label == "Files"
    assert menu.item("2").label == "Save"
    assert menu.item(2.0).label == "Save"
    assert menu.item(0) is None and menu.item(99) is None and menu.item("x") is None


def test_menu_resolves_labels_through_normalisation():
    menu = DecisionMenu(Observation(elements=(Element("Save ."),)))
    assert menu.find("save").label == "Save ."
    assert menu.find("nothing").__class__ is type(None)


def test_json_objects_are_extracted_in_order_and_the_last_wins():
    raw = 'First {"kind": "click", "index": 1} then {"kind": "done", "reason": "ok"}'
    assert [o["kind"] for o in extract_json_objects(raw)] == ["click", "done"]
    # a model that echoes its prompt puts its own answer last
    assert extract_json_object(raw)["kind"] == "done"
    with pytest.raises(BrainError, match="no JSON object"):
        extract_json_object("no json here")


def test_parse_decision_json_index():
    action = parse_decision('{"kind": "click", "index": 2, "reason": "saves"}', FRAME)
    assert action.kind == "click" and action.label == "Save"


def test_parse_decision_echo_of_the_prompt_is_rejected():
    echoed = '{"kind": "click", "index": 1} {"kind": "click", "index": 2}'
    assert parse_decision(echoed, FRAME).label == "Save"          # last object wins
    with pytest.raises(BrainError, match="echo"):
        parse_decision(DecisionMenu(FRAME).render(), FRAME)        # 3+ labels mentioned


def test_parse_decision_bare_number_and_prose_label():
    assert parse_decision("2", FRAME).label == "Save"
    assert parse_decision('Click on "Save" in the dialogue.', FRAME).label == "Save"
    assert parse_decision("The goal is already done.", FRAME).kind == "done"


def test_parse_decision_rejects_an_index_outside_the_menu():
    with pytest.raises(BrainError, match="not in the"):
        parse_decision('{"kind": "click", "index": 99}', FRAME)


def test_parse_decision_rejects_a_label_that_is_not_on_screen():
    with pytest.raises(BrainError, match="not on screen"):
        parse_decision('{"kind": "click", "label": "Format disk"}', FRAME)


def test_parse_decision_tolerates_a_garbled_verb_when_a_label_is_present():
    # 256M models really do emit {"kind": "file", "label": "Save"}
    action = parse_decision('{"kind": "file", "label": "Save"}', FRAME)
    assert action.kind == "click" and action.label == "Save"


def test_parse_decision_empty_and_useless_replies():
    for raw in ["", "   ", "I am not sure what to do, sorry."]:
        with pytest.raises(BrainError):
            parse_decision(raw, FRAME)


class _ScriptedModel(ModelBrain):
    name = "scripted"

    def __init__(self, replies, **kwargs):
        super().__init__("fake/model", **kwargs)
        self._replies = list(replies)
        self.prompts = []

    def _generate(self, prompt, observation):
        self.prompts.append(prompt)
        return self._replies.pop(0)


def test_model_brain_tells_the_model_what_it_already_tried():
    brain = _ScriptedModel(['{"kind": "click", "index": 2}', '{"kind": "done", "reason": "ok"}'])
    first = brain.decide("save", FRAME)
    assert first.label == "Save"
    assert "Already attempted" not in brain.prompts[0]  # nothing tried yet

    second = brain.decide("save", FRAME)                 # same screen -> no progress
    prompt = brain.prompts[1]
    assert "Already attempted" in prompt
    assert "screen did NOT change" in prompt
    assert "Choose a DIFFERENT action" in prompt
    assert second.kind == "done"


def test_model_brain_history_is_bounded():
    brain = _ScriptedModel(['{"kind": "wait", "seconds": 1}'] * 6)
    for _ in range(6):
        brain.decide("wait", FRAME)
    assert len(brain.history) == 3


def test_prompt_carries_the_goal_the_menu_and_the_done_option():
    brain = _ScriptedModel(['{"kind": "done"}'])
    brain.decide("save the document", FRAME)
    prompt = brain.prompts[0]
    assert "save the document" in prompt
    assert "2. Save  (at" in prompt
    assert '"kind": "done"' in prompt
    assert "<number from the list>" in prompt      # placeholders, not worked examples


def test_text_only_brains_stop_when_there_is_nothing_to_click():
    class NeedsIndex(_ScriptedModel):
        requires_index = True

    with pytest.raises(GoalUnreachable, match="no clickable elements"):
        NeedsIndex(['{"kind": "done"}']).decide("save", Observation(elements=()))


def test_a_right_answer_in_the_wrong_field_is_recovered():
    """Real 0.5B output: named the correct element but put it in "reason"."""
    action = parse_decision('{"kind": "click", "index": 0, "reason": "Save"}', FRAME)
    assert action.kind == "click" and action.label == "Save"
    assert "recovered" in action.reason


def test_recovery_still_refuses_an_unresolvable_index():
    with pytest.raises(BrainError, match="not in the"):
        parse_decision('{"kind": "click", "index": 42, "reason": "no idea"}', FRAME)


def test_recovery_is_not_ambiguous():
    # two labels named and no usable index -> refuse rather than guess
    with pytest.raises(BrainError):
        parse_decision('{"kind": "click", "index": 0, "reason": "Save or Cancel"}', FRAME)