from __future__ import annotations

import json

from synaptopus.examples import ParityPressureState, XorDemoState, build_parity_pressure_demo, build_xor_novelty_demo
from synaptopus.runtime import run_until_acceptance_count
from synaptopus.snapshots import (
    create_demo_snapshot,
    load_demo_snapshot_json,
    restore_demo_snapshot,
    resume_demo_snapshot,
    save_demo_snapshot_json,
)


def test_xor_snapshot_round_trip_restores_state_and_system(tmp_path) -> None:
    system = build_xor_novelty_demo()
    record = run_until_acceptance_count(
        system,
        XorDemoState(),
        accepted_count=1,
        max_attempts_per_accept=4,
    )

    snapshot = create_demo_snapshot(
        "xor_novelty",
        system=system,
        record=record,
        parameters={"accepted_count": 1},
    )
    destination = tmp_path / "snapshot.json"
    save_demo_snapshot_json(snapshot, destination)

    loaded = load_demo_snapshot_json(destination)
    restored_system, restored_record = restore_demo_snapshot(loaded)

    assert destination.exists()
    assert restored_system is not None
    assert restored_record.final_state == record.final_state
    assert tuple(step.candidate for step in restored_record.accepted) == ((0, 1),)


def test_parity_snapshot_resume_matches_continuous_run() -> None:
    checkpoint_system = build_parity_pressure_demo()
    partial = run_until_acceptance_count(
        checkpoint_system,
        ParityPressureState(),
        accepted_count=2,
        max_attempts_per_accept=12,
    )
    snapshot = create_demo_snapshot(
        "parity_pressure",
        system=checkpoint_system,
        record=partial,
    )

    resumed_system, resumed = resume_demo_snapshot(
        snapshot,
        additional_accepted_count=2,
        max_attempts_per_accept=12,
    )
    continuous = run_until_acceptance_count(
        build_parity_pressure_demo(),
        ParityPressureState(),
        accepted_count=4,
        max_attempts_per_accept=12,
    )

    assert resumed_system is not None
    assert tuple(step.candidate for step in resumed.accepted) == tuple(
        step.candidate for step in continuous.accepted
    )
    assert resumed.attempt_count == continuous.attempt_count
    assert resumed.final_state == continuous.final_state


def test_snapshot_json_uses_demo_snapshot_artifact_type(tmp_path) -> None:
    system = build_xor_novelty_demo()
    record = run_until_acceptance_count(
        system,
        XorDemoState(),
        accepted_count=1,
        max_attempts_per_accept=4,
    )
    snapshot = create_demo_snapshot("xor_novelty", system=system, record=record)
    destination = tmp_path / "snapshot.json"

    save_demo_snapshot_json(snapshot, destination)
    payload = json.loads(destination.read_text(encoding="utf-8"))

    assert payload["artifact_type"] == "demo_snapshot"
    assert payload["payload"]["demo_name"] == "xor_novelty"