temp/features/steps/langgraph_state_branch_coverage_steps.py

"""
Behave step definitions to cover remaining branches in langgraph/state.py.
"""

import json
import tempfile
from pathlib import Path

from behave import given, then, when
from behave.runner import Context

from cleveragents.langgraph.state import GraphState, StateManager, StateUpdateMode


def _ensure(context: Context) -> None:
    context.state_instances = getattr(context, "state_instances", {}) or {}
    context.state_managers = getattr(context, "state_managers", {}) or {}
    context.test_data = getattr(context, "test_data", {}) or {}
    context.results = getattr(context, "results", {}) or {}


@given('I have a GraphState with metadata {metadata_json} and current node "{node}"')
def step_graphstate_replace(context: Context, metadata_json: str, node: str):
    _ensure(context)
    metadata = json.loads(metadata_json)
    context.state_instances["replace_state"] = GraphState(
        metadata=metadata, current_node=node
    )


@when(
    'I replace the state with updates containing metadata {metadata_json} and current node "{node}"'
)
def step_replace_updates(context: Context, metadata_json: str, node: str):
    _ensure(context)
    state = context.state_instances["replace_state"]
    updates = {"metadata": json.loads(metadata_json), "current_node": node}
    state.update(updates, StateUpdateMode.REPLACE)


@then("the state's metadata should equal {metadata_json}")
def step_verify_metadata_equals(context: Context, metadata_json: str):
    state = context.state_instances.get("replace_state")
    assert state is not None
    expected = json.loads(metadata_json)
    assert state.metadata == expected


@then('the state\'s current node should equal "{node}"')
def step_verify_current_node_equals(context: Context, node: str):
    state = context.state_instances.get("replace_state")
    assert state is not None
    assert state.current_node == node


@given("I have a GraphState with metadata {metadata_json}")
def step_graphstate_metadata(context: Context, metadata_json: str):
    _ensure(context)
    context.state_instances["merge_meta"] = GraphState(
        metadata=json.loads(metadata_json)
    )


@given("I have merge metadata updates {metadata_json}")
def step_merge_metadata_updates(context: Context, metadata_json: str):
    _ensure(context)
    context.test_data["merge_metadata"] = json.loads(metadata_json)


@when("I merge metadata updates into the state")
def step_merge_metadata(context: Context):
    _ensure(context)
    state = context.state_instances["merge_meta"]
    state.update(
        {"metadata": context.test_data["merge_metadata"]}, StateUpdateMode.MERGE
    )


@then('the state\'s metadata should contain keys "{key1}" and "{key2}"')
def step_verify_metadata_keys(context: Context, key1: str, key2: str):
    state = context.state_instances.get("merge_meta")
    assert state is not None
    assert key1 in state.metadata
    assert key2 in state.metadata


@given("I have a GraphState with 1 message")
def step_graphstate_one_message(context: Context):
    _ensure(context)
    context.state_instances["append_single"] = GraphState(messages=[{"id": 0}])


@when("I append a single message with id {message_id:d} into the state")
def step_append_single_message(context: Context, message_id: int):
    _ensure(context)
    state = context.state_instances["append_single"]
    state.update({"messages": {"id": message_id}}, StateUpdateMode.APPEND)


@then("the state's messages should end with id {message_id:d}")
def step_verify_messages_end(context: Context, message_id: int):
    state = context.state_instances.get("append_single")
    assert state is not None
    assert state.messages[-1]["id"] == message_id


@then("the message list length should be 2")
def step_verify_message_length_two(context: Context):
    state = context.state_instances.get("append_single")
    assert state is not None
    assert len(state.messages) == 2


@given("I have a StateManager without checkpointing")
def step_manager_without_checkpoint(context: Context):
    _ensure(context)
    context.state_managers["no_checkpoint"] = StateManager()


@when("I invoke checkpoint saving manually")
def step_invoke_save_checkpoint(context: Context):
    manager = context.state_managers["no_checkpoint"]
    manager._save_checkpoint()


@then("the update count should remain at 0")
def step_verify_update_count_zero(context: Context):
    manager = context.state_managers["no_checkpoint"]
    assert manager.update_count == 0


