from __future__ import annotations

from dataclasses import dataclass

from synaptopus.runtime import (
    StepTrace,
    merge_execution_records,
    run_until_acceptance,
    run_until_acceptance_count,
)


@dataclass(frozen=True)
class CounterState:
    value: int = 0


class EvenAcceptanceSystem:
    def step(self, state: CounterState) -> StepTrace[CounterState, int, dict[str, int]]:
        next_state = CounterState(state.value + 1)
        return StepTrace(
            previous_state=state,
            next_state=next_state,
            candidate=next_state.value,
            accepted=next_state.value % 2 == 0,
            elapsed_seconds=0.0,
            metadata={"value": next_state.value},
        )


def test_run_until_acceptance_collects_rejected_attempts() -> None:
    record = run_until_acceptance(EvenAcceptanceSystem(), CounterState(), max_attempts=5)

    assert record.accepted_count == 1
    assert record.attempt_count == 2
    assert record.final_state == CounterState(2)
    assert record.accepted[0].candidate == 2
    assert record.attempts[0].accepted is False
    assert record.attempts[1].accepted is True


def test_run_until_acceptance_count_collects_multiple_acceptances() -> None:
    record = run_until_acceptance_count(
        EvenAcceptanceSystem(),
        CounterState(),
        accepted_count=3,
        max_attempts_per_accept=5,
    )

    assert record.accepted_count == 3
    assert record.attempt_count == 6
    assert tuple(step.candidate for step in record.accepted) == (2, 4, 6)
    assert record.final_state == CounterState(6)


def test_merge_execution_records_appends_continuation() -> None:
    first = run_until_acceptance_count(
        EvenAcceptanceSystem(),
        CounterState(),
        accepted_count=1,
        max_attempts_per_accept=5,
    )
    second = run_until_acceptance_count(
        EvenAcceptanceSystem(),
        first.final_state,
        accepted_count=2,
        max_attempts_per_accept=5,
    )

    merged = merge_execution_records(first, second)

    assert merged.accepted_count == 3
    assert merged.attempt_count == 6
    assert tuple(step.candidate for step in merged.accepted) == (2, 4, 6)
    assert merged.final_state == CounterState(6)