"""Commit and reopen the same complete state with either optional SQL backend.

Install sqlite or duckdb extras, then run ``python examples/persistence.py
--backend sqlite`` (or ``duckdb``). This synchronous example owns the database
on its current thread; embedded async applications use Client's managed worker.
"""

from __future__ import annotations

import argparse
from dataclasses import replace
from datetime import UTC, datetime, timedelta
from pathlib import Path
from tempfile import TemporaryDirectory

from rpkiparrot import MemoryStore
from rpkiparrot.config import FreshnessPolicy
from rpkiparrot.models import Afi, Aspa, AspaRecord, EvaluationContext, PayloadKind
from rpkiparrot.persistence import DuckDBBackend, SQLiteBackend, state_digest, state_from_snapshot
from rpkiparrot.readers import parse_json


def main(backend_name: str) -> None:
    """Show atomic commit, original receipt on retry, and verified persisted state."""
    at = datetime(2026, 1, 1, tzinfo=UTC)
    memory = MemoryStore(context=EvaluationContext.offline(at=at))
    dataset = parse_json(
        '{"roas":[{"asn":"AS64496","prefix":"192.0.2.0/24","maxLength":24}]}',
        format="routinator-json",
    )
    # A host can provide different original authorizations for each address
    # family. Modern producer JSON omits AFI and explicitly means both families.
    dataset = replace(
        dataset,
        present=dataset.present | {PayloadKind.ASPA},
        aspa_families=frozenset(Afi),
        aspas=(
            AspaRecord(aspa=Aspa(customer=64496, providers=frozenset({64497}), afi=Afi.IPV4)),
            AspaRecord(aspa=Aspa(customer=64496, providers=frozenset({64498}), afi=Afi.IPV6)),
        ),
    )
    snapshot = memory.replace_source(
        "chosen-input", dataset, freshness=FreshnessPolicy(valid_until=at + timedelta(hours=1))
    )
    state = state_from_snapshot(snapshot)
    backend_type = SQLiteBackend if backend_name == "sqlite" else DuckDBBackend
    with TemporaryDirectory() as directory:
        path = Path(directory) / "payload.db"
        owner = backend_type(path)
        owner.open()
        try:
            receipt = owner.commit(state, expected=None)
            assert owner.commit(state, expected=None) == receipt
        finally:
            owner.close()
        restored_owner = backend_type(path)
        restored_owner.open()
        try:
            # A recognized transient read failure raises PersistenceError with
            # code="persistence_failed", retryable=True; the host may retry.
            # A corruption error must not be handled by deleting this database
            # or trusting unverified records. This example propagates failures.
            restored = restored_owner.read_state()
            assert restored is not None and restored.receipt == receipt
            assert restored.sources == state.sources
            assert restored.sources[0].dataset.aspa_families == frozenset(Afi)
            assert [record.aspa.afi for record in restored.sources[0].dataset.aspas] == list(Afi)
            assert restored.aspa_active_groups == {afi: "memory" for afi in Afi}
            assert receipt.schema_version == 2
            assert state_digest(restored) == receipt.state_digest
            assert restored.sources[0].info.expires_at == at + timedelta(hours=1)
            print(f"{backend_name}: verified generation {receipt.snapshot_id.generation}")
        finally:
            restored_owner.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--backend", choices=("sqlite", "duckdb"), default="sqlite")
    main(parser.parse_args().backend)