# tests/test_event_sourcing.py
from __future__ import annotations

import asyncio
from dataclasses import dataclass
from typing import Any

from casty import ActorRef, ActorSystem, Behaviors
from casty.actor import Behavior, SnapshotEvery
from casty.journal import InMemoryJournal, PersistedEvent, Snapshot


@dataclass(frozen=True)
class Increment:
    amount: int


@dataclass(frozen=True)
class GetValue:
    reply_to: ActorRef[int]


type CounterMsg = Increment | GetValue


@dataclass(frozen=True)
class Incremented:
    amount: int


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


def apply_event(state: CounterState, event: Incremented) -> CounterState:
    return CounterState(value=state.value + event.amount)


async def handle_command(ctx: Any, state: CounterState, cmd: CounterMsg) -> Any:
    match cmd:
        case Increment(amount):
            return Behaviors.persisted(events=[Incremented(amount)])
        case GetValue(reply_to):
            reply_to.tell(state.value)
            return Behaviors.same()
    return Behaviors.unhandled()


def counter_entity(entity_id: str, journal: InMemoryJournal) -> Behavior[CounterMsg]:
    return Behaviors.event_sourced(
        entity_id=entity_id,
        journal=journal,
        initial_state=CounterState(value=0),
        on_event=apply_event,
        on_command=handle_command,
    )


def counter_entity_with_snapshots(
    entity_id: str, journal: InMemoryJournal
) -> Behavior[CounterMsg]:
    return Behaviors.event_sourced(
        entity_id=entity_id,
        journal=journal,
        initial_state=CounterState(value=0),
        on_event=apply_event,
        on_command=handle_command,
        snapshot_policy=SnapshotEvery(n_events=3),
    )


async def test_event_sourced_basic_persist_and_read() -> None:
    journal = InMemoryJournal()
    async with ActorSystem(name="test") as system:
        ref = system.spawn(counter_entity("c1", journal), "counter")
        await asyncio.sleep(0.1)

        ref.tell(Increment(10))
        ref.tell(Increment(5))
        await asyncio.sleep(0.1)

        value = await system.ask(ref, lambda r: GetValue(reply_to=r), timeout=5.0)
        assert value == 15

    # Verify events were persisted
    events = await journal.load("c1")
    assert len(events) == 2
    assert events[0].event == Incremented(10)
    assert events[1].event == Incremented(5)


async def test_event_sourced_recovery_from_journal() -> None:
    journal = InMemoryJournal()

    # Pre-populate journal with events (simulating previous run)
    await journal.persist(
        "c1",
        [
            PersistedEvent(sequence_nr=1, event=Incremented(100), timestamp=1.0),
            PersistedEvent(sequence_nr=2, event=Incremented(50), timestamp=2.0),
        ],
    )

    async with ActorSystem(name="test") as system:
        ref = system.spawn(counter_entity("c1", journal), "counter")
        await asyncio.sleep(0.1)

        # State should be recovered: 100 + 50 = 150
        value = await system.ask(ref, lambda r: GetValue(reply_to=r), timeout=5.0)
        assert value == 150


async def test_snapshot_taken_after_n_events() -> None:
    journal = InMemoryJournal()
    async with ActorSystem(name="test") as system:
        ref = system.spawn(counter_entity_with_snapshots("c1", journal), "counter")
        await asyncio.sleep(0.1)

        for _i in range(3):
            ref.tell(Increment(1))
        await asyncio.sleep(0.1)

        snapshot = await journal.load_snapshot("c1")
        assert snapshot is not None
        assert snapshot.state == CounterState(value=3)
        assert snapshot.sequence_nr == 3


async def test_recovery_from_snapshot_plus_events() -> None:
    journal = InMemoryJournal()

    # Pre-populate with snapshot + newer events
    await journal.save_snapshot(
        "c1",
        Snapshot(sequence_nr=3, state=CounterState(value=30), timestamp=3.0),
    )
    await journal.persist(
        "c1",
        [
            PersistedEvent(sequence_nr=1, event=Incremented(10), timestamp=1.0),
            PersistedEvent(sequence_nr=2, event=Incremented(10), timestamp=2.0),
            PersistedEvent(sequence_nr=3, event=Incremented(10), timestamp=3.0),
            PersistedEvent(sequence_nr=4, event=Incremented(7), timestamp=4.0),
        ],
    )

    async with ActorSystem(name="test") as system:
        ref = system.spawn(counter_entity("c1", journal), "counter")
        await asyncio.sleep(0.1)

        # Should recover from snapshot (30) + event 4 (7) = 37
        value = await system.ask(ref, lambda r: GetValue(reply_to=r), timeout=5.0)
        assert value == 37


