from __future__ import annotations

import json
from dataclasses import dataclass

from synaptopus.runtime import StepTrace, run_until_acceptance_count
from synaptopus.serialization import (
    deserialize_execution_record,
    save_execution_record_json,
    serialize_execution_record,
)


@dataclass(frozen=True)
class SequenceState:
    accepted: tuple[int, ...] = ()
    attempts: int = 0


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


def test_execution_record_serialization_converts_dataclasses() -> None:
    record = run_until_acceptance_count(
        EvenAcceptanceSystem(),
        SequenceState(),
        accepted_count=2,
        max_attempts_per_accept=4,
    )

    serialized = serialize_execution_record(record)

    assert serialized.final_state["accepted"] == [2, 4]
    assert serialized.attempts[0].previous_state["attempts"] == 0
    assert serialized.attempts[1].candidate == 2
    assert serialized.accepted[0].metadata["attempt"] == 2


def test_execution_record_serialization_saves_json(tmp_path) -> None:
    record = run_until_acceptance_count(
        EvenAcceptanceSystem(),
        SequenceState(),
        accepted_count=1,
        max_attempts_per_accept=4,
    )
    destination = tmp_path / "trace.json"

    save_execution_record_json(record, destination)
    loaded = json.loads(destination.read_text(encoding="utf-8"))

    assert loaded["artifact_type"] == "execution_trace"
    assert loaded["payload"]["accepted"][0]["candidate"] == 2
    assert loaded["payload"]["final_state"]["accepted"] == [2]


def test_execution_record_deserialization_restores_typed_record() -> None:
    record = run_until_acceptance_count(
        EvenAcceptanceSystem(),
        SequenceState(),
        accepted_count=2,
        max_attempts_per_accept=4,
    )

    serialized = serialize_execution_record(record)
    restored = deserialize_execution_record(
        serialized,
        state_decoder=lambda data: SequenceState(
            accepted=tuple(int(value) for value in data["accepted"]),
            attempts=int(data["attempts"]),
        ),
        candidate_decoder=lambda data: int(data),
        metadata_decoder=lambda data: None if data is None else {"attempt": int(data["attempt"])},
    )

    assert restored.final_state == record.final_state
    assert tuple(step.candidate for step in restored.accepted) == (2, 4)
    assert restored.attempts[0].metadata == {"attempt": 1}