# tests/test_event_sourced_sharding.py
from __future__ import annotations

import asyncio
from dataclasses import dataclass
from typing import Any

from casty import Behaviors, ActorRef, Behavior
from casty.core.journal import InMemoryJournal
from casty.cluster.system import ClusteredActorSystem
from casty.cluster.envelope import ShardEnvelope


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


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


type AccountMsg = Deposit | GetBalance


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


@dataclass(frozen=True)
class AccountState:
    balance: int


def apply_event(state: AccountState, event: Deposited) -> AccountState:
    return AccountState(balance=state.balance + event.amount)


async def handle_command(ctx: Any, state: AccountState, cmd: AccountMsg) -> Any:
    match cmd:
        case Deposit(amount):
            return Behaviors.persisted(events=[Deposited(amount)])
        case GetBalance(reply_to):
            reply_to.tell(state.balance)
            return Behaviors.same()
    return Behaviors.unhandled()


# Shared journal so we can inspect events after test
_journal = InMemoryJournal()


def account_entity(entity_id: str) -> Behavior[AccountMsg]:
    return Behaviors.event_sourced(
        entity_id=entity_id,
        journal=_journal,
        initial_state=AccountState(balance=0),
        on_event=apply_event,
        on_command=handle_command,
    )


async def test_event_sourced_sharding_workflow() -> None:
    global _journal
    _journal = InMemoryJournal()

    async with ClusteredActorSystem(
        name="bank", host="127.0.0.1", port=25520, node_id="node-1"
    ) as system:
        proxy = system.spawn(
            Behaviors.sharded(entity_factory=account_entity, num_shards=10),
            "accounts",
        )
        await asyncio.sleep(0.1)

        proxy.tell(ShardEnvelope("alice", Deposit(100)))
        proxy.tell(ShardEnvelope("alice", Deposit(50)))
        proxy.tell(ShardEnvelope("bob", Deposit(200)))
        await asyncio.sleep(0.3)

        balance_alice = await system.ask(
            proxy,
            lambda r: ShardEnvelope("alice", GetBalance(reply_to=r)),
            timeout=5.0,
        )
        assert balance_alice == 150

        balance_bob = await system.ask(
            proxy,
            lambda r: ShardEnvelope("bob", GetBalance(reply_to=r)),
            timeout=5.0,
        )
        assert balance_bob == 200

    # Verify events were persisted
    alice_events = await _journal.load("alice")
    assert len(alice_events) == 2
    bob_events = await _journal.load("bob")
    assert len(bob_events) == 1


async def test_sharded_with_replication_config() -> None:
    """Verify that sharded entities accept ReplicationConfig."""
    from casty.core.replication import ReplicationConfig

    _local_journal = InMemoryJournal()

    def replicated_account(entity_id: str) -> Behavior[AccountMsg]:
        return Behaviors.event_sourced(
            entity_id=entity_id,
            journal=_local_journal,
            initial_state=AccountState(balance=0),
            on_event=apply_event,
            on_command=handle_command,
        )

    async with ClusteredActorSystem(
        name="bank-repl", host="127.0.0.1", port=25530, node_id="node-1"
    ) as system:
        proxy = system.spawn(
            Behaviors.sharded(
                entity_factory=replicated_account,
                num_shards=10,
                replication=ReplicationConfig(replicas=1),
            ),
            "accounts",
        )
        await asyncio.sleep(0.1)

        proxy.tell(ShardEnvelope("alice", Deposit(100)))
        await asyncio.sleep(0.2)

        balance = await system.ask(
            proxy,
            lambda r: ShardEnvelope("alice", GetBalance(reply_to=r)),
            timeout=5.0,
        )
        assert balance == 100