async def test_event_sourced_pushes_to_replicas() -> None:
    """Primary entity pushes persisted events to replica refs."""
    from casty.replication import ReplicateEvents

    journal = InMemoryJournal()
    received_replications: list[Any] = []

    async def mock_replica_handler(ctx: Any, msg: Any) -> Any:
        match msg:
            case ReplicateEvents():
                received_replications.append(msg)
        return Behaviors.same()

    async with ActorSystem(name="test") as system:
        replica_ref = system.spawn(Behaviors.receive(mock_replica_handler), "replica")
        await asyncio.sleep(0.1)

        def counter_with_replicas(entity_id: str) -> Behavior[CounterMsg]:
            return Behaviors.event_sourced(
                entity_id=entity_id,
                journal=journal,
                initial_state=CounterState(value=0),
                on_event=apply_event,
                on_command=handle_command,
                replica_refs=[replica_ref],
            )

        ref = system.spawn(counter_with_replicas("c1"), "counter")
        await asyncio.sleep(0.1)

        ref.tell(Increment(10))
        await asyncio.sleep(0.2)

        assert len(received_replications) == 1
        assert received_replications[0].entity_id == "c1"
        assert len(received_replications[0].events) == 1


async def test_event_sourced_waits_for_acks() -> None:
    """When min_acks > 0, primary waits for replica acks."""
    from casty.replication import ReplicateEvents, ReplicateEventsAck, ReplicationConfig

    journal = InMemoryJournal()
    replica_journal = InMemoryJournal()
    ack_events: list[Any] = []

    async def acking_replica_handler(ctx: Any, msg: Any) -> Any:
        match msg:
            case ReplicateEvents(entity_id=eid, events=events, reply_to=reply_to):
                await replica_journal.persist(eid, events)
                ack_events.append(msg)
                if reply_to is not None:
                    highest_seq = events[-1].sequence_nr if events else 0
                    reply_to.tell(ReplicateEventsAck(entity_id=eid, sequence_nr=highest_seq))
        return Behaviors.same()

    async with ActorSystem(name="test") as system:
        replica_ref = system.spawn(Behaviors.receive(acking_replica_handler), "replica")
        await asyncio.sleep(0.1)

        def counter_with_acks(entity_id: str) -> Behavior[CounterMsg]:
            return Behaviors.event_sourced(
                entity_id=entity_id,
                journal=journal,
                initial_state=CounterState(value=0),
                on_event=apply_event,
                on_command=handle_command,
                replica_refs=[replica_ref],
                replication=ReplicationConfig(replicas=1, min_acks=1, ack_timeout=5.0),
            )

        ref = system.spawn(counter_with_acks("c1"), "counter")
        await asyncio.sleep(0.1)

        ref.tell(Increment(10))
        await asyncio.sleep(0.3)

        # Verify replication was done
        assert len(ack_events) == 1

        # Verify state is correct (meaning ack was processed and didn't break anything)
        value = await system.ask(ref, lambda r: GetValue(reply_to=r), timeout=5.0)
        assert value == 10


async def test_event_sourced_fire_and_forget_replication() -> None:
    """When min_acks=0, primary doesn't wait for acks."""
    from casty.replication import ReplicateEvents, ReplicationConfig

    journal = InMemoryJournal()
    replications: list[Any] = []

    async def slow_replica_handler(ctx: Any, msg: Any) -> Any:
        match msg:
            case ReplicateEvents():
                await asyncio.sleep(0.5)  # Simulate slow replica
                replications.append(msg)
        return Behaviors.same()

    async with ActorSystem(name="test") as system:
        replica_ref = system.spawn(Behaviors.receive(slow_replica_handler), "replica")
        await asyncio.sleep(0.1)

        def counter_fire_forget(entity_id: str) -> Behavior[CounterMsg]:
            return Behaviors.event_sourced(
                entity_id=entity_id,
                journal=journal,
                initial_state=CounterState(value=0),
                on_event=apply_event,
                on_command=handle_command,
                replica_refs=[replica_ref],
                replication=ReplicationConfig(replicas=1, min_acks=0),
            )

        ref = system.spawn(counter_fire_forget("c1"), "counter")
        await asyncio.sleep(0.1)

        ref.tell(Increment(10))
        await asyncio.sleep(0.1)

        # Primary should respond quickly even though replica is slow
        value = await system.ask(ref, lambda r: GetValue(reply_to=r), timeout=5.0)
        assert value == 10