@given(
    "I have a checkpoint file with execution count {count:d} and metadata {metadata_json}"
)
def step_prepare_checkpoint_file(context: Context, count: int, metadata_json: str):
    _ensure(context)
    checkpoint_dir = Path(tempfile.mkdtemp())
    checkpoint_file = checkpoint_dir / "checkpoint_test.json"
    data = {
        "state": GraphState(metadata=json.loads(metadata_json)).to_dict(),
        "timestamp": "test",
        "update_count": count,
    }
    checkpoint_file.write_text(json.dumps(data), encoding="utf-8")
    context.test_data["checkpoint_file"] = checkpoint_file


@when("I load the checkpoint into a StateManager")
def step_load_checkpoint(context: Context):
    manager = StateManager()
    checkpoint_file = context.test_data["checkpoint_file"]
    manager.load_checkpoint(checkpoint_file)
    context.results["loaded_manager"] = manager


@then("the manager state should include metadata {metadata_json}")
def step_verify_loaded_metadata(context: Context, metadata_json: str):
    manager: StateManager = context.results["loaded_manager"]
    expected = json.loads(metadata_json)
    assert manager.state.metadata == expected


@then("the manager update count should be {count:d}")
def step_verify_loaded_update_count(context: Context, count: int):
    manager: StateManager = context.results["loaded_manager"]
    assert manager.update_count == count


@when("I request the latest checkpoint")
def step_request_latest_checkpoint(context: Context):
    manager = context.state_managers.get("no_checkpoint") or context.results.get(
        "latest_manager"
    )
    context.results["latest_checkpoint"] = (
        manager.get_latest_checkpoint() if manager else None
    )


@then("no checkpoint path should be returned")
def step_verify_latest_none(context: Context):
    assert context.results["latest_checkpoint"] is None


@given("I have a StateManager with two checkpoints written at different times")
def step_manager_with_two_checkpoints(context: Context):
    _ensure(context)
    checkpoint_dir = Path(tempfile.mkdtemp())
    manager = StateManager(checkpoint_dir=checkpoint_dir)

    # Write first checkpoint
    manager.state.metadata = {"first": True}
    manager.update_count = 1
    manager._save_checkpoint()

    # Write second checkpoint with newer mtime
    manager.state.metadata = {"second": True}
    manager.update_count = 2
    manager._save_checkpoint()

    context.results["latest_manager"] = manager


@then("the newest checkpoint path should be returned")
def step_verify_latest_checkpoint(context: Context):
    manager: StateManager = context.results["latest_manager"]
    latest = manager.get_latest_checkpoint()
    assert latest is not None
    checkpoint_files = list(manager.checkpoint_dir.glob("checkpoint_*.json"))
    assert latest == max(checkpoint_files, key=lambda p: p.stat().st_mtime)


@given("I have a StateManager with time travel disabled")
def step_manager_time_travel_disabled(context: Context):
    _ensure(context)
    context.state_managers["tt_disabled"] = StateManager(enable_time_travel=False)


@when("I request time travel")
def step_request_time_travel(context: Context):
    manager = context.state_managers["tt_disabled"]
    context.results["tt_result"] = manager.time_travel()


@then("the returned time travel state should be None")
def step_verify_time_travel_none(context: Context):
    assert context.results["tt_result"] is None


@given("I have a StateManager with time travel history entries")
def step_manager_with_history(context: Context):
    _ensure(context)
    manager = StateManager(enable_time_travel=True)
    manager.update_state({"execution_count": 1, "metadata": {"keep": "me"}})
    manager.update_state({"execution_count": 2, "metadata": {"keep": "me2"}})
    context.state_managers["history_reset"] = manager


@when("I clear the history and then reset with metadata {metadata_json}")
def step_clear_and_reset(context: Context, metadata_json: str):
    manager = context.state_managers["history_reset"]
    manager.clear_history()
    manager.reset(GraphState(metadata=json.loads(metadata_json)))
    context.results["reset_manager"] = manager


@then("the history should be empty after reset")
def step_verify_history_empty(context: Context):
    manager: StateManager = context.results["reset_manager"]
    assert len(manager.history) == 0


@then("the execution count should be zero after reset")
def step_verify_execution_zero(context: Context):
    manager: StateManager = context.results["reset_manager"]
    assert manager.update_count == 0


@then("the state metadata should include {metadata_json}")
def step_verify_state_metadata_after_reset(context: Context, metadata_json: str):
    manager: StateManager = context.results["reset_manager"]
    expected = json.loads(metadata_json)
    assert manager.state.metadata == expected