from __future__ import annotations

from tests.harness.test_runner import TestCase

SUITE = "test_restart"


def _clean_restart(ctx):
    ok = ctx.agent.restart(timeout=30)
    assert ok, "restart failed"
    resp = ctx.client.status()
    assert resp.get("success") is True
    return True, "restart succeeded", {}


def _status_after_restart(ctx):
    before = ctx.client.status()
    assert before.get("success") is True
    ok = ctx.agent.restart(timeout=30)
    assert ok, "restart failed"
    after = ctx.client.status()
    assert after.get("success") is True
    assert before.get("pubkey") == after.get("pubkey"), "pubkey changed after restart"
    return True, "status stable after restart", {"pubkey": after.get("pubkey")}


def _conversation_after_restart(ctx):
    ok = ctx.agent.restart(timeout=30)
    assert ok, "restart failed"
    resp = ctx.client.prompt("After restart, say hello.", max_turns=3)
    assert resp.get("success") is True
    assert str(resp.get("final_response", "")).strip(), "empty final response"
    return True, "conversation works post-restart", {}


def get_tests():
    return [
        TestCase(SUITE, "clean_restart", "Stop/start restart", _clean_restart, requires_restart=True),
        TestCase(SUITE, "status_after_restart", "Status stable after restart", _status_after_restart),
        TestCase(SUITE, "conversation_after_restart", "Prompt after restart", _conversation_after_restart),
    ]
