"""The stepwise API: a console needs to run the loop one turn at a time.

`run()` is the batch form. A live console needs the opposite - the caller owns the
loop, so a button press advances exactly one step and a pause switch stops between
them. Both are the same rules; these tests pin that down, including an
equivalence check so the two forms cannot drift apart.
"""

import pytest

from tiny_computer_agent import (
    AgentLoop, Element, FakeScreen, MockBrain, PlannerBrain, Policy,
)


def script():
    return MockBrain([
        {"kind": "wait", "seconds": 0.1},
        {"kind": "click", "label": "Save"},
        {"kind": "key", "keys": "ctrl+s"},
        {"kind": "done", "reason": "saved"},
    ])


def screen():
    return FakeScreen([
        (Element("Files", x=10, y=6, width=60, height=20),),
        (Element("Save", x=320, y=120, width=80, height=24),),
    ])


def shape(episode):
    """Everything about an episode that a user could notice."""
    return (
        episode.stop_reason,
        episode.success,
        [
            (s.index,
             s.action.describe() if s.action else None,
             s.error,
             s.result.summary() if s.result else None)
            for s in episode.steps
        ],
    )


# -- stepping ---------------------------------------------------------------

def test_begin_reads_the_screen_once_and_step_once_advances_one_turn():
    screen_ = screen()
    loop = AgentLoop(screen=screen_, brain=script()).begin("save")
    assert screen_.captures == 1                    # one read, before the first decision
    assert loop.finished is False
    assert loop.episode.steps == []

    first = loop.step_once()
    assert first.index == 0 and first.action.kind == "wait"
    assert len(loop.episode.steps) == 1
    assert loop.finished is False                   # three turns still to go

    loop.step_once()
    assert [s.action.kind for s in loop.episode.steps] == ["wait", "click"]


def test_stepping_past_the_end_returns_none_instead_of_raising():
    loop = AgentLoop(screen=screen(), brain=script()).begin("save")
    while loop.step_once() is not None:
        pass
    assert loop.finished is True
    assert loop.episode.success is True
    assert loop.step_once() is None
    assert loop.step_once() is None


def test_step_once_before_begin_does_nothing():
    loop = AgentLoop(screen=screen(), brain=script())
    assert loop.step_once() is None
    assert loop.episode is None


def test_the_budget_is_enforced_one_step_at_a_time():
    loop = AgentLoop(
        screen=screen(),
        brain=MockBrain([{"kind": "wait", "seconds": 0.2 + i} for i in range(10)]),
        policy=Policy(step_budget=3),
    ).begin("save")
    taken = 0
    while loop.step_once() is not None:
        taken += 1
    assert taken == 3
    assert "budget of 3" in loop.episode.stop_reason


# -- pause and resume -------------------------------------------------------

def test_pause_stops_between_steps_and_leaves_the_episode_open():
    loop = AgentLoop(screen=screen(), brain=script()).begin("save")
    loop.step_once()
    loop.pause()
    episode = loop.run()                       # returns immediately
    assert len(episode.steps) == 1
    assert loop.finished is False              # not finished, only paused
    assert loop.paused is True


def test_resume_continues_the_same_episode_to_the_same_end():
    loop = AgentLoop(screen=screen(), brain=script()).begin("save")
    loop.step_once()
    loop.pause()
    loop.run()
    loop.resume()
    episode = loop.run()
    assert episode.success is True
    assert [s.action.kind for s in episode.steps] == ["wait", "click", "key", "done"]
    assert loop.paused is False


def test_a_paused_run_is_equivalent_to_an_uninterrupted_one():
    straight = AgentLoop(screen=screen(), brain=script()).run("save")

    interrupted = AgentLoop(screen=screen(), brain=script()).begin("save")
    interrupted.step_once()
    interrupted.pause()
    interrupted.run()
    interrupted.resume()
    interrupted.run()

    assert shape(interrupted.episode) == shape(straight)


# -- the two forms cannot drift apart ---------------------------------------

def test_run_is_exactly_begin_plus_step_once_in_a_loop():
    batch = AgentLoop(screen=screen(), brain=PlannerBrain()).run("click 'Save'")

    stepwise = AgentLoop(screen=screen(), brain=PlannerBrain()).begin("click 'Save'")
    while stepwise.step_once() is not None:
        pass

    assert shape(stepwise.episode) == shape(batch)


def test_verify_still_owns_success_when_stepping():
    texts = ["Unsaved changes", "before", "Saved"]

    class TextScreen:
        def __init__(self):
            self.n = 0

        def capture(self):
            from tiny_computer_agent import Observation
            text = texts[min(self.n, len(texts) - 1)]
            self.n += 1
            return Observation(index=self.n - 1, elements=(Element("Save"),), text=text)

    loop = AgentLoop(
        screen=TextScreen(),
        brain=MockBrain([{"kind": "wait", "seconds": 0.1}] * 4),
        verify=lambda observation: "saved" in observation.text.lower(),
    ).begin("save")
    while loop.step_once() is not None:
        pass
    assert loop.episode.success is True
    assert loop.episode.stop_reason == "verified after 2 action(s)"


def test_a_kill_switch_stops_a_stepwise_run_before_the_first_action(tmp_path):
    marker = tmp_path / "STOP"
    marker.write_text("stop")
    loop = AgentLoop(
        screen=screen(),
        brain=MockBrain([{"kind": "click", "label": "Save"}]),
        policy=Policy(kill_switch_path=str(marker)),
    ).begin("save")
    assert loop.step_once() is None
    assert loop.episode.steps == []
    assert loop.episode.stop_reason == "kill switch engaged"
    assert loop.finished is True


def test_a_dead_screen_finishes_the_episode_at_begin():
    class Blind:
        def capture(self):
            raise RuntimeError("X connection failed")

    loop = AgentLoop(screen=Blind(), brain=script()).begin("anything")
    assert loop.finished is True
    assert loop.episode.steps == []
    assert loop.episode.stop_reason.startswith("capture failed: RuntimeError")
    assert loop.step_once() is None